Update AlbumentationsX augmentations for 2.3.7

pyproject.toml

@@ -53,7 +53,7 @@ dependencies = [
     "wandb>=0.24.0",
     "nvidia-ml-py>=13.590.48",
     "dynamic-network-architectures>=0.4.3",
-    "albumentationsx>=2.0.16",
+    "albumentationsx==2.3.7",
     "antspyx>=0.6.2",
     "torchmetrics>=1.8.2",
     "ipython>=8.37.0",

src/density_estimator/datasets/density_dataset.py

@@ -218,16 +218,18 @@ def get_transforms(
             [
                 # Safe 16-bit alternative to ColorJitter
                 A.RandomBrightnessContrast(
-                    brightness_limit=0.3, contrast_limit=0.3, p=0.7
+                    brightness_range=(-0.3, 0.3),
+                    contrast_range=(-0.3, 0.3),
+                    p=0.7,
                 ),
-                A.GaussianBlur(blur_limit=(3, 5), sigma_limit=(0.1, 1.5), p=0.3),
-                A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
+                A.GaussianBlur(blur_range=(3, 5), sigma_range=(0.1, 1.5), p=0.3),
+                A.GaussNoise(std_range=(0.2, 0.44), p=0.3),
                 A.CoarseDropout(
-                    max_holes=4,
-                    max_height=16,
-                    max_width=16,
-                    fill_value=fill_value,
-                    mask_fill_value=0,  # CRITICAL: Zeroes out the density target securely
+                    num_holes_range=(4, 4),
+                    hole_height_range=(16, 16),
+                    hole_width_range=(16, 16),
+                    fill=fill_value,
+                    fill_mask=0,  # Zero the corresponding density target and ROI mask.
                     p=0.3,
                 ),
             ]
@@ -241,8 +243,8 @@ def get_transforms(
             min_height=img_size,
             min_width=img_size,
             border_mode=cv2.BORDER_CONSTANT,
-            fill=fill_value,  # 'value' instead of 'fill' in newer albumentations
-            fill_mask=0,  # 'mask_value' instead of 'fill_mask'
+            fill=fill_value,
+            fill_mask=0,
         ),
         A.Normalize(mean=norm_mean, std=norm_std, max_pixel_value=fill_value),
         ToTensorV2(transpose_mask=True),  # Ensure mask is (C, H, W)

src/segmentation/finetuning/stardist_finetuning.py

@@ -1,15 +1,12 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import sys
+from glob import glob
 from pathlib import Path as StdPath
 
+import albumentations as A
 import matplotlib
 import numpy as np
-
-matplotlib.rcParams["image.interpolation"] = "none"
-
-from glob import glob
-
 from csbdeep.utils import Path, normalize
 from csbdeep.utils.tf import limit_gpu_memory
 from stardist import (
@@ -22,8 +19,22 @@
 from tifffile import imread
 from tqdm import tqdm
 
+matplotlib.rcParams["image.interpolation"] = "none"
+
 np.random.seed(42)
 lbl_cmap = random_label_cmap()
+augmentation = A.Compose(
+    [
+        A.SquareSymmetry(p=1),
+        A.RandomBrightnessContrast(
+            brightness_range=(-0.2, 0.2),
+            contrast_range=(-0.4, 1.0),
+            p=1,
+        ),
+        A.GaussNoise(std_range=(0, 0.02), per_channel=True, p=1),
+    ],
+    seed=42,
+)
 
 X = sorted(
     glob(
@@ -52,7 +63,7 @@
     )
     sys.stdout.flush()
 
-X = [normalize(x, 1, 99.8, axis=axis_norm) for x in tqdm(X)]
+X = [normalize(x, 1, 99.8, axis=axis_norm, clip=True) for x in tqdm(X)]
 Y = [fill_label_holes(y) for y in tqdm(Y)]
 
 assert len(X) > 1, "not enough training data"
@@ -100,35 +111,13 @@
     )
 
 
-def random_fliprot(img, mask):
-    assert img.ndim >= mask.ndim
-    axes = tuple(range(mask.ndim))
-    perm = tuple(np.random.permutation(axes))
-    img = img.transpose(perm + tuple(range(mask.ndim, img.ndim)))
-    mask = mask.transpose(perm)
-    for ax in axes:
-        if np.random.rand() > 0.5:
-            img = np.flip(img, axis=ax)
-            mask = np.flip(mask, axis=ax)
-    return img, mask
-
-
-def random_intensity_change(img):
-    img = img * np.random.uniform(0.6, 2) + np.random.uniform(-0.2, 0.2)
-    return img
-
-
 def augmenter(x, y):
     """Augmentation of a single input/label image pair.
     x is an input image
     y is the corresponding ground-truth label image
     """
-    x, y = random_fliprot(x, y)
-    x = random_intensity_change(x)
-    # add some gaussian noise
-    sig = 0.02 * np.random.uniform(0, 1)
-    x = x + sig * np.random.normal(0, 1, x.shape)
-    return x, y
+    augmented = augmentation(image=x, mask=y)
+    return augmented["image"], augmented["mask"]
 
 
 model.train(X_trn, Y_trn, validation_data=(X_val, Y_val), augmenter=augmenter)