refactor: train and more aggressive augmentation

src/aind_exaspim_neuron_segmentation/machine_learning/augmentation.py

@@ -148,7 +148,7 @@ def __call__(self, input_img, label_mask):
             Rotated label mask.
         """
         for axes in self.axes:
-            if random.random() > 0.2:
+            if random.random() <= 0.8:
                 angle = random.uniform(*self.angles)
                 input_img = rotate3d(input_img, angle, axes)
                 label_mask = rotate3d(label_mask, angle, axes)
@@ -249,15 +249,15 @@ class RandomNoise3D:
     Adds random Gaussian noise to a 3D image.
     """
 
-    def __init__(self, max_std=0.15):
+    def __init__(self, max_std=0.16):
         """
         Initializes a RandomNoise3D transformer.
 
         Parameters
         ----------
         max_std : float, optional
             Maximum standard deviation of the Gaussian noise distribution.
-            Default is 0.05.
+            Default is 0.2.
         """
         self.max_std = max_std
 
@@ -275,7 +275,7 @@ def __call__(self, img):
         numpy.ndarray
             Noisy image.
         """
-        std = self.max_std * random.random()
+        std = random.uniform(0, self.max_std)
         noise = np.random.normal(0, std, img.shape)
         return img + noise
 
@@ -285,15 +285,15 @@ class RandomSmooth3D:
     Applies Gaussian smoothing to a 3D image.
     """
 
-    def __init__(self, max_sigma=1.0):
+    def __init__(self, max_sigma=0.8):
         """
         Initializes a GaussianSmooth3D transformer.
 
         Parameters
         ----------
         max_sigma : float, optional
             Maximum standard deviation of the Gaussian kernel.
-            Default is 1.0.
+            Default is 1.
         """
         self.max_sigma = max_sigma
 
@@ -311,7 +311,7 @@ def __call__(self, img):
         numpy.ndarray
             Smoothed image.
         """
-        sigma = self.max_sigma * random.random()
+        sigma = random.uniform(0, self.max_sigma)
         return gaussian_filter(img, sigma=sigma)
 
 

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=500,
+        brightness_clip=300,
         normalization_percentiles=(1, 99.9),
         patch_shape=(96, 96, 96),
     ):
@@ -162,7 +162,7 @@ def __init__(
         input_img_paths,
         label_mask_paths,
         affinity_mode=True,
-        brightness_clip=1000,
+        brightness_clip=300,
         normalization_percentiles=(1, 99.9),
         patch_shape=(96, 96, 96),
         transform=None
@@ -180,7 +180,7 @@ def __init__(
             If True, the model predicts affinities; if False, it predicts
             foreground–background. Default is True.
         brightness_clip : float, optional
-            Maximum brightness value for voxel intensities. Default is 1000.
+            Maximum brightness value for voxel intensities. Default is 300.
         normalization_percentiles, Tuple[int]
             Lower and upper percentiles used for normalization. Default is
             (1, 99.5).
@@ -321,7 +321,7 @@ def __init__(
         input_img_paths,
         label_mask_paths,
         affinity_mode=True,
-        brightness_clip=1000,
+        brightness_clip=300,
         normalization_percentiles=(1, 99.5),
         patch_shape=(96, 96, 96),
     ):
@@ -338,7 +338,7 @@ def __init__(
             If True, the model predicts affinities; if False, it predicts
             foreground–background. Default is True.
         brightness_clip : float, optional
-            Maximum brightness value for voxel intensities. Default is 1000.
+            Maximum brightness value for voxel intensities. Default is 300.
         normalization_percentiles, Tuple[int]
             Lower and upper percentiles used for normalization. Default is
             (1, 99.5).