Refactor more augmentation (#557)

* refactor: more aggressive image augmentation

* refactor: more aggressive image augmentation

---------

Co-authored-by: anna-grim <[email protected]>

Author: anna-grim

src/neuron_proofreader/machine_learning/augmentation.py

@@ -84,17 +84,11 @@ def __call__(self, patches):
         patches : numpy.ndarray
             Image with the shape (2, H, W, D), where "patches[0, ...]" is from
             the input image and "patches[1, ...]" is from the segmentation.
-
-        Returns
-        -------
-        patches : numpy.ndarray
-            Flipped 3D image and segmentation patch.
         """
         for axis in self.axes:
             if random.random() > 0.5:
                 patches[0, ...] = np.flip(patches[0, ...], axis=axis)
                 patches[1, ...] = np.flip(patches[1, ...], axis=axis)
-        return patches
 
 
 class RandomRotation3D:
@@ -125,18 +119,12 @@ def __call__(self, patches):
         patches : numpy.ndarray
             Image with the shape (2, H, W, D), where "patches[0, ...]" is from
             the input image and "patches[1, ...]" is from the segmentation.
-
-        Returns
-        -------
-        patches : numpy.ndarray
-            Rotated 3D image and segmentation patch.
         """
         for axes in self.axes:
             if random.random() < 0.5:
                 angle = random.uniform(*self.angles)
                 patches[0, ...] = rotate3d(patches[0, ...], angle, axes)
                 patches[1, ...] = rotate3d(patches[1, ...], angle, axes, True)
-        return patches
 
 
 class RandomScale3D:
@@ -216,30 +204,25 @@ def __call__(self, img_patch):
         ----------
         img_patch : numpy.ndarray
             Image to which contrast will be added.
-
-        Returns
-        -------
-        numpy.ndarray
-            Contrasted 3D image.
         """
         factor = random.uniform(*self.factor_range)
-        return np.clip(img_patch * factor, 0, 1)
+        img_patch = np.clip(img_patch * factor, 0, 1)
 
 
 class RandomNoise3D:
     """
     Adds random Gaussian noise to a 3D image.
     """
 
-    def __init__(self, max_std=0.2):
+    def __init__(self, max_std=0.3):
         """
         Initializes a RandomNoise3D transformer.
 
         Parameters
         ----------
         max_std : float, optional
             Maximum standard deviation of the Gaussian noise distribution.
-            Default is 0.16.
+            Default is 0.3.
         """
         self.max_std = max_std
 
@@ -251,16 +234,10 @@ def __call__(self, img_patch):
         ----------
         img_patch : np.ndarray
             Image to which noise will be added.
-
-        Returns
-        -------
-        img_patch : numpy.ndarray
-            Noisy 3D image.
         """
         std = self.max_std * random.random()
-        noise = np.random.uniform(-std, std, img_patch.shape)
-        img_patch += noise
-        return np.clip(img_patch, 0, 1)
+        img_patch += np.random.uniform(-std, std, img_patch.shape)
+        img_patch = np.clip(img_patch, 0, 1)
 
 
 # --- Helpers ---

src/neuron_proofreader/merge_proofreading/merge_datasets.py

@@ -446,10 +446,10 @@ def get_img_patch(self, brain_id, center):
         -------
         img_patch : numpy.ndarray
             Extracted image patch, which has been normalized and clipped to a
-            maximum value of 300.
+            maximum value of 400.
         """
         img_patch = self.img_readers[brain_id].read(center, self.patch_shape)
-        img_patch = img_util.normalize(np.minimum(img_patch, 300))
+        img_patch = img_util.normalize(np.minimum(img_patch, 400))
         return img_patch
 
     def get_segment_mask(self, subgraph):
@@ -953,11 +953,11 @@ def _load_multimodal_batch(self, batch_idxs):
             # Store results
             patches = np.zeros((len(batch_idxs),) + self.patches_shape)
             labels = np.zeros((len(batch_idxs), 1))
-            point_clouds = np.zeros((len(batch_idxs), 3, 3600))
+            point_clouds = np.zeros((len(batch_idxs), 3600, 3))
             for thread in as_completed(pending.keys()):
                 i = pending.pop(thread)
                 patches[i], subgraph, labels[i] = thread.result()
-                point_clouds[i] = subgraph_to_point_cloud(subgraph) #.T
+                point_clouds[i] = subgraph_to_point_cloud(subgraph).T
 
         # Set batch dictionary
         batch = ml_util.TensorDict(

src/neuron_proofreader/merge_proofreading/merge_inference.py

@@ -319,7 +319,7 @@ def get_batch(self, img, offset, patch_centers, nodes):
         batch = np.empty((len(patch_centers), 2,) + self.patch_shape)
         for i, center in enumerate(patch_centers):
             s = img_util.get_slices(center, self.patch_shape)
-            batch[i, 0, ...] = img_util.normalize(np.minimum(img[s], 300))
+            batch[i, 0, ...] = img_util.normalize(np.minimum(img[s], 400))
             batch[i, 1, ...] = label_mask[s]
         return nodes, torch.tensor(batch, dtype=torch.float)
 
@@ -331,10 +331,10 @@ def get_multimodal_batch(self, img, offset, patch_centers, nodes):
 
         # Populate batch array
         patches = np.empty((batch_size, 2,) + self.patch_shape)
-        point_clouds = np.empty((batch_size, 3, 3200), dtype=np.float32)
+        point_clouds = np.empty((batch_size, 3, 3600), dtype=np.float32)
         for i, (center, node) in enumerate(zip(patch_centers, nodes)):
             s = img_util.get_slices(center, self.patch_shape)
-            patches[i, 0, ...] = img_util.normalize(np.minimum(img[s], 300))
+            patches[i, 0, ...] = img_util.normalize(np.minimum(img[s], 400))
             patches[i, 1, ...] = label_mask[s]
 
             subgraph = self.graph.get_rooted_subgraph(node, self.subgraph_radius)