albumentations-team / albumentations

Fast and flexible image augmentation library. Paper about the library: https://www.mdpi.com/2078-2489/11/2/125
https://albumentations.ai
MIT License
14.25k stars 1.65k forks source link

feature request: full lens distortion control #876

Open flugenheimer opened 3 years ago

flugenheimer commented 3 years ago

First of I really like what you guys are doing here with this library. Keep up the good work!

It would be nice to have full control over the lens distortion parameters. I see that there is already barrel/pinch distortion def albumentations.augmentations.functional.optical_distortion(img, k=0, dx=0, dy=0, interpolation=1, border_mode=4, value=None)

This is however not the full story: there is a nice visualization at the MATLAB page here: https://se.mathworks.com/help/vision/ug/camera-calibration.html image

I found an implementation of lens distortion with the radial distortions k1, k2 and k3, as well as tangential distortion p1 and p2 here: https://towardsdatascience.com/image-augmentation-mastering-15-techniques-and-useful-functions-with-python-codes-44c3f8c1ea1f

class LensDistortion(object):
    def __init__(self, d_coef=(0.15, 0.15, 0.1, 0.1, 0.05)):
        self.d_coef = np.array(d_coef)

    def __call__(self, X, Y):

        # get the height and the width of the image
        h, w = X.shape[:2]

        # compute its diagonal
        f = (h ** 2 + w ** 2) ** 0.5

        # set the image projective to carrtesian dimension
        K = np.array([[f, 0, w / 2],
                      [0, f, h / 2],
                      [0, 0,     1]])

        d_coef = self.d_coef * np.random.random(5) # value
        d_coef = d_coef * (2 * (np.random.random(5) < 0.5) - 1) # sign
        # Generate new camera matrix from parameters
        M, _ = cv2.getOptimalNewCameraMatrix(K, d_coef, (w, h), 0)

        # Generate look-up tables for remapping the camera image
        remap = cv2.initUndistortRectifyMap(K, d_coef, None, M, (w, h), 5)

        # Remap the original image to a new image
        X = cv2.remap(X, *remap, cv2.INTER_LINEAR)
        Y = cv2.remap(Y, *remap, cv2.INTER_LINEAR)
        return X, Y

Modifying this to fit albumentation to control lens distortion would be a nice addition.

flugenheimer commented 3 years ago

it should be stated that d_coef=(k1, k2, p1, p2, k3)