Building an Albumentations pipeline

Setup#

pip install "albumentations==2.0.8"

Written and run against albumentations 2.0.8, numpy 2.2.6, opencv 5.0.0, torch 2.13.0 (CPU). Installing albumentations pulls in opencv-python-headless and albucore. ToTensorV2 also needs torch.

Version 2.0 changed things. always_apply was removed, and Compose gained a seed argument. Code written for 1.x can raise or warn. Pin your version.

One pipeline, four kinds of label#

import numpy as np, albumentations as A
from albumentations.pytorch import ToTensorV2

print("albumentations", A.__version__)
rng = np.random.default_rng(0)
image = rng.integers(0, 256, (60, 80, 3), dtype=np.uint8)      # HWC uint8: what Albumentations wants
mask  = np.zeros((60, 80), dtype=np.uint8); mask[10:30, 20:50] = 1
bboxes = [[20, 10, 50, 30]]                                     # pascal_voc: x_min y_min x_max y_max
labels = ["cat"]

tf = A.Compose(
    [
        A.HorizontalFlip(p=1.0),
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=1.0),
        A.Affine(rotate=(15, 15), p=1.0),
        A.Resize(48, 64),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ],
    bbox_params=A.BboxParams(format="pascal_voc", label_fields=["labels"], min_visibility=0.3),
    seed=0,
)

out = tf(image=image, mask=mask, bboxes=bboxes, labels=labels)
print("keys        :", sorted(out.keys()))
print("image       :", type(out["image"]).__name__, tuple(out["image"].shape), out["image"].dtype)
print("mask        :", type(out["mask"]).__name__, tuple(out["mask"].shape), out["mask"].dtype)
print("bboxes      :", [[round(float(v), 1) for v in b] for b in out["bboxes"]])
print("labels      :", out["labels"])
print("image mean  :", round(float(out["image"].mean()), 3))

albumentations 2.0.8
keys        : ['bboxes', 'image', 'labels', 'mask']
image       : Tensor (3, 48, 64) torch.float32
mask        : Tensor (48, 64) torch.uint8
bboxes      : [[20.1, 4.4, 47.5, 26.1]]
labels      : ['cat']
image mean  : 0.567

Reproduces exactly, because Compose(seed=0) was set.

Seven details in that output#

The transform is called with keyword arguments and returns a dict. tf(image=..., mask=..., bboxes=..., labels=...). Passing the image positionally raises. This is the single biggest difference from torchvision.

The image went in as HWC and came out as CHW. (60, 80, 3) numpy became (3, 48, 64) torch. ToTensorV2 does the transpose. Every other transform in the chain sees HWC numpy.

The mask did not become float and did not get normalised. It stayed uint8 with shape (48, 64). A.Normalize and ToTensorV2 know a mask is a label, not a picture. Class indices survive intact, which is what a segmentation loss needs.

The box moved through four transforms and came out sensible. It started at [20, 10, 50, 30] on an 80x60 canvas. It ended at [20.1, 4.4, 47.5, 26.1] on a 64x48 canvas, after a mirror, a rotation and a resize. Doing that by hand is the part people get wrong.

labels came back beside the boxes because it was named in label_fields. Anything listed there is filtered in lockstep with the boxes.

min_visibility=0.3 sets the survival rule. A box whose visible area drops below 30% of its original area is dropped.

image mean 0.567 is after normalisation with ImageNet statistics, on random noise. It is a sanity check that Normalize ran, not a meaningful quantity.

The behaviour that justifies the library#

import numpy as np, albumentations as A

rng = np.random.default_rng(0)
image = rng.integers(0, 256, (60, 80, 3), dtype=np.uint8)
bboxes = [[20, 10, 50, 30], [60, 40, 78, 58]]      # second box lives in the bottom-right
labels = ["cat", "dog"]

crop = A.Compose([A.Crop(x_min=0, y_min=0, x_max=55, y_max=35)],
                 bbox_params=A.BboxParams(format="pascal_voc",
                                          label_fields=["labels"],
                                          min_visibility=0.3))
out = crop(image=image, bboxes=bboxes, labels=labels)
print("boxes kept :", [[round(float(v)) for v in b] for b in out["bboxes"]])
print("labels kept:", out["labels"])
print("image now  :", out["image"].shape)

boxes kept : [[20, 10, 50, 30]]
labels kept: ['cat']
image now  : (35, 55, 3)

The crop kept the top-left 55x35 region. The dog sat at x from 60 to 78, entirely outside it. Its box went, and "dog" went with it.

Write this by hand and the usual bug is filtering the boxes while forgetting the labels. Then box zero is the cat and label zero is still "dog". Training proceeds, loss decreases, and the detector learns to call cats dogs.

Choosing a bbox format#

Declare the format your annotations are already in. Converting by hand before the pipeline is a step you do not need and a bug you do not want.

Common mistakes#

Feeding BGR images from cv2.imread and normalising with RGB statistics. OpenCV loads BGR. Fix: cv2.cvtColor(img, cv2.COLOR_BGR2RGB) right after reading, once, in your Dataset. See OpenCV.

Passing float images in the 0-255 range. Several transforms branch on dtype. Fix: keep uint8 until A.Normalize.

Copying 1.x code with always_apply=True. In 2.0.8 this does not raise. It emits UserWarning: Argument(s) 'always_apply' are not valid for transform BasicTransform, then ignores it. The transform falls back to its default p=0.5 and fires half the time. Your "always applied" resize now runs on half your images. Fix: p=1.0, and treat warnings from this library as errors in CI.

Forgetting label_fields. Boxes get filtered, your separate label list does not, and the two silently misalign. Fix: always pass every per-box array through label_fields.

Augmenting the validation set. Fix: build two Composes. The evaluation one holds resize, normalise and ToTensorV2 only.

Assuming reproducibility without asking for it. Fix: A.Compose([...], seed=0) for tests and debugging. Leave it unset in real training, or every epoch sees identical augmentations. See random seeds and reproducibility.

Try it yourself#

Change min_visibility=0.3 to 0.05 in the second script and re-run. Then set the crop to x_max=65 so the dog is half in frame. Find the visibility threshold at which the dog's box survives. Decide which value you would want on a real detector.

What to learn next#

• Test-time augmentation — using the same transforms at prediction time instead of training time.
• Object detection — the task whose labels make all of this necessary.
• Image segmentation — where the label is a full mask that has to move with the pixels.