Error in keypoints when using large optical distortions

[ThomasieDK]

Describe the bug

When applying large deformations using opticaldistortion the keypoints are not transformed correctly.

To Reproduce

OS: Windows
Python version 3.11
Albumentations V2.0.0

Steps to reproduce the behavior:

import cv2
import numpy as np
import albumentations as A
from albumentations.core.composition import KeypointParams
import matplotlib.pyplot as plt

# ---------------------------
# A) Generate a Checkerboard
# ---------------------------
def generate_checkerboard(rows=8, cols=8, square_size=50):
    """
    Single-channel checkerboard: 0=black squares, 255=white squares.
    'rows' and 'cols' = number of squares in each dimension.
    """
    height = rows * square_size
    width = cols * square_size
    cb = np.ones((height, width), dtype=np.uint8) * 255

    for r in range(rows):
        for c in range(cols):
            if (r + c) % 2 == 0:
                y1 = r * square_size
                x1 = c * square_size
                cb[y1:y1 + square_size, x1:x1 + square_size] = 0
    return cb


def get_internal_checker_corners(rows, cols, square_size):
    """
    Return the (rows-1)*(cols-1) 'internal' corners typically used for calibration.
    E.g. for 8x8 squares, we get 7x7=49 points.
    """
    points = []
    for r in range(1, rows):
        for c in range(1, cols):
            x = c * square_size
            y = r * square_size
            points.append((x, y))
    return points


transform = A.Compose([
    # A.RandomCrop(width=330, height=330),
    # A.RandomBrightnessContrast(p=0.2),
    A.OpticalDistortion(p=1, mode='fisheye', distort_limit=(-1.05, -1.05),
    interpolation = 4,mask_interpolation=4)
], keypoint_params=A.KeypointParams(format='xy', remove_invisible=True,
                                    angle_in_degrees=True,
                                    check_each_transform=True))

image = generate_checkerboard(rows=16, cols=16, square_size=25)
keypoints = get_internal_checker_corners(rows=16, cols=16, square_size=25)
keypoints_orig = [(x, y, 0, 1) for (x, y) in keypoints]
transformed = transform(image=image, keypoints=keypoints)
transformed_image = transformed['image']
transformed_keypoints = transformed['keypoints']
corners_distorted = [(x, y) for (x, y) in transformed_keypoints]
corners_distorted = np.array(corners_distorted)
keypoints = np.array(keypoints_orig)
fig, axs = plt.subplots(1, 2)
axs[0].imshow(image)
axs[0].scatter(keypoints[:, 0], keypoints[:, 1])
axs[1].imshow(transformed_image)
axs[1].scatter(corners_distorted[:, 0], corners_distorted[:, 1])
plt.show()

Expected behavior

The keypoints should be aligned with the internal checkers.

Actual behavior

The transformation between the image and the keypoints are not aligned

Screenshots

[RhysEvan]

I have found the origin of the issue but am sadly not able to think of a way to correct it...

What I did to specifically achieve this is go into albumentations>augmentations>geomtric>functional.py>remap_keypoints
The original code:

@handle_empty_array("keypoints")
def remap_keypoints(
    keypoints: np.ndarray,
    map_x: np.ndarray,
    map_y: np.ndarray,
    image_shape: tuple[int, int],
) -> np.ndarray:
    height, width = image_shape[:2]

    # Create inverse mappings
    x_inv = np.arange(width).reshape(1, -1).repeat(height, axis=0)
    y_inv = np.arange(height).reshape(-1, 1).repeat(width, axis=1)

    # Extract x and y coordinates
    x, y = keypoints[:, 0], keypoints[:, 1]

    # Clip coordinates to image boundaries
    x = np.clip(x, 0, width - 1, out=x)
    y = np.clip(y, 0, height - 1, out=y)

    # Convert to integer indices
    x_idx, y_idx = x.astype(int), y.astype(int)

    # Apply the inverse mapping
    new_x = x_inv[y_idx, x_idx] + (x - map_x[y_idx, x_idx])
    new_y = y_inv[y_idx, x_idx] + (y - map_y[y_idx, x_idx])

    # Clip the new coordinates to ensure they're within the image bounds
    new_x = np.clip(new_x, 0, width - 1, out=new_x)
    new_y = np.clip(new_y, 0, height - 1, out=new_y)

    # Create the transformed keypoints array
    return np.column_stack([new_x, new_y, keypoints[:, 2:]])

after "extensive" testing and math checking I simply decided to remove all the extras and just use the map_x and map_y for the reprojection:

@handle_empty_array("keypoints")
def remap_keypoints(
    keypoints: np.ndarray,
    map_x: np.ndarray,
    map_y: np.ndarray,
    image_shape: tuple[int, int],
) -> np.ndarray:
    height, width = image_shape[:2]

    # Create inverse mappings
    x_inv = np.arange(width).reshape(1, -1).repeat(height, axis=0)
    y_inv = np.arange(height).reshape(-1, 1).repeat(width, axis=1)
    # Extract x and y coordinates
    x, y = keypoints[:, 0], keypoints[:, 1]
    print("original; ", str(x[0]))
    # Clip coordinates to image boundaries
    x = np.clip(x, 0, width - 1, out=x)
    y = np.clip(y, 0, height - 1, out=y)
    print("clipped; ", str(x[0]))
    # Convert to integer indices
    x_idx, y_idx = x.astype(int), y.astype(int)
    print("inted; ", str(x[0]))
    # Apply the inverse mapping
    new_x = map_x[y_idx, x_idx]
    new_y = map_y[y_idx, x_idx]
    # new_x = x_inv[y_idx, x_idx] + (x - map_x[y_idx, x_idx])
    # new_y = y_inv[y_idx, x_idx] + (y - map_y[y_idx, x_idx])
    print("inv_x; ",str(x_inv[y_idx[0], x_idx[0]]))
    print("map_x; ", str(map_x[y_idx, x_idx][0]))
    print("new_x; ", str(new_x[0]))
    # Clip the new coordinates to ensure they're within the image bounds
    new_x = np.clip(new_x, 0, width - 1, out=new_x)
    new_y = np.clip(new_y, 0, height - 1, out=new_y)
    print("clipped_new_x; ", str(new_x[0]))

    # Create the transformed keypoints array
    return np.column_stack([new_x, new_y, keypoints[:, 2:]])

Which resulted in the inverted distortion result of all the checkpoints. g.e.:

When directly comparing that to the remap function that is actually applied to the image it is clear that cv2.remap must be doing something to invert the xy_map resulting in the inverting issue.

[ternaus]

@ThomasieDK

Thanks, love such bug reports. Show deep domain knowledge of the reporter and effort he put into debug.

Will look into it and fix.