refactor: more aggressive data augmentation

src/aind_exaspim_neuron_segmentation/machine_learning/augmentation.py

@@ -8,7 +8,7 @@
 
 """
 
-from scipy.ndimage import rotate, zoom
+from scipy.ndimage import gaussian_filter, rotate, zoom
 
 import numpy as np
 import random
@@ -31,8 +31,11 @@ def __init__(self):
             RandomFlip3D(),
             RandomRotation3D(),
         ]
-        self.intensity_transforms = transforms.Compose(
-                [RandomContrast3D(), RandomNoise3D()]
+        self.intensity_transforms1 = transforms.Compose(
+                [RandomNoise3D(), RandomContrast3D()]
+            )
+        self.intensity_transforms2 = transforms.Compose(
+                [RandomSmooth3D(), RandomContrast3D()]
             )
 
     def __call__(self, input_img, label_mask):
@@ -58,7 +61,10 @@ def __call__(self, input_img, label_mask):
             input_img, label_mask = transform(input_img, label_mask)
 
         # Intensity transforms
-        input_img = self.intensity_transforms(input_img)
+        if random.random() < 0.5:
+            input_img = self.intensity_transforms1(input_img)
+        else:
+            input_img = self.intensity_transforms2(input_img)
         return input_img, label_mask
 
 
@@ -142,7 +148,7 @@ def __call__(self, input_img, label_mask):
             Rotated label mask.
         """
         for axes in self.axes:
-            if random.random() > 0.4:
+            if random.random() > 0.2:
                 angle = random.uniform(*self.angles)
                 input_img = rotate3d(input_img, angle, axes)
                 label_mask = rotate3d(label_mask, angle, axes)
@@ -243,7 +249,7 @@ class RandomNoise3D:
     Adds random Gaussian noise to a 3D image.
     """
 
-    def __init__(self, max_std=0.05):
+    def __init__(self, max_std=0.15):
         """
         Initializes a RandomNoise3D transformer.
 
@@ -269,11 +275,46 @@ def __call__(self, img):
         numpy.ndarray
             Noisy image.
         """
-        std = self.max_std * random.random()
+        std = self.max_std #* random.random()
         noise = np.random.normal(0, std, img.shape)
         return img + noise
 
 
+class RandomSmooth3D:
+    """
+    Applies Gaussian smoothing to a 3D image.
+    """
+
+    def __init__(self, max_sigma=1.0):
+        """
+        Initializes a GaussianSmooth3D transformer.
+
+        Parameters
+        ----------
+        max_sigma : float, optional
+            Maximum standard deviation of the Gaussian kernel.
+            Default is 1.0.
+        """
+        self.max_sigma = max_sigma
+
+    def __call__(self, img):
+        """
+        Applies Gaussian smoothing to an image.
+
+        Parameters
+        ----------
+        img : numpy.ndarray
+            3D image to smooth.
+
+        Returns
+        -------
+        numpy.ndarray
+            Smoothed image.
+        """
+        sigma = self.max_sigma #* random.random()
+        return gaussian_filter(img, sigma=sigma)
+
+
 # --- Helpers ---
 def rotate3d(img, angle, axes):
     """

src/aind_exaspim_neuron_segmentation/machine_learning/data_handling.py

@@ -28,7 +28,7 @@ def __init__(
         input_img_paths,
         label_mask_paths,
         affinity_mode=True,
-        brightness_clip=1000,
+        brightness_clip=500,
         normalization_percentiles=(1, 99.9),
         patch_shape=(96, 96, 96),
     ):
@@ -48,7 +48,7 @@ def __init__(
             Maximum brightness value for voxel intensities. Default is 1000.
         normalization_percentiles, Tuple[int]
             Lower and upper percentiles used for normalization. Default is
-            (1, 99.5).
+            (1, 99.9).
         patch_shape : Tuple[int], optional
             Shape of 3D patches to extract from the images. Default is
             (96, 96, 96).
@@ -124,13 +124,11 @@ def get_input_patch(self, i, center):
         patch : numpy.ndarray
             Normalized input patch with shape (1, D, H, W).
         """
-        # Get image patch
         patch = self.get_patch(self.input_imgs[i], center)
         patch = np.minimum(patch, self.brightness_clip)
-
-        # Normalize image patch
-        mn, mx = np.percentile(patch, self.normalization_percentiles)
-        patch = np.clip((patch - mn) / (mx - mn + 1e-8), 0, 1)
+        patch = img_util.normalize(
+            patch, percentiles=self.normalization_percentiles
+        )
         return patch
 
     def get_label_patch(self, i, center):