Refactor ffn

src/neuron_proofreader/utils/img_util.py

@@ -659,22 +659,25 @@ def is_precomputed(img_path):
         return False
 
 
-def normalize(img):
+def normalize(img, percentiles=(1, 99.5)):
     """
-    Normalizes an image so that the minimum and maximum intensity values are 0
-    and 1.
+    Normalizes an image using a percentile-based scheme and clips values to
+    [0, 1].
 
     Parameters
     ----------
     img : numpy.ndarray
         Image to be normalized.
+    percentiles : Tuple[float], optional
+        Upper and lower percentiles used to normalize the given image. Default
+        is (1, 99.5).
 
     Returns
     -------
     img : numpy.ndarray
         Normalized image.
     """
-    mn, mx = np.percentile(img, [1, 99.9])
+    mn, mx = np.percentile(img, percentiles)
     return np.clip((img - mn) / (mx - mn + 1e-5), 0, 1)
 
 

src/neuron_proofreader/utils/ml_util.py

@@ -37,22 +37,23 @@ def __init__(self, input_dim, output_dim, n_layers):
         super().__init__()
 
         # Instance attributes
+        assert n_layers > 1
         self.net = self.build_network(input_dim, output_dim, n_layers)
 
     def build_network(self, input_dim, output_dim, n_layers):
         # Set input/output dimensions
         input_dim_i = input_dim
-        output_dim_i = input_dim // 2
+        output_dim_i = max(input_dim // 2, 4)
 
         # Build architecture
         layers = []
         for i in range(n_layers):
             mlp = init_mlp(input_dim_i, input_dim_i * 2, output_dim_i)
             layers.append(mlp)
 
-            input_dim_i = input_dim_i // 2
+            input_dim_i = output_dim_i
             output_dim_i = (
-                output_dim_i // 2 if i < n_layers - 2 else output_dim
+                max(output_dim_i // 2, 4) if i < n_layers - 2 else output_dim
             )
 
         # Initialize weights
@@ -167,7 +168,7 @@ def load_model(model, model_path, device="cuda"):
 
 def tensor_to_list(tensor):
     """
-    Converts a tensor to a list.
+    Converts the given tensor to a list.
 
     Parameters
     ----------