Pr 787

src/neuron_proofreader/configs.py

@@ -23,6 +23,12 @@
 
 class Config(ABC):
 
+    def keys(self):
+        return asdict(self).keys()
+
+    def __getitem__(self, key):
+        return asdict(self)[key]
+
     def to_dict(self):
         """
         Converts configuration attributes to a dictionary.
@@ -90,15 +96,19 @@ class ImageConfig(Config):
     ----------
     brightness_clip : int
         Intensity value that voxel brightness is clipped to.
+    fov_merge : Tuple[int], optional
+        Patch shape loaded from the image for merge detection. Resized to
+        patch_shape before the model. None loads at patch_shape directly.
     percentiles : Tuple[float], optional
         Percentiles used to normalize patches.
     patch_shape : Tuple[int]
-        Shape of patch to be read from image.
+        Shape of patch expected by the model.
     transform : bool
         Indication of whether to use image augmentation.
     """
 
     brightness_clip: int = 400
+    fov_merge: Tuple[int, int, int] = None
     img_path: str = None
     name: str = "image_config"
     percentiles: Tuple[float, float] = (1, 99.5)
@@ -127,6 +137,10 @@ class MergeInferenceConfig(Config):
         memory; must be set before running learned merge detection.
     min_search_size : float
         Minimum fragment cable length (in microns) to include in the search.
+    modality : str
+        Input modality expected by the merge model. Options are "image" for
+        image-only models and "image_graph" for models that also consume a
+        skeleton TreeSample per site (e.g. JointProofreader).
     model_config_path : str, optional
         Path to the merge model config JSON file.
     model_path : str, optional
@@ -145,6 +159,7 @@ class MergeInferenceConfig(Config):
 
     batch_size: int = None
     min_search_size: float = 0
+    modality: str = "image"
     model_config_path: str = None
     model_path: str = None
     name: str = "merge_inference_config"
@@ -167,6 +182,8 @@ class SplitInferenceConfig(Config):
         Increment that acceptance threshold is lowered by each round.
     min_threshold : float
         Minimum confidence threshold for accepting a proposal.
+    model_config_path : str, optional
+        Path to the split model config JSON file.
     patch_shape : Tuple[int], optional
         Patch shape for image sampling, overrides ImageConfig.patch_shape if
         set.
@@ -176,16 +193,22 @@ class SplitInferenceConfig(Config):
         split model was trained with.
     removal_threshold : float
         Proposals with model predictions below this value are removed.
+    use_tree_features : bool
+        If True, each proposal carries an Arborist TreeSample alongside its
+        image patch. Required by models whose split head encodes skeleton
+        geometry (e.g. JointProofreader).
     """
 
     batch_size: int = None
     dt: float = 0.05
     min_threshold: float = 0.8
+    model_config_path: str = None
     model_path: str = None
     name: str = "split_inference_config"
     patch_shape: Tuple[int, int, int] = None
     percentiles: Tuple[float, float] = (1, 99.9)
     removal_threshold: float = 0.3
+    use_tree_features: bool = False
 
 
 @dataclass

src/neuron_proofreader/machine_learning/image_dataloader.py

@@ -276,21 +276,23 @@ class DetectionPatchLoader(PatchLoader):
 
     # --- Implementation of Abstract Inferface ---
     def __call__(self, node):
-        # Get patches
         center, shape = self.compute_patch_specs(node)
         img = self.read_image(center, shape)
         mask = self.create_mask(center, shape, node)
-        patches = self.stack(img, mask)
 
-        # Check whether to apply image augmentation
+        if self.fov_merge is not None:
+            img = img_util.resize(img, self.patch_shape)
+            mask = img_util.resize_nearest(mask, self.patch_shape)
+
+        patches = self.stack(img, mask)
         if self.transform:
             patches = self.transform(patches)
         return node, patches
 
     def compute_patch_specs(self, node):
         voxel = self.graph.node_voxel(node)
         voxel = self.adjust_voxel(voxel)
-        return voxel, self.patch_shape
+        return voxel, self.fov_merge or self.patch_shape
 
 
 class DetectionBatchLoader(PatchLoader):

src/neuron_proofreader/machine_learning/train.py

@@ -73,6 +73,7 @@ def __init__(
         max_epochs=200,
         pos_weight=None,
         n_mips=0,
+        n_mistake_mips=64,
         use_amp=True,
     ):
         """
@@ -102,6 +103,10 @@ def __init__(
             Number of randomly sampled validation predictions to save as MIPs
             each epoch, named by outcome (e.g. "false_positive{idx}.png").
             Default is 0 (disabled).
+        n_mistake_mips : int, optional
+            Number of mistake MIPs (false positives and false negatives) to
+            save per epoch via reservoir sampling. Saved under
+            mistake_mips/epoch_NNN/ in the log directory. Default is 64.
         use_amp : bool, optional
             If True, enables automatic mixed precision (float16) training.
             Default is True.
@@ -119,8 +124,10 @@ def __init__(
         self.log_dir = log_dir
         self.max_epochs = max_epochs
         self.mips_dir = os.path.join(log_dir, "mips")
+        self.mistake_mips_dir = os.path.join(log_dir, "mistake_mips")
         self.model_name = model_name
         self.n_mips = n_mips
+        self.n_mistake_mips = n_mistake_mips
 
         pw = (
             torch.tensor([pos_weight], device=device)
@@ -202,6 +209,9 @@ def _run_epoch(self, dataloader, epoch, train):
                 util.mkdir(self.mips_dir, True)
                 n = len(dataloader.dataset)
                 self._mip_idxs = set(random.sample(range(n), min(self.n_mips, n)))
+            if self.n_mistake_mips > 0:
+                self._mistake_buf = []
+                self._mistake_seen = 0
 
         metrics = ml_util.BinaryMetricAccumulator()
         for x, y in dataloader:
@@ -222,10 +232,13 @@ def _run_epoch(self, dataloader, epoch, train):
 
             if not train:
                 self._save_mips(x, y, y_pred, idx_offset)
+                self._collect_mistakes(x, y, y_pred)
                 idx_offset += len(y)
 
         stats = metrics.compute(min_recall=self.enforced_recall)
         self.update_tensorboard(stats, epoch, prefix)
+        if not train and self.n_mistake_mips > 0:
+            self._flush_mistake_mips(epoch)
         return stats
 
     def forward_pass(self, x, y):
@@ -348,6 +361,37 @@ def _save_mips(self, x, y, y_pred, idx_offset):
                 x[i, 0], 2 * x[i, 1], output_path
             )
 
+    def _collect_mistakes(self, x, y, y_pred):
+        """
+        Adds false positives and false negatives from this batch to a
+        reservoir-sampled buffer capped at n_mistake_mips images.
+        """
+        if self.n_mistake_mips == 0:
+            return
+        imgs = ml_util.to_cpu(x["img"] if isinstance(x, dict) else x, True)
+        for i in range(len(y)):
+            outcome = classify_prediction(y[i], y_pred[i])
+            if outcome not in ("false_positive", "false_negative"):
+                continue
+            self._mistake_seen += 1
+            item = (imgs[i], outcome)
+            if len(self._mistake_buf) < self.n_mistake_mips:
+                self._mistake_buf.append(item)
+            else:
+                j = random.randint(0, self._mistake_seen - 1)
+                if j < self.n_mistake_mips:
+                    self._mistake_buf[j] = item
+
+    def _flush_mistake_mips(self, epoch):
+        """
+        Saves the buffered mistake MIPs for this epoch to disk.
+        """
+        epoch_dir = os.path.join(self.mistake_mips_dir, f"epoch_{epoch:03d}")
+        util.mkdir(epoch_dir)
+        for i, (img, outcome) in enumerate(self._mistake_buf):
+            path = os.path.join(epoch_dir, f"{outcome}_{i:02d}.png")
+            img_util.plot_image_and_segmentation_mips(img[0], 2 * img[1], path)
+
     def save_config(self):
         """
         Saves trainer configuration to a JSON file in the log directory.
@@ -414,6 +458,7 @@ def __init__(
         max_epochs=200,
         pos_weight=None,
         n_mips=0,
+        n_mistake_mips=64,
     ):
         """
         Instantiates a DistributedTrainer object.
@@ -433,6 +478,8 @@ def __init__(
         pos_weight : float or None, optional
             Weight applied to the positive class in BCEWithLogitsLoss. Default
             is None (no reweighting).
+        n_mistake_mips : int, optional
+            Number of mistake MIPs to save per epoch. Default is 64.
         """
         # Call parent class
         super().__init__(
@@ -444,6 +491,7 @@ def __init__(
             max_epochs=max_epochs,
             pos_weight=pos_weight,
             n_mips=n_mips,
+            n_mistake_mips=n_mistake_mips,
         )
 
         # Check that multiple GPUs are available

src/neuron_proofreader/merge_proofreading/search_datasets.py

@@ -290,7 +290,7 @@ def get_patch_and_arborist(self, nodes, img, offset):
             s = img_util.get_slices(center, self.patch_shape)
             patch = img[(slice(0, 2), *s)]
             subgraph = self.graph.rooted_subgraph(node, self.subgraph_radius)
-            tree_sample = subgraph_to_tree_sample(subgraph, node)
+            tree_sample = subgraph_to_tree_sample(subgraph, 0)
             yield node, patch, tree_sample
 
     # --- Helpers ---
@@ -353,5 +353,5 @@ def get_patch(self, node, img):
     def get_patch_and_arborist(self, node, img):
         patch = torch.from_numpy(img).float()
         subgraph = self.graph.rooted_subgraph(node, self.subgraph_radius)
-        tree_sample = subgraph_to_tree_sample(subgraph, node)
+        tree_sample = subgraph_to_tree_sample(subgraph, 0)
         yield node, patch, tree_sample