Fast Implementation of the 1 Parameter Division Distortion Model

I have implemented some code that allows me to estimate the center of distortion for an image along with a distortion parameter that attempts to remove any radial distortion present in the image. It is modelled after the one-parameter division model by Fitzgibbon:

xu = xd / (1 + K1 * r^2)

yu = yd / (1 + K1 * r^2)

I am wondering what the best way is to process this transform as fast as possible. I was thinking there should be a way to use numpy vectorization to create a mapping for a given image that I can use to perform the undistortion procedure. I have experience with OpenCV (undistort(), initUndistortRectifyMap(), etc.), however, these methods require an estimate of the camera focal properties (fx, fy) which I do not have.

I could implement a nested for loop structure to calculate all of the undistorted destinations for a given input image size, however, this would be very inefficient. Is there an approach I can take with numpy to build this mapping?

Here is how I transform a single point from the distorted state to the undistorted state.

def get_undistorted(pd, dc, k):
    xd, yd = pd
    dcx, dcy = dc
    r2 = (dcx - xd)**2 + (dcy - yd)**2
    xu = dcx + (xd - dcx)/(1 + k*r2)
    yu = dcy + (yd - dcy)/(1 + k*r2)
    return xu, yu

I can't promise that this is the most efficient way to apply the transform, but it is fully vectorised, and the interpolation is configurable. The coordinate space is pre-computed so that only the interpolation operation itself needs to be done once per frame.

from typing import NamedTuple

import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from scipy.interpolate import RegularGridInterpolator


class Transform(NamedTuple):
    mn_high: tuple[int, int]
    interp_grid: tuple[np.ndarray, np.ndarray]
    yxdistort: np.ndarray
    yxfixed: np.ndarray
    rect_axes: tuple[np.ndarray, np.ndarray]

    @classmethod
    def setup(
        cls,
        dc: np.ndarray,
        k: float,
        mn: tuple[int, int],
        scale: float,
    ) -> 'Transform':
        ylong = np.arange(0, mn[0], scale)
        xlong = np.arange(0, mn[1], scale)
        yxrect_d3 = np.stack(
            np.meshgrid(xlong, ylong), axis=-1,
        )
        yxrect = yxrect_d3.reshape((-1, 2))[:, ::-1]

        interp_grid = tuple(
            dim.ravel()
            for dim in np.indices(dimensions=mn, sparse=True)
        )

        return cls(
            mn_high=yxrect_d3.shape[:2],
            interp_grid=interp_grid,
            yxdistort=redistort(dc=dc, k=k, vu=yxrect),
            yxfixed=undistort(dc=dc, k=k, yx=yxrect),
            rect_axes=(ylong, xlong),
        )

    def distort_image(self, image: np.ndarray) -> np.ndarray:
        distort_interp = RegularGridInterpolator(
            points=self.interp_grid, values=image, bounds_error=False,
        )
        return distort_interp(self.yxdistort).reshape(
            (*self.mn_high, -1)
        )

    def undistort_image(self, image: np.ndarray) -> np.ndarray:
        undistort_interp = RegularGridInterpolator(
            points=self.rect_axes, values=image, bounds_error=False,
        )
        return undistort_interp(self.yxfixed).reshape(
            (*self.mn_high, -1)
        )


def undistort(
    dc: np.ndarray,  # centre point, *2
    k: float,        # distort parameter aka. lambda
    yx: np.ndarray,  # distorted, n * 2
) -> np.ndarray:     # n * 2
    r2 = ((yx - dc)**2).sum(axis=1, keepdims=True)
    vu = dc + (yx - dc)/(1 + k*r2)
    return vu


def redistort(
    dc: np.ndarray,  # centre point, *2
    k: float,        # distort parameter aka. lambda
    vu: np.ndarray,  # non-distorted, n * 2
) -> np.ndarray:     # n * 2
    inner = k*((vu - dc)**2).sum(axis=1, keepdims=True)

    return (
        k*(
            dc*(
                + (vu**2).sum(axis=1, keepdims=True)
                - 2*dc*vu
                + dc**2
                + dc[::-1]**2
            )
            - 2*dc.prod()*vu[:, ::-1]
        )
        + 0.5*(vu - dc)*(
            1 - np.sqrt(1 - 4*inner)
        )
    )/inner


def regression_test() -> None:
    dcx = 320/2
    dcy = 240/2
    xd = 5
    yd = 17
    lambd = 1e-6

    r2 = (dcx - xd)**2 + (dcy - yd)**2
    xu = dcx + (xd - dcx)/(1 + lambd*r2)
    yu = dcy + (yd - dcy)/(1 + lambd*r2)

    dc = np.array((dcy, dcx))
    actual = undistort(
        dc=dc, k=lambd, yx=np.array((
            (yd, xd),
        )),
    )
    assert np.allclose(actual, (yu, xu))

    vu = np.array((
        (yu, xu),
        (40, 30),
        (25, 20),
    ))
    expected = np.array((
        (         yd,          xd),
        (38.04372287, 26.82104966),
        (22.11282237, 15.74521192),
    ))
    actual = redistort(dc, lambd, vu)
    assert np.allclose(expected, actual, rtol=0, atol=1e-8)


def estimate_resolution_loss(
    mn: tuple[int, int],
    dc: np.ndarray,
    k: float,
) -> float:
    # Pretty bad! This is not efficient.
    n = 200
    middle_strip = np.empty((n, 2))
    middle_strip[:, 0] = np.linspace(0, mn[0]-1, n)
    middle_strip[:, 1] = mn[1]/2

    distorted = redistort(dc=dc, k=k, vu=middle_strip)
    y, x = distorted.T
    y = y[np.isfinite(y)]

    res_reduce = mn[0]/(y.max() - y.min())
    return res_reduce


def roundtrip_demo(filename: str) -> None:
    # Once per parameter+resolution set
    imorig = np.array(Image.open(filename))
    mn = imorig.shape[:2]
    dc = 0.5*np.array(mn) + 0.1  # avoid a divide-by-zero
    k = 1e-6
    res_reduce = estimate_resolution_loss(mn=mn, dc=dc, k=k)
    transform = Transform.setup(mn=mn, dc=dc, k=k, scale=res_reduce)

    # Once every frame
    imdistort = transform.distort_image(imorig/255)
    imfixed = transform.undistort_image(imdistort)

    fig, (ax_orig, ax_distort, ax_fixed) = plt.subplots(ncols=3)
    ax_orig.imshow(imorig)
    ax_distort.imshow(imdistort)
    ax_fixed.imshow(imfixed)
    ax_orig.set_title('Original')
    ax_distort.set_title('Distorted')
    ax_fixed.set_title('Corrected')
    plt.show()


if __name__ == '__main__':
    regression_test()
    roundtrip_demo('baby goat.jpg')

More can be done to handle the edge better, etc.