From 2300e62b91d2e48ac85fccd79f1de7cb14fd2800 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Tue, 3 Feb 2026 02:43:21 +0000
Subject: [PATCH 1/2] refactor: feedfoward and normalization

---
 src/neuron_proofreader/utils/img_util.py | 11 +++++++----
 src/neuron_proofreader/utils/ml_util.py  |  7 ++++---
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/neuron_proofreader/utils/img_util.py b/src/neuron_proofreader/utils/img_util.py
index c70d114..6f344c8 100644
--- a/src/neuron_proofreader/utils/img_util.py
+++ b/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)
 
 
diff --git a/src/neuron_proofreader/utils/ml_util.py b/src/neuron_proofreader/utils/ml_util.py
index 14cbe77..63eb008 100644
--- a/src/neuron_proofreader/utils/ml_util.py
+++ b/src/neuron_proofreader/utils/ml_util.py
@@ -37,12 +37,13 @@ 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 = []
@@ -50,9 +51,9 @@ def build_network(self, input_dim, output_dim, 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

From fc8eebfbbf866769f8e536b18392f1502f03f00e Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Tue, 3 Feb 2026 02:44:50 +0000
Subject: [PATCH 2/2] refactor: doc

---
 src/neuron_proofreader/utils/ml_util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/neuron_proofreader/utils/ml_util.py b/src/neuron_proofreader/utils/ml_util.py
index 63eb008..86507b5 100644
--- a/src/neuron_proofreader/utils/ml_util.py
+++ b/src/neuron_proofreader/utils/ml_util.py
@@ -168,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
     ----------