[Enhancement] Simplify number13's synchronized targets and square symmetries

[ternaus]

Hi,

Thank you for publishing the full SpaceNet 8 winning solutions. Albumentations is doing the synchronization work for the pre/post imagery and segmentation masks in 02-number13; AlbumentationsX 2.3.5 can express those variable target lists directly and collapse the orientation block to one transform.

The checked-in trainer uses SN8Dataset with one image and two masks for foundation training, then two images and one mask for flood training:

if training_type == 'foundation':
    data_to_load = ["preimg", "building", "road"]
    gens = ds.get_generators_sn8(config, data_to_load=data_to_load)

if training_type == 'flood':
    data_to_load = ["preimg", "postimg", "flood"]
    gens = ds.get_generators_sn8(config, data_to_load=data_to_load, do_imbalanced=True)

The dataset builds image0, mask0, and mask1 registrations from that list:

if 'preimg' in data_to_load and 'postimg' in data_to_load:
    additional_targets = {'image0': 'image'}
    for i in range(len(data_to_load)-3):
        additional_targets[f'mask{i}'] = 'mask'

    self.aug_train = Compose(tx, additional_targets=additional_targets)
else:
    additional_targets = {}
    for i in range(len(data_to_load)-2):
        additional_targets[f'mask{i}'] = 'mask'
    self.aug_train = Compose(tx, additional_targets=additional_targets)

It then spells out every supported image/mask count at the call site:

def augment_train(self, img, mask):
    if len(img)==1:
        if len(mask)==1:
            augmented = self.aug_train(image=img[0], mask=mask[0])
        elif len(mask)==2:
            augmented = self.aug_train(image=img[0], mask=mask[0], mask0=mask[1])
        elif len(mask)==3:
            augmented = self.aug_train(image=img[0], mask=mask[0], mask0=mask[1], mask1=mask[2])
    else:
        if len(mask)==1:
            augmented = self.aug_train(image=img[0], image0=img[1], mask=mask[0])
        elif len(mask)==2:
            augmented = self.aug_train(image=img[0], image0=img[1], mask=mask[0], mask0=mask[1])
        elif len(mask)==3:
            augmented = self.aug_train(image=img[0], image0=img[1], mask=mask[0], mask0=mask[1], mask1=mask[2])
    return augmented

The returned dictionary is split back into image and mask lists by inspecting each generated key:

augmented = self.augment_train(returned_images, returned_masks)
returned_images = []
returned_masks = []
for k in augmented.keys():
    if 'image' in k:
        returned_images.append(self.tr(Image.fromarray(augmented[k])))
    else:
        mask = augmented[k]
        if len(mask.shape) == 2:
            mask = mask[:, :, np.newaxis]
        returned_masks.append(np.transpose(mask, (2, 0, 1)).astype('float32'))

That routing preserves one sampled crop, orientation, and color change across every related input. In 2.3.5, the native images and masks targets carry the same contract as stacked NumPy arrays, so the branches and generated key names are unnecessary.

There is a second simplification in the augmentation list:

tx = [
    ShiftScaleRotate(shift_limit=0.2, scale_limit=0, rotate_limit=0),
    OneOf([RandomSizedCrop(min_max_height=(int(height * 0.8), int(height * 1.2)), w2h_ratio=1., height=height,
                           width=width, p=0.9),
           RandomCrop(height=height, width=width, p=0.1)], p=1),
    Rotate(limit=10, p=0.2, border_mode=cv2.BORDER_CONSTANT, value=0),
    HorizontalFlip(),
    RandomRotate90(),
    Transpose(),
    OneOf([RGBShift(), RandomBrightnessContrast(), RandomGamma(), ColorJitter()], p=0.5),
]
if not has_offnadir:
    tx.append(VerticalFlip())

The three independent p=0.5 operations—HorizontalFlip, Transpose, and the final VerticalFlip—map their eight boolean combinations one-to-one onto the eight elements of the square-symmetry group for every fixed outcome of RandomRotate90. The resulting orientation is therefore uniform even if the quarter-turn step’s apply probability differs between package versions. SquareSymmetry(p=1) expresses that same uniform eight-element distribution in one step.

With the current 2.3.5 argument names, the pipeline and dataset path can be written as:

from albumentations import (
    Affine, ColorJitter, Compose, OneOf, RandomBrightnessContrast,
    RandomCrop, RandomGamma, RandomSizedCrop, RGBShift, Rotate,
    SquareSymmetry,
)

tx = [
    Affine(
        translate_percent=(-0.2, 0.2),
        scale=(1.0, 1.0),
        rotate=(0.0, 0.0),
        shear=(0.0, 0.0),
        fill=0,
        fill_mask=0,
    ),
    OneOf([
        RandomSizedCrop(
            min_max_height=(int(height * 0.8), int(height * 1.2)),
            size=(height, width),
            w2h_ratio=1.0,
            p=0.9,
        ),
        RandomCrop(height=height, width=width, p=0.1),
    ], p=1),
    Rotate(
        angle_range=(-10, 10),
        border_mode=cv2.BORDER_CONSTANT,
        fill=0,
        fill_mask=0,
        p=0.2,
    ),
    SquareSymmetry(p=1),
    OneOf([RGBShift(), RandomBrightnessContrast(), RandomGamma(), ColorJitter()], p=0.5),
]

self.aug_train = Compose(tx)

def augment_train(self, img, mask):
    return self.aug_train(
        images=np.stack(img),
        masks=np.stack(mask),
    )

The output handling becomes explicit about the two plural targets:

augmented = self.augment_train(returned_images, returned_masks)

returned_images = [
    self.tr(Image.fromarray(image))
    for image in augmented["images"]
]
returned_masks = [
    np.transpose(
        mask[:, :, np.newaxis] if mask.ndim == 2 else mask,
        (2, 0, 1),
    ).astype("float32")
    for mask in augmented["masks"]
]

The sequence guide documents the images=np.ndarray and masks=np.ndarray call: parameters are sampled once and applied consistently along the first dimension. This applies to paired pre/post satellite images as well as video frames.

I tested the published PyPI wheel and tag albumentationsx==2.3.5 with 640×640 RGB uint8 inputs and the repository’s complete migrated augmentation list. For all six supported combinations of one or two images with one, two, or three masks, the named-target and stacked-target routes matched element-for-element across 32 seeds: 192 full-pipeline comparisons. Another 32 runs preserved the flood path’s two RGB images and one two-channel mask as (2, 512, 512, 3) and (1, 512, 512, 2) uint8 arrays. I also enumerated the 32 deterministic geometry paths formed by four fixed quarter-turns and the eight flip/transpose choices: every square-symmetry element occurred four times, with all eight appearing once for each fixed quarter-turn. The current Affine shift-only form matched the corresponding 2.3.5 ShiftScaleRotate form across 32 seeds.

The transform class names in the original 2022 code did not all change. The wider 2.3.5 constructor updates shown above are the recommended shift-only Affine expression, height and width → size for RandomSizedCrop, and limit/value → angle_range/fill for Rotate. The full training and validation path should still be tested with the original data and models.

The requirements file installs the legacy albumentations package without a version pin, so I did not open a dependency-changing PR.

pip uninstall albumentations
pip install -U albumentationsx

The Python module path remains albumentations, so the existing from albumentations import ... import style stays the same.

The packages use different licenses: the legacy albumentations package 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.

I’m glad Albumentations was useful in this solution. 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.