[Enhancement] Replace the imgaug-backed training transforms with native transforms

[ternaus]

Hi,

Albumentations has been part of LaMa's training recipe since the original release, and I am genuinely glad it was useful there. The active distortions path still needs imgaug only for IAAPerspective2 and IAAAffine2; AlbumentationsX 2.3.5 covers both directly with Perspective and Affine.

The README's Places365 training command selects lama-fourier, whose data default selects abl-04-256-mh-dist; that configuration activates transform_variant: distortions. The same route is shown for custom-dataset training in the README.

The current IAAAffine2 constructor and processor are:

def __init__(
    self,
    scale=(0.7, 1.3),
    translate_percent=None,
    translate_px=None,
    rotate=0.0,
    shear=(-0.1, 0.1),
    order=1,
    cval=0,
    mode="reflect",
    always_apply=False,
    p=0.5,
):
    super(IAAAffine2, self).__init__(always_apply, p)
    self.scale = dict(x=scale, y=scale)
    self.translate_percent = to_tuple(translate_percent, 0)
    self.translate_px = to_tuple(translate_px, 0)
    self.rotate = to_tuple(rotate)
    self.shear = dict(x=shear, y=shear)
    self.order = order
    self.cval = cval
    self.mode = mode

@property
def processor(self):
    return iaa.Affine(
        self.scale,
        self.translate_percent,
        self.translate_px,
        self.rotate,
        self.shear,
        self.order,
        self.cval,
        self.mode,
    )

def get_transform_init_args_names(self):
    return ("scale", "translate_percent", "translate_px", "rotate", "shear", "order", "cval", "mode")

The current IAAPerspective2 constructor and processor are:

def __init__(self, scale=(0.05, 0.1), keep_size=True, always_apply=False, p=0.5,
             order=1, cval=0, mode="replicate"):
    super(IAAPerspective2, self).__init__(always_apply, p)
    self.scale = to_tuple(scale, 1.0)
    self.keep_size = keep_size
    self.cval = cval
    self.mode = mode

@property
def processor(self):
    return iaa.PerspectiveTransform(self.scale, keep_size=self.keep_size, mode=self.mode, cval=self.cval)

def get_transform_init_args_names(self):
    return ("scale", "keep_size")

They are called at the start of the active pipeline:

elif transform_variant == 'distortions':
    transform = A.Compose([
        IAAPerspective2(scale=(0.0, 0.06)),
        IAAAffine2(scale=(0.7, 1.3),
                   rotate=(-40, 40),
                   shear=(-0.1, 0.1)),
        A.PadIfNeeded(min_height=out_size, min_width=out_size),
        A.OpticalDistortion(),
        A.RandomCrop(height=out_size, width=out_size),
        A.HorizontalFlip(),
        A.CLAHE(),
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2),
        A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=30, val_shift_limit=5),
        A.ToFloat()
    ])

The native replacement can preserve the configured geometry and the old reflective padding while updating the argument names used elsewhere in this branch:

-from saicinpainting.training.data.aug import IAAAffine2, IAAPerspective2

 elif transform_variant == 'distortions':
     transform = A.Compose([
-        IAAPerspective2(scale=(0.0, 0.06)),
-        IAAAffine2(scale=(0.7, 1.3),
-                   rotate=(-40, 40),
-                   shear=(-0.1, 0.1)),
-        A.PadIfNeeded(min_height=out_size, min_width=out_size),
-        A.OpticalDistortion(),
+        A.Perspective(scale=(0.0, 0.06), keep_size=True, border_mode=cv2.BORDER_REPLICATE),
+        A.Affine(scale=(0.7, 1.3), rotate=(-40, 40), shear=(-0.1, 0.1), keep_ratio=False, border_mode=cv2.BORDER_REFLECT_101),
+        A.PadIfNeeded(min_height=out_size, min_width=out_size, border_mode=cv2.BORDER_REFLECT_101),
+        A.OpticalDistortion(distort_range=(-0.05, 0.05), border_mode=cv2.BORDER_REFLECT_101),
         A.RandomCrop(height=out_size, width=out_size),
         A.HorizontalFlip(),
         A.CLAHE(),
-        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2),
-        A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=30, val_shift_limit=5),
+        A.RandomBrightnessContrast(brightness_range=(-0.2, 0.2), contrast_range=(-0.2, 0.2)),
+        A.HueSaturationValue(hue_shift_range=(-5, 5), sat_shift_range=(-30, 30), val_shift_range=(-5, 5)),
         A.ToFloat()
     ])

keep_ratio=False is the only non-default scale setting needed here: Affine expands the tuple range for both axes and samples X and Y independently. The shear tuple is also sampled independently for both axes. Its default image interpolation is cv2.INTER_LINEAR, matching the old order=1; imgaug's "reflect" maps to cv2.BORDER_REFLECT_101, and "replicate" maps to cv2.BORDER_REPLICATE. Both transforms keep their existing p=0.5 defaults. Once the other transform variants are migrated too, saicinpainting/training/data/aug.py and the direct imgaug==0.4.0 environment dependency can be removed.

This addresses both compatibility failures already reported in the tracker. #301 hits the removed DualIAATransform import with a newer Albumentations, while #345 reaches NumPy 2's removed np.sctypes API through Albumentations 0.5.2's import of imgaug. Pinning the old packages works around one side at a time; replacing the two IAA wrappers removes the dependency that causes both failures.

I tested the original code with albumentations==0.5.2 and imgaug==0.4.0, then tested the replacement against the published albumentationsx==2.3.5 wheel and matching tag. Across 128 seeds each, RGB uint8 inputs shaped 256×256, 317×509, and 512×512 preserved their shape and dtype through each old and new geometric transform. The complete updated distortions branch produced 256×256×3 float32 in [0, 1] for all 384 calls.

A separate 4,096-seed check kept X/Y scale within 0.7..1.3, rotation within -40..40, both shear axes within -0.1..0.1, and perspective scale within 0..0.06; X and Y scale were sampled independently, and all four apply/skip combinations occurred. At the same 4,096 seeds, the compact tuple form and the expanded X/Y dictionaries produced identical applied configurations, matrices, and images. The random-number stream and sampled pixels change after migration, so an old and a new seeded run should not be expected to match element-for-element.

The active training datasets pass one RGB image to the transform and generate the inpainting mask afterward, so the test above covers that path. Any future use of these helpers with masks, boxes, or keypoints should be tested separately during migration.

The repository pins albumentations==0.5.2, and the Conda environment also pins imgaug==0.4.0, so I did not open a dependency-change PR. A complete migration should apply the same current argument names to the other transform variants, including RandomScale, RandomBrightnessContrast, and HueSaturationValue, and preserve the old reflective defaults for PadIfNeeded and OpticalDistortion before retraining comparisons.

If you want to try the maintained package, the Python import remains import albumentations as A:

pip uninstall albumentations
pip install -U albumentationsx

The packages use different licenses: albumentations==0.5.2 is MIT, while albumentationsx==2.3.5 is AGPL-3.0-only. The license guide explains the terms. Albumentations, LLC also offers commercial licenses with alternative terms.

If you have feedback, complaints, or proposals for AlbumentationsX, please open an issue. I read the tracker every day.

If this note is useful, a star or sponsorship would mean a lot.