From af67ad0f428847964fe6bfd2dcb0f8828c36e398 Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Wed, 29 Jul 2026 18:37:06 +0300 Subject: [PATCH 1/3] Fix AlbumentationsX augmentation parameters --- pyproject.toml | 2 +- .../datasets/density_dataset.py | 22 ++++++++++--------- uv.lock | 15 +++++++------ 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bda889b..446ab7e 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/density_estimator/datasets/density_dataset.py b/src/density_estimator/datasets/density_dataset.py index 2077f53..c227021 100644 --- a/src/density_estimator/datasets/density_dataset.py +++ b/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) diff --git a/uv.lock b/uv.lock index 3172125..5299904 100644 --- a/uv.lock +++ b/uv.lock @@ -45,7 +45,7 @@ wheels = [ [[package]] name = "albucore" -version = "0.1.5" +version = "0.2.5" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numkong" }, @@ -53,17 +53,18 @@ dependencies = [ { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "stringzilla" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7d/23/d1eff908effdec4d6c79ef8566978f3cb55435ad19de0f090ebcd7794923/albucore-0.1.5.tar.gz", hash = "sha256:7f188c437405dbe132c168cb291f3b015f5f6c56d6c333277ab42ba635270ff4", size = 73106, upload-time = "2026-04-07T05:18:27.183Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ac/69/c5846a4e520935a8125c5a83be11326eae838a1dfb280c7b38fa1b9c5f33/albucore-0.2.5.tar.gz", hash = "sha256:6653da52223987af3bcf264ff4c1672595ed8c13b6f66ffd28c5a353a68cb43f", size = 153595, upload-time = "2026-07-22T13:55:02.478Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/10/0a/454d887ff1e19f7c7729f88acebfadc4b3e428d0172f1f385cd3d88dd6e1/albucore-0.1.5-py3-none-any.whl", hash = "sha256:13948e6bb80d55715ce2d17e47831f8f4e934d72e65bb2289ab69d5a2eda265a", size = 42044, upload-time = "2026-04-07T05:18:25.826Z" }, + { url = "https://files.pythonhosted.org/packages/6e/ef/dd7fddfc3ce4d25d51923eb68a09e9a29e23740f731d1b9116d99c338f8a/albucore-0.2.5-py3-none-any.whl", hash = "sha256:946a28bbe1c6078a72ef9220955239db7d4b4fec4ba90b552c25e3daa5d2f2ae", size = 42535, upload-time = "2026-07-22T13:55:00.939Z" }, ] [[package]] name = "albumentationsx" -version = "2.1.3" +version = "2.3.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "albucore" }, + { name = "numkong" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pydantic" }, @@ -71,9 +72,9 @@ dependencies = [ { name = "scipy" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7d/f3/40690cfde37f083b57a675bcc9ad1ba5857a50ea6f42a331a010588461e7/albumentationsx-2.1.3.tar.gz", hash = "sha256:143387bff9d21f7a11b124b01879e78d709e06765af5b32ea6b0f33caa46f581", size = 541986, upload-time = "2026-04-08T05:59:07.081Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/c2/83e17e9f8a554220a13458e5d0b06bb800a945ebfc5bda26399f0fda0cca/albumentationsx-2.3.7.tar.gz", hash = "sha256:ba0f9dfb33efbd893fe37a41176f5190df15d658840d73b87fc2da879bad7c96", size = 790168, upload-time = "2026-07-29T13:36:01.061Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/fc/efba215e0ed09d860e757c10a34a473d27ff0d8fbebf32a820bef584efe1/albumentationsx-2.1.3-py3-none-any.whl", hash = "sha256:5a4b6fa59f9b9458d57255cd26da6c5b951a45acf30b4e8dab50d381ba0540c5", size = 470192, upload-time = "2026-04-08T05:59:05.301Z" }, + { url = "https://files.pythonhosted.org/packages/df/e9/ff8d5103ae29fdb01e32d2e06272edcd6ec05761d59830d7959ff63076dc/albumentationsx-2.3.7-py3-none-any.whl", hash = "sha256:85a43a2cc1d329a4d4fdff3534b9fa6c545ae729f716f99287938bd7a53a41be", size = 563487, upload-time = "2026-07-29T13:35:59.359Z" }, ] [[package]] @@ -2838,7 +2839,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "albumentationsx", specifier = ">=2.0.16" }, + { name = "albumentationsx", specifier = "==2.3.7" }, { name = "antspyx", specifier = ">=0.6.2" }, { name = "autograd", specifier = ">=1.8.0" }, { name = "cellpose", extras = ["gui"], git = "https://www.github.com/mouseland/cellpose.git" }, From 7c31617b747844eaabe29583d1ae4a2f785b34bb Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Wed, 29 Jul 2026 18:57:53 +0300 Subject: [PATCH 2/3] Use SquareSymmetry for StarDist augmentation --- .../finetuning/stardist_finetuning.py | 26 +++++-------------- 1 file changed, 7 insertions(+), 19 deletions(-) diff --git a/src/segmentation/finetuning/stardist_finetuning.py b/src/segmentation/finetuning/stardist_finetuning.py index 69fb34e..4895480 100644 --- a/src/segmentation/finetuning/stardist_finetuning.py +++ b/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,11 @@ from tifffile import imread from tqdm import tqdm +matplotlib.rcParams["image.interpolation"] = "none" + np.random.seed(42) lbl_cmap = random_label_cmap() +geometric_augmenter = A.Compose([A.SquareSymmetry(p=1)], seed=42) X = sorted( glob( @@ -100,19 +100,6 @@ ) -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 @@ -123,7 +110,8 @@ def augmenter(x, y): x is an input image y is the corresponding ground-truth label image """ - x, y = random_fliprot(x, y) + augmented = geometric_augmenter(image=x, mask=y) + x, y = augmented["image"], augmented["mask"] x = random_intensity_change(x) # add some gaussian noise sig = 0.02 * np.random.uniform(0, 1) From 8ea099e2a03cdfddb4db946a10cbc9a994ce02bc Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Wed, 29 Jul 2026 19:27:12 +0300 Subject: [PATCH 3/3] Use native StarDist augmentation pipeline --- .../finetuning/stardist_finetuning.py | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/src/segmentation/finetuning/stardist_finetuning.py b/src/segmentation/finetuning/stardist_finetuning.py index 4895480..59363f8 100644 --- a/src/segmentation/finetuning/stardist_finetuning.py +++ b/src/segmentation/finetuning/stardist_finetuning.py @@ -23,7 +23,18 @@ np.random.seed(42) lbl_cmap = random_label_cmap() -geometric_augmenter = A.Compose([A.SquareSymmetry(p=1)], seed=42) +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,23 +111,13 @@ ) -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 """ - augmented = geometric_augmenter(image=x, mask=y) - x, y = augmented["image"], augmented["mask"] - 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)