Cleanup unused code (#791)

* fix: forward use_anisotropy through ProposalGraph and the pipeline

ProposalGraph.__init__ accepted use_anisotropy via **kwargs but never
passed it to FragmentsGraph, and ProofreadPipeline.load_graph did not
forward GraphConfig.use_anisotropy either. The flag was therefore
silently ignored and SWC coordinates were always scaled by the
anisotropy on load, which is wrong for SWCs already in world
coordinates such as pipeline outputs.

Callers reading voxel-space SWCs must now pass use_anisotropy=True
(the default) explicitly in their GraphConfig to keep scaling.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix: sparse search progress bar counted unsearched fragments

SparseSearchDataset.estimate_iterations returned every branching node in
the graph, but the dataset only visits fragments longer than
min_search_size. The tqdm total was therefore inflated (45k vs 28k
visited on one whole-brain graph) and the ETA pessimistic.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* refactor: deduplicate Dense/Sparse search datasets

Move patch-loader construction, search_mode, and default step_size into
class attributes handled by SearchDataset.__init__, and define the
patch + rooted-subgraph input once as get_patch_and_subgraph (renamed from
get_patch_and_arborist). The base class now defaults prefetch to 64.

Co-Authored-By: Claude <noreply@anthropic.com>

* refactor: drop superseded models, split io_util out of util, add DatasetConfig

- Remove the old vision_models.py and gnn_models.py and promote
  new_vision_models.py / new_gnn_models.py to those names.
- Move file, zip, GCS and S3 I/O from utils/util.py into utils/io_util.py
  and update callers.
- Add DatasetConfig (brain_id, segmentation_id, img_path, fragments_path)
  and give ProposalsConfig a name so every config can be saved.

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Anna Grim <anna.grim@alleninstitute.org>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: 3 people

src/neuron_proofreader/configs.py

@@ -18,7 +18,7 @@
 from neuron_proofreader.machine_learning.image_augmentation import (
     ImageTransforms,
 )
-from neuron_proofreader.utils import util
+from neuron_proofreader.utils import io_util
 
 
 class Config(ABC):
@@ -52,7 +52,31 @@ def save(self, output_dir):
             Path to directory to save JSON file.
         """
         path = os.path.join(output_dir, f"{self.name}.json")
-        util.write_json(path, self.to_dict())
+        io_util.write_json(path, self.to_dict())
+
+
+@dataclass
+class DatasetConfig(Config):
+    """
+    Configuration class for the dataset that is proofread.
+
+    Attributes
+    ----------
+    brain_id : str
+        Unique identifier of the whole-brain dataset.
+    fragments_path : str
+        Path to the directory (local or cloud) containing fragment SWC files.
+    img_path : str
+        Path to the whole-brain image.
+    segmentation_id : str
+        Unique identifier of the segmentation that fragments come from.
+    """
+
+    brain_id: str = None
+    fragments_path: str = None
+    img_path: str = None
+    name: str = "dataset_config"
+    segmentation_id: str = None
 
 
 @dataclass
@@ -237,6 +261,7 @@ class ProposalsConfig(Config):
     allow_nonleaf_proposals: bool = False
     max_proposals_per_leaf: int = 3
     min_size_with_proposals: float = 0
+    name: str = "proposals_config"
     trim_endpoints_bool: bool = True
     search_radius: float = 25
     search_scaling_factor: float = 1.5

src/neuron_proofreader/fragments_graph.py

@@ -15,7 +15,7 @@
 
 from arborist.skeleton_graph import SkeletonGraph
 from arborist.utils.graph_loading import GraphLoader
-from neuron_proofreader.utils import geometry_util, img_util, util
+from neuron_proofreader.utils import geometry_util, img_util, io_util
 
 
 class FragmentsGraph(SkeletonGraph):
@@ -226,8 +226,8 @@ def clip_to_bbox(self, metadata_path):
         """
         Clips skeletons to the bounding box defined in a metadata JSON file.
         """
-        if util.check_gcs_file_exists(metadata_path):
-            metadata = util.read_json(metadata_path)
+        if io_util.check_gcs_file_exists(metadata_path):
+            metadata = io_util.read_json(metadata_path)
             origin = metadata["chunk_origin"][::-1]
             shape = metadata["chunk_shape"][::-1]
             nodes = list()

src/neuron_proofreader/machine_learning/image_dataloader.py

@@ -17,7 +17,7 @@
 from neuron_proofreader.machine_learning.image_augmentation import (
     ImageTransforms,
 )
-from neuron_proofreader.utils import geometry_util, img_util, util
+from neuron_proofreader.utils import geometry_util, img_util, io_util
 
 
 class TensorStoreImage:
@@ -56,7 +56,7 @@ def __init__(self, img_path):
             Path to image.
         """
         # Load image
-        bucket_name, inner_path = util.parse_cloud_path(img_path)
+        bucket_name, inner_path = io_util.parse_cloud_path(img_path)
         self.img_path = img_path
         self.img = ts.open(
             {

src/neuron_proofreader/machine_learning/train.py

@@ -26,7 +26,7 @@
 import torch.nn as nn
 import torch.optim as optim
 
-from neuron_proofreader.utils import img_util, ml_util, util
+from neuron_proofreader.utils import img_util, io_util, ml_util
 
 
 class Trainer:
@@ -115,7 +115,7 @@ def __init__(
         if exp_name is None:
             exp_name = "session-" + datetime.today().strftime("%Y%m%d_%H%M")
         log_dir = os.path.join(output_dir, exp_name)
-        util.mkdir(log_dir)
+        io_util.mkdir(log_dir)
 
         # Instance attributes
         self.best_f1 = 0
@@ -206,7 +206,7 @@ def _run_epoch(self, dataloader, epoch, train):
         else:
             self.model.eval()
             if self.n_mips > 0:
-                util.mkdir(self.mips_dir, True)
+                io_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:
@@ -387,7 +387,7 @@ 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)
+        io_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)

src/neuron_proofreader/merge_proofreading/merge_datamodules.py

@@ -41,7 +41,7 @@
 from neuron_proofreader.fragments_graph import FragmentsGraph
 from arborist.data.augmentation import GraphTransforms
 from arborist.utils.swc_loading import Reader
-from neuron_proofreader.utils import ml_util, util
+from neuron_proofreader.utils import io_util, ml_util, util
 from neuron_proofreader.utils.graph_util import subgraph_to_tree_sample
 
 
@@ -407,7 +407,7 @@ def save_config(self, path):
             }
             for bd in self.datasets
         ]
-        util.write_json(path, config)
+        io_util.write_json(path, config)
 
     # --- Dataset Interface ---
     def __len__(self):
@@ -704,8 +704,8 @@ def create_dataset_collection(
             img_config.set_val_mode()
 
     # Load image prefixes
-    bucket, root_prefix = util.parse_cloud_path(sites_root_path)
-    img_prefixes = util.read_json(img_prefixes_path)
+    bucket, root_prefix = io_util.parse_cloud_path(sites_root_path)
+    img_prefixes = io_util.read_json(img_prefixes_path)
 
     # Iterate over brains
     datasets = list()
@@ -750,6 +750,6 @@ def create_dataset_collection(
 
 def get_segmentation_id(sites_path, brain_id):
     brain_sites_path = os.path.join(sites_path, brain_id)
-    subprefixes = util.list_gcs_subprefixes(brain_sites_path)
+    subprefixes = io_util.list_gcs_subprefixes(brain_sites_path)
     assert len(subprefixes) == 1
     return subprefixes[0].split("/")[-2]

src/neuron_proofreader/merge_proofreading/merge_inference.py

@@ -39,7 +39,7 @@
     SparseSearchDataset,
     multimodal_collate,
 )
-from neuron_proofreader.utils import img_util, util
+from neuron_proofreader.utils import img_util, io_util, util
 
 
 class MergeProofreader(ABC):
@@ -372,7 +372,7 @@ def save_parameters(self):
             "search_mode": self.dataset.search_mode,
             "subgraph_radius": self.dataset.subgraph_radius,
         }
-        util.write_json(json_path, parameters)
+        io_util.write_json(json_path, parameters)
 
     def save_predictions(self):
         nodes = np.array(self.visited_sites, dtype=int)

src/neuron_proofreader/merge_proofreading/search_datasets.py

@@ -32,7 +32,7 @@
 def multimodal_collate(batch):
     """
     Collate a batch of (node, img_patch, tree_sample) tuples produced by
-    get_patch_and_arborist.
+    get_patch_and_subgraph.
 
     Returns
     -------
@@ -48,13 +48,18 @@ def multimodal_collate(batch):
 # --- Datasets ---
 class SearchDataset(IterableDataset, ABC):
 
+    search_mode = None
+    patch_loader_cls = None
+    default_step_size = None
+
     def __init__(
         self,
         graph,
         img_config,
         modality="image",
         min_search_size=0,
-        prefetch=32,
+        prefetch=64,
+        step_size=None,
         subgraph_radius=100,
     ):
         # Call parent class
@@ -64,13 +69,15 @@ def __init__(
         self.graph = graph
         self.modality = modality
         self.min_size = min_search_size
+        self.patch_loader = self.patch_loader_cls(graph, img_config)
         self.patch_shape = img_config.patch_shape
         self.prefetch = prefetch
+        self.step_size = step_size or self.default_step_size
         self.subgraph_radius = subgraph_radius
 
         # Input getter
         if modality != "image":
-            self.get_input = self.get_patch_and_arborist
+            self.get_input = self.get_patch_and_subgraph
         else:
             self.get_input = self.get_patch
         self.collate_fn = multimodal_collate if modality != "image" else None
@@ -136,6 +143,11 @@ def generate_component_sites(self, root):
         """
         pass
 
+    def get_patch_and_subgraph(self, *args):
+        for node, patch in self.get_patch(*args):
+            subgraph = self.graph.rooted_subgraph(node, self.subgraph_radius)
+            yield node, patch, subgraph_to_tree_sample(subgraph, 0)
+
     # --- Helpers ---
     def __getattr__(self, name):
         return getattr(self.graph, name)
@@ -214,31 +226,9 @@ def _is_near_leaf_slow(self, node, threshold=32):
 class DenseSearchDataset(SearchDataset):
 
     max_batch_span = 512
-
-    def __init__(
-        self,
-        graph,
-        img_config,
-        modality="image",
-        min_search_size=0,
-        prefetch=64,
-        step_size=40,
-        subgraph_radius=100,
-    ):
-        # Call parent class
-        super().__init__(
-            graph,
-            img_config,
-            modality=modality,
-            min_search_size=min_search_size,
-            prefetch=prefetch,
-            subgraph_radius=subgraph_radius,
-        )
-
-        # Instance attributes
-        self.patch_loader = DetectionBatchLoader(self.graph, img_config)
-        self.search_mode = "dense"
-        self.step_size = step_size
+    search_mode = "dense"
+    patch_loader_cls = DetectionBatchLoader
+    default_step_size = 40
 
     def generate_component_sites(self, root):
         """
@@ -289,16 +279,6 @@ def get_patch(self, nodes, img, offset):
             s = img_util.get_slices(center, self.patch_shape)
             yield node, img[(slice(0, 2), *s)]
 
-    def get_patch_and_arborist(self, nodes, img, offset):
-        img = torch.from_numpy(img).float()
-        voxels = np.array([self.node_voxel(i) for i in nodes], dtype=int)
-        for node, center in zip(nodes, voxels - 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, 0)
-            yield node, patch, tree_sample
-
     # --- Helpers ---
     def estimate_iterations(self):
         """
@@ -319,29 +299,9 @@ def estimate_iterations(self):
 
 class SparseSearchDataset(SearchDataset):
 
-    def __init__(
-        self,
-        graph,
-        img_config,
-        modality="image",
-        min_search_size=0,
-        prefetch=64,
-        step_size=10,
-        subgraph_radius=100,
-    ):
-        # Call parent class
-        super().__init__(
-            graph,
-            img_config,
-            modality=modality,
-            min_search_size=min_search_size,
-            prefetch=prefetch,
-            subgraph_radius=subgraph_radius,
-        )
-
-        # Instance attributes
-        self.patch_loader = DetectionPatchLoader(self.graph, img_config)
-        self.search_mode = "sparse"
+    search_mode = "sparse"
+    patch_loader_cls = DetectionPatchLoader
+    default_step_size = 10
 
     def estimate_iterations(self):
         # Only fragments above min_search_size are searched
@@ -358,9 +318,3 @@ def generate_component_sites(self, root):
 
     def get_patch(self, node, img):
         yield node, torch.from_numpy(img).float()
-
-    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, 0)
-        yield node, patch, tree_sample