From 65dd8feadda514bf5f240a2e0bb98561ef64edc4 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Fri, 10 Jul 2026 18:12:25 +0000
Subject: [PATCH 1/3] refactor: split datasets

---
 src/neuron_proofreader/configs.py             |  11 +-
 src/neuron_proofreader/skeleton_graph.py      |   2 +-
 .../split_proofreading/split_datasets.py      | 127 +++++-------------
 3 files changed, 44 insertions(+), 96 deletions(-)

diff --git a/src/neuron_proofreader/configs.py b/src/neuron_proofreader/configs.py
index ec28a2ec..612e1abb 100644
--- a/src/neuron_proofreader/configs.py
+++ b/src/neuron_proofreader/configs.py
@@ -46,6 +46,15 @@ def save(self, output_dir):
         path = os.path.join(output_dir, f"{self.name}.json")
         util.write_json(path, self.to_dict())
 
+    def __repr__(self):
+        fields = self.to_dict()
+        width = max(len(k) for k in fields)
+        lines = "\n".join(
+            f"  {k:<{width}} = {v!r}"
+            for k, v in fields.items()
+        )
+        return f"{self.__class__.__name__}(\n{lines}\n)"
+
 
 @dataclass
 class GraphConfig(Config):
@@ -124,7 +133,7 @@ class ProposalsConfig(Config):
     allow_nonleaf_proposals : bool
         Indication of whether to generate proposals between leaf and nodes
         with degree 2.
-    proposals_per_leaf : int
+    max_proposals_per_leaf : int
         Maximum number of proposals generated at leaf nodes.
     trim_endpoints_bool : bool
         Indication of whether trim endpoints of isolated leaf-to-leaf
diff --git a/src/neuron_proofreader/skeleton_graph.py b/src/neuron_proofreader/skeleton_graph.py
index 1bd13a90..cdf5a222 100644
--- a/src/neuron_proofreader/skeleton_graph.py
+++ b/src/neuron_proofreader/skeleton_graph.py
@@ -627,7 +627,7 @@ def clip_to_bbox(self, metadata_path):
             clip skeletons to.
         """
         bucket_name, path = util.parse_cloud_path(metadata_path)
-        if util.check_gcs_file_exists(bucket_name, path):
+        if util.check_gcs_exists(bucket_name, path):
             # Extract bounding box
             metadata = util.read_json_from_gcs(bucket_name, path)
             origin = metadata["chunk_origin"][::-1]
diff --git a/src/neuron_proofreader/split_proofreading/split_datasets.py b/src/neuron_proofreader/split_proofreading/split_datasets.py
index 635ef531..f19dbb36 100644
--- a/src/neuron_proofreader/split_proofreading/split_datasets.py
+++ b/src/neuron_proofreader/split_proofreading/split_datasets.py
@@ -31,46 +31,31 @@
 class FragmentsDataset(IterableDataset):
     """
     A dataset object that contains a graph built from fragments corresponding
-    to a single brain. Note that this dataset supports fragments extracted
-    from either a block or whole-brain.
+    to a single brain.
     """
 
-    def __init__(
-        self,
-        fragments_graph,
-        img_config,
-        batch_size=32,
-        gt_path=None,
-        prefetch=4,
-    ):
+    def __init__(self, fragments_graph, img_config, proposals_per_batch=64):
         """
         Instantiates a FragmentsDataset object.
 
         Parameters
         ----------
-        fragments_path : str
-            Path to predicted SWC files to be loaded.
-        img_path : str
-            Path to the raw image associated with the fragments.
-        graph_config : GraphConfig
+        fragments_graph : str
             ...
-        gt_path : str, optional
-            Path to ground-truth SWC files to be loaded. Default is None.
+        img_config : ImageConfig
+            Configuration object containing image parameters.
+        proposals_per_batch : int, optional
+            Maximum number of proposals in subgraphs yielded per batch.
+            Default is 64.
         """
-        # Instance attributes
-        self.batch_size = batch_size
         self.graph = fragments_graph
-        self.gt_path = gt_path
-        self.prefetch = prefetch
-        self.transform = img_config.transform
-
-        # Feature extractor
         self.feature_extractor = FeaturePipeline(
             self.graph,
             img_config.img_path,
             brightness_clip=img_config.brightness_clip,
             patch_shape=img_config.patch_shape,
         )
+        self.proposals_per_batch = proposals_per_batch
 
     # --- Get Data ---
     def __iter__(self):
@@ -101,7 +86,9 @@ def get_sampler(self):
         sampler : SubgraphSampler
             Subgraph sampler that is used to iterate over dataset.
         """
-        sampler = SubgraphSampler(self.graph, max_proposals=self.batch_size)
+        sampler = SubgraphSampler(
+            self.graph, max_proposals=self.proposals_per_batch
+        )
         return iter(sampler)
 
 
@@ -111,64 +98,19 @@ class FragmentsDatasetCollection(IterableDataset):
     to different whole-brain images.
     """
 
-    def __init__(self, shuffle=True):
+    def __init__(self, datasets, shuffle=True):
         """
         Instantiates a FragmentsDatasetCollection object.
 
         Parameters
         ----------
         shuffle : bool, optional
-            Indication of whether to shuffle examples. Default is True.
+            True if examples should be shuffled each epoch. Default is True.
         """
         # Instance attributes
-        self.datasets = dict()
+        self.datasets = datasets
         self.shuffle = shuffle
 
-    def add_dataset(
-        self,
-        key,
-        fragments_path,
-        img_path,
-        config,
-        gt_path=None,
-        metadata_path=None,
-        prefetch=4,
-        soma_centroids=list(),
-    ):
-        """
-        Adds a dataset to the collection of datasets.
-
-        Parameters
-        ----------
-        key : hashable
-            Unique identifier of the dataset to be added.
-        fragments_path : str
-            Path to predicted SWC files to be loaded.
-        img_path : str
-            Path to the raw image associated with the fragments.
-        config : Config
-            Configuration object containing parameters and settings.
-        gt_path : str, optional
-            Path to ground-truth SWC files to be loaded. Default is None.
-        metadata_path : str, optional
-            Patch to JSON file containing metadata on block that fragments
-            were extracted from. Default is None.
-        prefetch : int, optional
-            Number of batches to prefetch. Default is 4.
-        soma_centroids : List[Tuple[int]], optional
-            Phyiscal coordinates of soma centroids. Default is an empty list.
-        """
-        assert key not in self.datasets, "Key has been used!"
-        self.datasets[key] = FragmentsDataset(
-            fragments_path,
-            img_path,
-            config,
-            gt_path=gt_path,
-            metadata_path=metadata_path,
-            prefetch=prefetch,
-            soma_centroids=soma_centroids,
-        )
-
     def __iter__(self):
         """
         Iterates over the datasets and yields model-ready inputs and targets.
@@ -180,7 +122,7 @@ def __iter__(self):
         targets : torch.Tensor
             Ground truth labels.
         """
-        # Initializations
+        # Create prefetching data structures
         samplers = self.init_samplers()
         queue = Queue(maxsize=self.prefetch * len(samplers))
         active_keys = set(samplers.keys())
@@ -214,18 +156,20 @@ def _worker(self, key, sampler, queue):
         finally:
             queue.put((key, StopIteration, None))
 
-    def generate_proposals(self, search_radius):
+    def generate_proposals(self, proposal_config):
         """
         Generates proposals for each dataset.
 
         Parameters
         ----------
-        search_radius : float
-            Search radius used to generate proposals.
+        proposal_config : ProposalConfig
+            Configuration object containing parameters for generating
+            proposals.
         """
         for key in tqdm(self.datasets, desc="Generate Proposals"):
             self.datasets[key].graph.generate_proposals(
-                search_radius, allow_nonleaf_proposals=True
+                proposal_config.search_radius,
+                **proposal_config,
             )
 
     # --- Helpers ---
@@ -322,39 +266,34 @@ def save_examples_summary(self, path):
 
 
 # -- Helpers --
-def generate_dataset_example_ids(bucket_name, dataset_prefix):
+def generate_example_ids(ds_prefix):
     """
     Generates dataset example identifiers.
 
     Parameters
     ----------
-    bucket_name : str
-        Name of the Google Cloud Storage bucket.
-    dataset_prefix : str
-        Root prefix under which dataset contents are organized.
+    ds_prefix : str
+        Root prefix under which datasets are organized.
 
     Yields
     -------
     Tuple[str]
-        Dataset example ID formatted as (brain_id, segmentation_id, block_id).
+        Dataset ID formatted as (brain_id, segmentation_id, block_id).
     """
-    brain_prefixes = util.list_gcs_subdirectories(bucket_name, dataset_prefix)
-    for brain_prefix in brain_prefixes:
+    bucket_name, _ = util.parse_cloud_path(ds_prefix)
+    for brain_prefix in util.list_gcs_subprefixes(ds_prefix):
         # Extract brain id
         brain_id = brain_prefix.split("/")[-2]
 
         # Iterate over segmentations
-        pred_prefix = os.path.join(brain_prefix, "pred_swcs/")
-        prefixes = util.list_gcs_subdirectories(bucket_name, pred_prefix)
-        for brain_segmentation_prefix in prefixes:
+        pred_prefix = os.path.join(bucket_name, brain_prefix, "pred_swcs")
+        for brain_seg_prefix in util.list_gcs_subprefixes(pred_prefix):
             # Extract segmentation id
-            segmentation_id = brain_segmentation_prefix.split("/")[-2]
+            segmentation_id = brain_seg_prefix.split("/")[-2]
 
             # Iterate over blocks
-            block_prefixes = util.list_gcs_subdirectories(
-                bucket_name, brain_segmentation_prefix
-            )
-            for block_prefix in block_prefixes:
+            ex_prefix = os.path.join(bucket_name, brain_seg_prefix)
+            for block_prefix in util.list_gcs_subprefixes(ex_prefix):
                 # Extract block id
                 block_id = block_prefix.split("/")[-2]
                 yield brain_id, segmentation_id, block_id

From 1e1ab28f02257340d6938f03245eab68ec9389e3 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Fri, 10 Jul 2026 22:29:34 +0000
Subject: [PATCH 2/3] optimized split train

---
 src/neuron_proofreader/configs.py             |  12 +-
 .../machine_learning/train.py                 |  24 +-
 .../models/geometric_gnn_models.py            | 585 ------------------
 src/neuron_proofreader/proposal_graph.py      |  30 +-
 src/neuron_proofreader/skeleton_graph.py      |   5 +-
 .../split_proofreading/proposal_generation.py |  70 +--
 .../split_proofreading/split_datamodules.py   | 262 ++++++++
 .../split_proofreading/split_datasets.py      | 299 ---------
 .../split_feature_extraction.py               |  57 +-
 .../split_proofreading/split_inference.py     |   2 +-
 src/neuron_proofreader/utils/ml_util.py       |   2 +
 src/neuron_proofreader/utils/util.py          |  65 +-
 12 files changed, 374 insertions(+), 1039 deletions(-)
 delete mode 100644 src/neuron_proofreader/models/geometric_gnn_models.py
 create mode 100644 src/neuron_proofreader/split_proofreading/split_datamodules.py
 delete mode 100644 src/neuron_proofreader/split_proofreading/split_datasets.py

diff --git a/src/neuron_proofreader/configs.py b/src/neuron_proofreader/configs.py
index 612e1abb..0e6edfb8 100644
--- a/src/neuron_proofreader/configs.py
+++ b/src/neuron_proofreader/configs.py
@@ -133,16 +133,20 @@ class ProposalsConfig(Config):
     allow_nonleaf_proposals : bool
         Indication of whether to generate proposals between leaf and nodes
         with degree 2.
+    initial_search_radius : float
+        Initial search radius for generating proposals.
     max_proposals_per_leaf : int
         Maximum number of proposals generated at leaf nodes.
     trim_endpoints_bool : bool
-        Indication of whether trim endpoints of isolated leaf-to-leaf
-        proposals.
+        True if endpoints of isolated leaf-to-leaf proposals should be 
+        trimmed.
     """
 
     allow_nonleaf_proposals: bool = False
+    initial_search_radius: float = 25
+    max_attempts: int = 2
     max_proposals_per_leaf: int = 3
     min_size_with_proposals: float = 0
-    trim_endpoints_bool: bool = True
-    search_radius: float = 25
+    name: str = "proposals_config"
     search_scaling_factor: float = 1.5
+    trim_endpoints_bool: bool = True
diff --git a/src/neuron_proofreader/machine_learning/train.py b/src/neuron_proofreader/machine_learning/train.py
index 83e19e19..094ab90e 100644
--- a/src/neuron_proofreader/machine_learning/train.py
+++ b/src/neuron_proofreader/machine_learning/train.py
@@ -117,7 +117,7 @@ def __init__(
         self.writer = SummaryWriter(log_dir=log_dir)
 
     # --- Core Routines ---
-    def run(self, train_dataloader, val_dataloader):
+    def __call__(self, train_dataloader, val_dataloader):
         """
         Runs the full training and validation loop.
 
@@ -128,9 +128,7 @@ def run(self, train_dataloader, val_dataloader):
         val_dataloader : torch.utils.data.Dataset
             Dataloader used for validation.
         """
-        exp_name = os.path.basename(os.path.normpath(self.log_dir))
-        val_dataloader.dataset.save_val_summary(self.log_dir)
-        print("\nExperiment:", exp_name)
+        print("\nExperiment:", os.path.basename(self.log_dir))
         for epoch in range(self.max_epochs):
             # Train-Validate
             print("\nEpoch", epoch)
@@ -181,11 +179,8 @@ def train_step(self, dataloader, epoch):
             self.scaler.step(self.optimizer)
             self.scaler.update()
 
-            # Compute metrics
-            hat_y = hat_y > 0
-            metrics.update(hat_y, y, loss)
-
-            # Update progress bar
+            # Updates
+            metrics.update(hat_y > 0, y, loss)
             if self.verbose:
                 pbar.update(1)
 
@@ -229,15 +224,12 @@ def validate_step(self, dataloader, epoch):
             with torch.inference_mode():
                 y, hat_y, loss = self.forward_pass(x, y)
 
-            # Compute metrics
-            hat_y = hat_y > 0
-            metrics.update(hat_y, y, loss)
-
             # Save MIPs of mistakes
             self._save_mistake_mips(x, y, hat_y, idx_offset)
             idx_offset += len(y)
 
-            # Update progress bar
+            # Updates
+            metrics.update(hat_y > 0, y, loss)
             if self.verbose:
                 pbar.update(1)
 
@@ -264,8 +256,8 @@ def forward_pass(self, x, y):
         loss : torch.Tensor
             Computed loss value.
         """
-        x = x.to(self.device, non_blocking=True)
-        y = y.to(self.device, non_blocking=True)
+        x = x.to(self.device)
+        y = y.to(self.device)
         with torch.autocast(device_type="cuda", dtype=torch.float16):
             hat_y = self.model(x)
             loss = self.criterion(hat_y, y)
diff --git a/src/neuron_proofreader/models/geometric_gnn_models.py b/src/neuron_proofreader/models/geometric_gnn_models.py
deleted file mode 100644
index 35398e15..00000000
--- a/src/neuron_proofreader/models/geometric_gnn_models.py
+++ /dev/null
@@ -1,585 +0,0 @@
-"""
-Created on Sat July 15 12:00:00 2025
-
-@author: Anna Grim
-@email: anna.grim@alleninstitute.org
-
-Code for graph neural network models that perform machine learning tasks
-within NeuronProofreader pipelines.
-
-"""
-
-from torch import nn
-from torch_geometric.nn import GATv2Conv
-
-import torch
-import torch.nn.functional as F
-
-from neuron_proofreader.models.vision_models import CNN3D
-from neuron_proofreader.utils.ml_util import FeedForwardNet
-
-
-# --- Multimodal GNN Architectures ---
-class VisionSkeleton(nn.Module):
-
-    def __init__(self, ggnn_name, patch_shape, output_dim=64):
-        # Call parent class
-        super().__init__()
-        assert ggnn_name in ["egnn"]
-
-        # Architecture
-        self.skeleton_model = SkeletonGNN(ggnn_name, output_dim=32)
-        self.vision_model = CNN3D(
-            patch_shape,
-            n_conv_layers=6,
-            n_feat_channels=20,
-            output_dim=output_dim,
-            use_double_conv=True,
-        )
-        self.output = FeedForwardNet(output_dim + 35, 1, 3)
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        # Modality-based embeddings
-        x_img = self.vision_model(x["img"])
-        x_skel = self.skeleton_model(*x["graph"])
-        x = torch.cat((x_img, x_skel), dim=1)
-
-        # Output layer
-        x = self.output(x)
-        return x
-
-
-class SkeletonGNN(nn.Module):
-
-    def __init__(self, ggnn_name, output_dim=64):
-        # Call parent class
-        super().__init__()
-
-        # Instance attributes
-        self.gnn_h = GAT(output_dim, 2 * output_dim, output_dim)
-        self.gnn_x = GAT(3, 16, 3)
-
-        # Set geometric gnn
-        if ggnn_name == "egnn":
-            self.geometric_gnn = EGNN(
-                in_node_dim=1, hidden_dim=32, out_node_dim=output_dim
-            )
-
-    # --- Core Routines ---
-    def forward(self, h, x, edge_index, batch):
-        # Node-level embeddings
-        h, x = self.geometric_gnn(h, x, edge_index)
-
-        # Graph-level embedddings
-        h_skels = list()
-        edge_index = edge_index
-        num_graphs = int(batch.max().item()) + 1
-        for graph_id in range(num_graphs):
-            # Extract subgraph
-            node_mask = batch == graph_id
-            h_g, x_g, edge_index_g = self.extract_subgraph(
-                h, x, edge_index, node_mask
-            )
-
-            # Pool node embeddings
-            h_g, x_g, edge_index_g = self.pool_nonbranching_paths(
-                h_g, x_g, edge_index_g
-            )
-
-            # Encode pooled graph
-            h_g = self.encode_pooled_graph(h_g, x_g, edge_index_g)
-            h_skels.append(h_g)
-        return torch.cat(h_skels, dim=0)
-
-    def pool_nonbranching_paths(self, h, x, edge_index):
-        # Extract adjacency matrix and degrees
-        num_nodes = h.size(0)
-        adj, deg = self.get_adj_and_deg(edge_index, num_nodes)
-
-        # Search graph
-        node_to_path = torch.full((num_nodes,), -1, device=h.device)
-        path_idx = 0
-        h_pooled = list()
-        x_pooled = list()
-        visited = set()
-        for start in range(num_nodes):
-            # Check whether to visit
-            if start in visited:
-                continue
-
-            # Case 1: Branch points are singleton paths
-            if deg[start] > 2:
-                node_to_path[start] = path_idx
-                h_pooled.append(h[start])
-                x_pooled.append(x[start])
-                path_idx += 1
-                visited.add(start)
-                continue
-
-            # Case 2: Non-branching path traversal
-            path = [start]
-            visited.add(start)
-            prev = None
-            cur = start
-            while True:
-                nbs = [n for n in adj[cur] if n != prev]
-                if len(nbs) != 1:
-                    break
-
-                nxt = nbs[0]
-                if nxt in visited or deg[nxt] > 2:
-                    break
-
-                path.append(nxt)
-                visited.add(nxt)
-                prev, cur = cur, nxt
-
-            h_pooled.append(h[path].mean(dim=0))
-            x_pooled.append(x[path].mean(dim=0))
-
-            for n in path:
-                node_to_path[n] = path_idx
-
-            path_idx += 1
-
-        # Finish
-        h_pooled = torch.stack(h_pooled, dim=0)
-        x_pooled = torch.stack(x_pooled, dim=0)
-        edge_index_pooled = self.get_edge_index_pooled(
-            edge_index, node_to_path
-        )
-        return h_pooled, x_pooled, edge_index_pooled
-
-    def get_adj_and_deg(self, edge_index, num_nodes):
-        # Compute node degrees
-        deg = torch.zeros(num_nodes, dtype=torch.long)
-        ones = torch.ones(edge_index.shape[1], dtype=torch.long)
-        deg.scatter_add_(0, edge_index[0], ones)
-        deg.scatter_add_(0, edge_index[1], ones)
-
-        # Build adjacency list
-        adj = [[] for _ in range(num_nodes)]
-        for u, v in edge_index.t().tolist():
-            adj[u].append(v)
-            adj[v].append(u)
-        return adj, deg
-
-    def encode_pooled_graph(self, h, x, edge_index):
-        # Message passing over pooled graph
-        h = self.gnn_h(h, edge_index)
-        x = self.gnn_x(x, edge_index)
-
-        # temp
-        if h.size(0) == 0 or x.size(0) == 0:
-            print(h.size(0), x.size(0))
-            raise RuntimeError("Empty tensor passed to graph pooling")
-
-        # Graph-level pooling
-        h = h.mean(dim=0, keepdim=True)
-        x = x.mean(dim=0, keepdim=True)
-        return torch.cat((h, x), dim=1)
-
-    # --- Helpers ---
-    def extract_subgraph(self, h, x, edge_index, node_mask):
-        # Build subgraph
-        node_ids = (node_mask).nonzero(as_tuple=True)[0]
-        h_g = h[node_ids]
-        x_g = x[node_ids]
-
-        # Remap nodes and edges
-        id_map = {int(n): i for i, n in enumerate(node_ids.tolist())}
-        edge_mask = node_mask[edge_index[0]] & node_mask[edge_index[1]]
-        edge_index_g = edge_index[:, edge_mask]
-        edge_index_g = torch.stack(
-            [
-                torch.tensor([id_map[int(u)] for u in edge_index_g[0]]),
-                torch.tensor([id_map[int(v)] for v in edge_index_g[1]]),
-            ],
-            dim=0,
-        )
-        return h_g, x_g, edge_index_g
-
-    @staticmethod
-    def get_edge_index_pooled(edge_index, node_to_path):
-        # Extract edges in pooled graph
-        src, dst = edge_index
-        src_p = node_to_path[src]
-        dst_p = node_to_path[dst]
-
-        # Remove intra-path edges
-        mask = src_p != dst_p
-        edge_index_pooled = torch.stack([src_p[mask], dst_p[mask]], dim=0)
-        return torch.unique(edge_index_pooled, dim=1)
-
-
-# --- Geometric GNN Architectures ---
-class EGNN(nn.Module):
-
-    def __init__(
-        self,
-        in_node_dim,
-        hidden_dim,
-        out_node_dim,
-        in_edge_dim=0,
-        device="cuda",
-        act_fn=nn.SiLU(),
-        n_layers=4,
-        residual=True,
-        attention=False,
-        normalize=False,
-        tanh=False,
-    ):
-        """
-        Instantiates an EGNN object.
-
-        Parameters
-        ----------
-        in_node_dim : int
-            Number of features for 'h' at the input.
-        hidden_dim : int
-            Number of hidden features.
-        out_node_dim : int
-            Number of features for 'h' at the output.
-        in_edge_dim : int, optional
-            Number of features for the edge features.
-        device : str
-            Device to load model and inputs. Default is "cuda".
-        act_fn : ...
-            Non-linearity
-        n_layers : int
-            Number of layer for the EGNN.
-        residual : bool
-            Indication of whether to use residual connections.
-        attention : bool
-            Indication of whether using attention mechanism.
-        normalize : bool
-            Normalizes the coordinates messages such that:
-                x^{l+1}_i = x^{l}_i + Σ(x_i - x_j)phi_x(m_ij)
-        tanh : ...
-            Sets a tanh activation function at the output of phi_x(m_ij).
-        """
-        # Call parent class
-        super(EGNN, self).__init__()
-
-        # Instance attributes
-        self.hidden_dim = hidden_dim
-        self.device = device
-        self.n_layers = n_layers
-        self.embedding_in = nn.Linear(in_node_dim, self.hidden_dim)
-        self.embedding_out = nn.Linear(self.hidden_dim, out_node_dim)
-
-        # Build architecture
-        for i in range(0, n_layers):
-            self.add_module(
-                "gcl_%d" % i,
-                E_GCL(
-                    self.hidden_dim,
-                    self.hidden_dim,
-                    self.hidden_dim,
-                    edges_in_dim=in_edge_dim,
-                    act_fn=act_fn,
-                    residual=residual,
-                    attention=attention,
-                    normalize=normalize,
-                    tanh=tanh,
-                ),
-            )
-        self.to(self.device)
-
-    # --- Core Routines ---
-    def forward(self, h, x, edge_index):
-        h = self.embedding_in(h)
-        for i in range(0, self.n_layers):
-            h, x, _ = self._modules["gcl_%d" % i](h, edge_index, x)
-        h = self.embedding_out(h)
-        return h, x
-
-
-class E_GCL(nn.Module):
-    """
-    Class that implements an equivariant convolutional layer (i.e. E(n)).
-    """
-
-    def __init__(
-        self,
-        input_dim,
-        output_dim,
-        hidden_dim,
-        edges_in_dim=0,
-        act_fn=nn.SiLU(),
-        residual=True,
-        attention=False,
-        normalize=False,
-        coords_agg="mean",
-        tanh=False,
-    ):
-        # Call parent class
-        super(E_GCL, self).__init__()
-
-        # Instance attributes
-        input_edge = input_dim * 2
-        self.residual = residual
-        self.attention = attention
-        self.normalize = normalize
-        self.coords_agg = coords_agg
-        self.tanh = tanh
-        self.epsilon = 1e-8
-        edge_coords_dim = 1
-
-        # Architecture
-        self.node_mlp = nn.Sequential(
-            nn.Linear(hidden_dim + input_dim, hidden_dim),
-            act_fn,
-            nn.Linear(hidden_dim, output_dim),
-        )
-        self.coord_mlp = nn.Sequential(
-            nn.Linear(hidden_dim, hidden_dim),
-            act_fn,
-            nn.Linear(hidden_dim, 1, bias=False),
-        )
-        self.edge_mlp = nn.Sequential(
-            nn.Linear(input_edge + edge_coords_dim + edges_in_dim, hidden_dim),
-            act_fn,
-            nn.Linear(hidden_dim, hidden_dim),
-            act_fn,
-        )
-        if self.attention:
-            self.att_mlp = nn.Sequential(
-                nn.Linear(hidden_dim, 1), nn.Sigmoid()
-            )
-
-    def edge_model(self, source, target, radial, edge_attr):
-        if edge_attr is None:
-            out = torch.cat([source, target, radial], dim=1)
-        else:
-            out = torch.cat([source, target, radial, edge_attr], dim=1)
-        out = self.edge_mlp(out)
-        if self.attention:
-            att_val = self.att_mlp(out)
-            out = out * att_val
-        return out
-
-    def node_model(self, x, edge_index, edge_attr, node_attr):
-        row, col = edge_index
-        agg = unsorted_segment_sum(edge_attr, row, num_segments=x.size(0))
-        if node_attr is not None:
-            agg = torch.cat([x, agg, node_attr], dim=1)
-        else:
-            agg = torch.cat([x, agg], dim=1)
-        out = self.node_mlp(agg)
-        if self.residual:
-            out = x + out
-        return out, agg
-
-    def coord_model(self, coord, edge_index, coord_diff, edge_feat):
-        row, col = edge_index
-        trans = coord_diff * self.coord_mlp(edge_feat)
-        if self.coords_agg == "sum":
-            agg = unsorted_segment_sum(trans, row, num_segments=coord.size(0))
-        elif self.coords_agg == "mean":
-            agg = unsorted_segment_mean(trans, row, num_segments=coord.size(0))
-        else:
-            raise Exception("Wrong coords_agg parameter" % self.coords_agg)
-        coord += agg
-        return coord
-
-    def coord2radial(self, edge_index, coord):
-        row, col = edge_index
-        coord_diff = coord[row] - coord[col]
-        radial = torch.sum(coord_diff**2, 1).unsqueeze(1)
-
-        if self.normalize:
-            norm = torch.sqrt(radial).detach() + self.epsilon
-            coord_diff = coord_diff / norm
-
-        radial = torch.zeros_like(radial, device="cuda")
-        return radial, coord_diff
-
-    def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None):
-        row, col = edge_index
-        radial, coord_diff = self.coord2radial(edge_index, coord)
-
-        edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)
-        coord = self.coord_model(coord, edge_index, coord_diff, edge_feat)
-        h, agg = self.node_model(h, edge_index, edge_feat, node_attr)
-
-        return h, coord, edge_attr
-
-
-def unsorted_segment_sum(data, segment_ids, num_segments):
-    result_shape = (num_segments, data.size(1))
-    result = data.new_full(result_shape, 0)  # Init empty result tensor.
-    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
-    result.scatter_add_(0, segment_ids, data)
-    return result
-
-
-def unsorted_segment_mean(data, segment_ids, num_segments):
-    result_shape = (num_segments, data.size(1))
-    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
-    result = data.new_full(result_shape, 0)  # Init empty result tensor.
-    count = data.new_full(result_shape, 0)
-    result.scatter_add_(0, segment_ids, data)
-    count.scatter_add_(0, segment_ids, torch.ones_like(data))
-    return result / count.clamp(min=1)
-
-
-# --- GNN Architectures ---
-class GAT(nn.Module):
-
-    def __init__(
-        self,
-        in_channels,
-        hidden_channels,
-        out_channels,
-        num_layers=2,
-        heads=4,
-        dropout=0.1,
-    ):
-        # Call parent class
-        super().__init__()
-
-        # Instance attributes
-        self.convs = nn.ModuleList()
-        self.dropout = dropout
-
-        # First layer
-        self.convs.append(
-            GATv2Conv(
-                in_channels,
-                hidden_channels,
-                heads=heads,
-                concat=True,
-                dropout=dropout,
-            )
-        )
-
-        # Hidden layers
-        for _ in range(num_layers - 2):
-            self.convs.append(
-                GATv2Conv(
-                    hidden_channels * heads,
-                    hidden_channels,
-                    heads=heads,
-                    concat=True,
-                    dropout=dropout,
-                )
-            )
-
-        # Output layer
-        self.convs.append(
-            GATv2Conv(
-                hidden_channels * heads,
-                out_channels,
-                heads=1,
-                concat=False,
-                dropout=dropout,
-            )
-        )
-
-    def forward(self, x, edge_index):
-        for conv in self.convs[:-1]:
-            x = conv(x, edge_index)
-            x = F.elu(x)
-        x = self.convs[-1](x, edge_index)
-        return x
-
-
-class GATGraphEncoder(nn.Module):
-
-    def __init__(
-        self,
-        in_channels,
-        hidden_channels,
-        out_channels,
-        heads=4,
-        num_layers=2,
-        dropout=0.2,
-    ):
-        # Call parent class
-        super().__init__()
-
-        # Instance attributes
-        self.gnn = GAT(
-            in_channels=in_channels,
-            hidden_channels=hidden_channels,
-            out_channels=hidden_channels,
-            num_layers=num_layers,
-            heads=heads,
-            dropout=dropout,
-        )
-        self.readout = nn.Linear(hidden_channels, out_channels)
-
-    def forward(self, x, edge_index):
-        x = self.gnn(x, edge_index)
-        x = x.mean(dim=0, keepdim=True)
-        return self.readout(x)
-
-
-# --- Helpers ---
-def get_edges(n_nodes):
-    rows, cols = [], []
-    for i in range(n_nodes):
-        for j in range(n_nodes):
-            if i != j:
-                rows.append(i)
-                cols.append(j)
-
-    edges = [rows, cols]
-    return edges
-
-
-def get_edges_batch(n_nodes, batch_size):
-    edges = get_edges(n_nodes)
-    edge_attr = torch.ones(len(edges[0]) * batch_size, 1)
-    edges = [torch.LongTensor(edges[0]), torch.LongTensor(edges[1])]
-    if batch_size == 1:
-        return edges, edge_attr
-    elif batch_size > 1:
-        rows, cols = [], []
-        for i in range(batch_size):
-            rows.append(edges[0] + n_nodes * i)
-            cols.append(edges[1] + n_nodes * i)
-        edges = [torch.cat(rows), torch.cat(cols)]
-    return edges, edge_attr
-
-
-def subgraph_to_data(subgraph):
-    h = torch.tensor(subgraph.node_radius[:, None], dtype=torch.float32)
-    x = torch.tensor(subgraph.node_xyz, dtype=torch.float32)
-    edges = torch.tensor(list(subgraph.edges), dtype=torch.long).T
-    return h, x, edges
-
-
-if __name__ == "__main__":
-    # Dummy parameters
-    batch_size = 8
-    n_nodes = 4
-    n_feat = 1
-    x_dim = 3
-
-    # Dummy variables h, x and fully connected edges
-    h = torch.ones(batch_size * n_nodes, n_feat)
-    x = torch.ones(batch_size * n_nodes, x_dim)
-    edges, edge_attr = get_edges_batch(n_nodes, batch_size)
-
-    # Initialize EGNN
-    egnn = EGNN(
-        in_node_dim=n_feat, hidden_dim=32, out_node_dim=1, in_edge_dim=1
-    )
-
-    # Run EGNN
-    h, x = egnn(h, x, edges)
diff --git a/src/neuron_proofreader/proposal_graph.py b/src/neuron_proofreader/proposal_graph.py
index f4bf813d..74ff9b9f 100644
--- a/src/neuron_proofreader/proposal_graph.py
+++ b/src/neuron_proofreader/proposal_graph.py
@@ -116,38 +116,22 @@ def add_proposal(self, i, j):
         self.node_proposals[j].add(i)
         self.proposals.add(frozenset({i, j}))
 
-    def generate_proposals(
-        self,
-        search_radius,
-        allow_nonleaf_proposals=False,
-        max_proposals_per_leaf=3,
-        min_size_with_proposals=0,
-    ):
+    def generate_proposals(self, config):
         """
         Generates proposals from leaf nodes.
 
         Parameters
         ----------
-        search_radius : float
-            Search radius used to generate proposals.
-        allow_nonleaf_proposals : bool, optional
-            Indication of whether to generate proposals between leaf and nodes
-            with degree 2. Default is False.
-        min_size_with_proposals : float
-            Minimum cable length (in microns) of connected components that
-            proposals are generated from. Default is 0.
+        config : ProposalConfig
+            Configuration object containing parameters for generating
+            proposals.
         """
         # Proposal generation
         assert len(self.kdtree.data) == self.number_of_nodes()
-        proposal_generator = ProposalGenerator(
-            self,
-            allow_nonleaf_proposals=allow_nonleaf_proposals,
-            max_proposals_per_leaf=max_proposals_per_leaf,
-            min_size_with_proposals=min_size_with_proposals,
-        )
-        proposals = proposal_generator(search_radius)
+        proposal_generator = ProposalGenerator(self, config)
+        proposals = proposal_generator()
 
-        self.search_radius = search_radius
+        self.search_radius = config.initial_search_radius
         self.store_proposals(proposals)
         self.trim_proposals()
 
diff --git a/src/neuron_proofreader/skeleton_graph.py b/src/neuron_proofreader/skeleton_graph.py
index cdf5a222..73c5226d 100644
--- a/src/neuron_proofreader/skeleton_graph.py
+++ b/src/neuron_proofreader/skeleton_graph.py
@@ -626,10 +626,9 @@ def clip_to_bbox(self, metadata_path):
             Path to JSON file containing origin and shape of bounding box to
             clip skeletons to.
         """
-        bucket_name, path = util.parse_cloud_path(metadata_path)
-        if util.check_gcs_exists(bucket_name, path):
+        if util.check_gcs_file_exists(metadata_path):
             # Extract bounding box
-            metadata = util.read_json_from_gcs(bucket_name, path)
+            metadata = util.read_json(metadata_path)
             origin = metadata["chunk_origin"][::-1]
             shape = metadata["chunk_shape"][::-1]
 
diff --git a/src/neuron_proofreader/split_proofreading/proposal_generation.py b/src/neuron_proofreader/split_proofreading/proposal_generation.py
index 9a70cc63..57569f4d 100644
--- a/src/neuron_proofreader/split_proofreading/proposal_generation.py
+++ b/src/neuron_proofreader/split_proofreading/proposal_generation.py
@@ -21,15 +21,7 @@ class ProposalGenerator:
     A class for generating proposals between fragments in a graph.
     """
 
-    def __init__(
-        self,
-        graph,
-        allow_nonleaf_proposals=False,
-        max_attempts=2,
-        max_proposals_per_leaf=3,
-        min_size_with_proposals=0,
-        search_scaling_factor=1.5,
-    ):
+    def __init__(self, graph, config):
         """
         Instantiates a ProposalGenerator object.
 
@@ -37,44 +29,29 @@ def __init__(
         ----------
         graph : ProposalGraph
             Graph that proposals will be generated for.
-        allow_nonleaf_proposals : bool, optional
-            Indication of whether to generate proposals between leaf and nodes
-            with degree 2. Default is False.
-        max_attempts : int, optional
-            Number of attempts made to generate proposals from a node with
-            increasing search radii. Default is 2.
-        max_proposals_per_leaf : bool, optional
-            Maximum number of proposals generated at each leaf. Default is 3.
-        min_size_with_proposals : float, optional
-            Minimum cable path length required for fragments that proposals
-            are generated from. Default is 0.
-        search_scaling_factor : 1.5, optional
-            Scaling actor used to enlarge search radius for each search.
-            Default is 2.
+        config : ProposalConfig
+            Configuration object containing parameters for generating
+            proposals.
         """
         # Instance attributes
-        self.allow_nonleaf_proposals = allow_nonleaf_proposals
+        self.allow_nonleaf_proposals = config.allow_nonleaf_proposals
         self.graph = graph
-        self.kdtree = None
-        self.max_attempts = max_attempts
-        self.max_proposals_per_leaf = max_proposals_per_leaf
-        self.min_size_with_proposals = min_size_with_proposals
-        self.search_scaling_factor = search_scaling_factor
+        self.initial_radius = config.initial_search_radius
+        self.max_attempts = config.max_attempts
+        self.max_proposals_per_leaf = config.max_proposals_per_leaf
+        self.min_size_with_proposals = config.min_size_with_proposals
+        self.search_scaling_factor = config.search_scaling_factor
 
-    def __call__(self, initial_radius):
+        # Set KD-Tree
+        self.set_kdtree()
+
+    def __call__(self):
         """
         Generates edge proposals between fragments within the given search
         radius.
-
-        Parameters
-        ----------
-        initial_radius : float
-            Initial search radius used to generate proposals between endpoints
-            of proposal.
         """
         # Initializations
-        self.set_kdtree()
-        iterator = self.graph.leaf_nodes()
+        iterator = self.valid_leafs()
         if self.graph.verbose:
             iterator = tqdm(iterator, desc="Proposal Generation")
 
@@ -82,20 +59,13 @@ def __call__(self, initial_radius):
         connections = dict()
         proposals = set()
         for leaf in iterator:
-            # Check if fragment satisfies size requirement
-            length = self.graph.cable_length(
-                max_depth=self.min_size_with_proposals, root=leaf
-            )
-            if length < self.min_size_with_proposals:
-                continue
-
             # Generate proposals
             cnt = 0
             node_candidates = list()
             while len(node_candidates) == 0 and cnt < self.max_attempts:
                 # Search for candidates
                 cnt += 1
-                radius = initial_radius * self.search_scaling_factor**cnt
+                radius = self.initial_radius * self.search_scaling_factor**cnt
                 node_candidates = self.find_node_candidates(leaf, radius)
 
                 # Parse candidates
@@ -336,6 +306,14 @@ def set_kdtree(self):
             leafs = np.array(self.graph.leaf_nodes())
             self.kdtree = KDTree(self.graph.node_xyz[leafs])
 
+    def valid_leafs(self):
+        valid_leafs = list()
+        min_sz = self.min_size_with_proposals
+        for i in self.graph.leaf_nodes():
+            if self.graph.cable_length(max_depth=min_sz, root=i) >= min_sz:
+                valid_leafs.append(i)
+        return valid_leafs
+
 
 # --- Trim Endpoints ---
 def trim_proposal_endpoints(graph, proposal, max_depth=20):
diff --git a/src/neuron_proofreader/split_proofreading/split_datamodules.py b/src/neuron_proofreader/split_proofreading/split_datamodules.py
new file mode 100644
index 00000000..d2eed972
--- /dev/null
+++ b/src/neuron_proofreader/split_proofreading/split_datamodules.py
@@ -0,0 +1,262 @@
+"""
+Created on Fri April 11 11:00:00 2024
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Implementation of dataset objects that contain graphs and facilitate feature
+generation for training and inference in split-correction tasks.
+
+"""
+
+from collections import deque
+from concurrent.futures import ThreadPoolExecutor, FIRST_COMPLETED, wait
+from queue import Queue
+from threading import Thread
+from torch.utils.data import Dataset, IterableDataset
+
+import numpy as np
+import os
+import random
+
+from neuron_proofreader.machine_learning.subgraph_sampler import (
+    SubgraphSampler,
+)
+from neuron_proofreader.split_proofreading.split_feature_extraction import (
+    FeaturePipeline,
+    HeteroGraphData,
+)
+from neuron_proofreader.utils import util
+
+
+# --- Datasets ---
+class FragmentsDataset(IterableDataset):
+    """
+    A dataset object that contains a graph built from fragments corresponding
+    to a single brain.
+    """
+
+    def __init__(self, fragments_graph, img_config):
+        """
+        Instantiates a FragmentsDataset object.
+
+        Parameters
+        ----------
+        fragments_graph : str
+            ...
+        img_config : ImageConfig
+            Configuration object containing image parameters.
+        """
+        self.graph = fragments_graph
+        self.feature_extractor = FeaturePipeline(
+            self.graph,
+            img_config.img_path,
+            brightness_clip=img_config.brightness_clip,
+            patch_shape=img_config.patch_shape,
+        )
+
+    # --- Get Data ---
+    def __iter__(self):
+        """
+        Iterates over the dataset and yields model-ready inputs and targets.
+
+        Yields
+        ------
+        inputs : HeteroGraphData
+            Input data.
+        targets : torch.Tensor
+            Ground truth labels.
+        """
+        for subgraph in self.sampler():
+            yield HeteroGraphData(self.feature_extractor(subgraph))
+
+    # --- Helpers ---
+    def __getattr__(self, name):
+        return getattr(self.graph, name)
+
+    def sampler(self, batch_size):
+        """
+        Gets a subgraph sampler used to iterate over dataset.
+
+        Returns
+        -------
+        sampler : SubgraphSampler
+            Subgraph sampler that is used to iterate over dataset.
+        """
+        return iter(SubgraphSampler(self.graph, max_proposals=batch_size))
+
+
+class FragmentsCollection(Dataset):
+    """
+    Stores FragmentDataset objects for a set of whole-brain images and
+    manages proposal generation and dataset-level statistics.
+    """
+
+    def __init__(self):
+        """
+        Parameters
+        ----------
+        datasets : Dict[Hashable, FragmentDataset]
+            Mapping from dataset key to FragmentDataset.
+        """
+        self.datasets = dict()
+
+    def add(self, key, dataset):
+        assert key not in self.datasets, "Dataset already exists!"
+        self.datasets[key] = dataset
+
+    def __getitem__(self, key):
+        return self.datasets[key]
+
+    # --- Helpers ---
+    def __getattr__(self, name):
+        return getattr(self.datasets, name)
+
+    def __len__(self):
+        return len(self.datasets)
+
+    def __repr__(self):
+        n_p = np.sum([ds.n_proposals() for ds in self.values()])
+        n_a = np.sum([len(ds.gt_accepts) for ds in self.values()])
+        return (
+            f"FragmentsCollection("
+            f"num_datasets={len(self)}, "
+            f"num_proposals={n_p}, "
+            f"percent_accepts={100 * n_a / (n_p + 1e-5):.2f})"
+        )
+
+
+class FragmentsLoader:
+    """
+    Prefetching loader that samples subgraphs from a FragmentsCollection.
+    A bounded pool of worker threads pulls batches from datasets in
+    round-robin order, overlapping I/O across up to `max_workers` datasets
+    at once, and feeds a bounded queue consumed by the training loop.
+    """
+
+    def __init__(self, collection, batch_size=32, prefetch=4, max_workers=4, shuffle=True):
+        """
+        Parameters
+        ----------
+        collection : FragmentsCollection
+        batch_size : int, optional
+            Max proposals per subgraph sampled from each dataset. Default is 32.
+        prefetch : int, optional
+            Number of batches to buffer ahead of consumption. Default is 4.
+        max_workers : int, optional
+            Number of datasets that may be fetched from concurrently. Bounded
+            regardless of collection size. Default is 4.
+        shuffle : bool, optional
+            If True, datasets are visited in random round-robin order each
+            epoch. Default is True.
+        """
+        self.batch_size = batch_size
+        self.collection = collection
+        self.prefetch = prefetch
+        self.max_workers = max_workers
+        self.shuffle = shuffle
+
+    def __iter__(self):
+        queue = Queue(maxsize=self.prefetch)
+        thread = Thread(target=self._produce, args=(queue,), daemon=True)
+        thread.start()
+
+        while True:
+            item = queue.get()
+            if item is StopIteration:
+                return
+            if isinstance(item, Exception):
+                raise item
+            yield item
+
+    def _produce(self, queue):
+        try:
+            samplers = self._init_samplers()
+            rotation = deque(samplers.keys())
+            if self.shuffle:
+                random.shuffle(rotation)
+
+            n_workers = min(self.max_workers, len(rotation)) or 1
+            with ThreadPoolExecutor(max_workers=n_workers) as executor:
+                pending = dict()  # future -> key
+                for _ in range(n_workers):
+                    self._submit_next(executor, rotation, samplers, pending)
+
+                while pending:
+                    done, _ = wait(pending.keys(), return_when=FIRST_COMPLETED)
+                    for future in done:
+                        key = pending.pop(future)
+                        batch = future.result()
+                        if batch is not None:
+                            queue.put(batch)
+                            rotation.append(key)  # dataset has more, requeue
+                        self._submit_next(executor, rotation, samplers, pending)
+        except Exception as e:
+            queue.put(e)
+            return
+        queue.put(StopIteration)
+
+    def _submit_next(self, executor, rotation, samplers, pending):
+        if not rotation:
+            return
+        key = rotation.popleft()
+        future = executor.submit(self._fetch_batch, key, samplers[key])
+        pending[future] = key
+
+    def _fetch_batch(self, key, sampler):
+        try:
+            subgraph = next(sampler)
+        except StopIteration:
+            return None
+        dataset = self.collection[key]
+        features = dataset.feature_extractor(subgraph)
+        data = HeteroGraphData(features)
+        return data.get_inputs(), data.get_targets()
+
+    def _init_samplers(self):
+        samplers = dict()
+        for key, dataset in self.collection.datasets.items():
+            samplers[key] = iter(
+                SubgraphSampler(dataset.graph, max_proposals=self.batch_size)
+            )
+        return samplers
+
+    def __len__(self):
+        batch_cnt = 0
+        for ds in self.collection.datasets.values():
+            batch_cnt += np.ceil(ds.n_proposals() / self.batch_size)
+        return int(batch_cnt)
+
+
+# -- Helpers --
+def generate_example_ids(ds_prefix):
+    """
+    Generates dataset example identifiers.
+
+    Parameters
+    ----------
+    ds_prefix : str
+        Root prefix under which datasets are organized.
+
+    Yields
+    -------
+    Tuple[str]
+        Dataset ID formatted as (brain_id, segmentation_id, block_id).
+    """
+    bucket_name, _ = util.parse_cloud_path(ds_prefix)
+    for brain_prefix in util.list_gcs_subprefixes(ds_prefix):
+        # Extract brain id
+        brain_id = brain_prefix.split("/")[-2]
+
+        # Iterate over segmentations
+        pred_prefix = os.path.join(bucket_name, brain_prefix, "pred_swcs")
+        for brain_seg_prefix in util.list_gcs_subprefixes(pred_prefix):
+            # Extract segmentation id
+            segmentation_id = brain_seg_prefix.split("/")[-2]
+
+            # Iterate over blocks
+            ex_prefix = os.path.join(bucket_name, brain_seg_prefix)
+            for block_prefix in util.list_gcs_subprefixes(ex_prefix):
+                # Extract block id
+                block_id = block_prefix.split("/")[-2]
+                yield brain_id, segmentation_id, block_id
diff --git a/src/neuron_proofreader/split_proofreading/split_datasets.py b/src/neuron_proofreader/split_proofreading/split_datasets.py
deleted file mode 100644
index f19dbb36..00000000
--- a/src/neuron_proofreader/split_proofreading/split_datasets.py
+++ /dev/null
@@ -1,299 +0,0 @@
-"""
-Created on Fri April 11 11:00:00 2024
-
-@author: Anna Grim
-@email: anna.grim@alleninstitute.org
-
-Implementation of dataset objects that contain graphs and facilitate feature
-generation for training and inference in split-correction tasks.
-
-"""
-
-from concurrent.futures import ThreadPoolExecutor
-from queue import Queue
-from torch.utils.data import IterableDataset
-from tqdm import tqdm
-
-import os
-import pandas as pd
-
-from neuron_proofreader.machine_learning.subgraph_sampler import (
-    SubgraphSampler,
-)
-from neuron_proofreader.split_proofreading.split_feature_extraction import (
-    FeaturePipeline,
-    HeteroGraphData,
-)
-from neuron_proofreader.utils import util
-
-
-# --- Datasets ---
-class FragmentsDataset(IterableDataset):
-    """
-    A dataset object that contains a graph built from fragments corresponding
-    to a single brain.
-    """
-
-    def __init__(self, fragments_graph, img_config, proposals_per_batch=64):
-        """
-        Instantiates a FragmentsDataset object.
-
-        Parameters
-        ----------
-        fragments_graph : str
-            ...
-        img_config : ImageConfig
-            Configuration object containing image parameters.
-        proposals_per_batch : int, optional
-            Maximum number of proposals in subgraphs yielded per batch.
-            Default is 64.
-        """
-        self.graph = fragments_graph
-        self.feature_extractor = FeaturePipeline(
-            self.graph,
-            img_config.img_path,
-            brightness_clip=img_config.brightness_clip,
-            patch_shape=img_config.patch_shape,
-        )
-        self.proposals_per_batch = proposals_per_batch
-
-    # --- Get Data ---
-    def __iter__(self):
-        """
-        Iterates over the dataset and yields model-ready inputs and targets.
-
-        Yields
-        ------
-        inputs : HeteroGraphData
-            Input data.
-        targets : torch.Tensor
-            Ground truth labels.
-        """
-        for subgraph in self.get_sampler():
-            features = self.feature_extractor(subgraph)
-            yield HeteroGraphData(features)
-
-    # --- Helpers ---
-    def __getattr__(self, name):
-        return getattr(self.graph, name)
-
-    def get_sampler(self):
-        """
-        Gets a subgraph sampler used to iterate over dataset.
-
-        Returns
-        -------
-        sampler : SubgraphSampler
-            Subgraph sampler that is used to iterate over dataset.
-        """
-        sampler = SubgraphSampler(
-            self.graph, max_proposals=self.proposals_per_batch
-        )
-        return iter(sampler)
-
-
-class FragmentsDatasetCollection(IterableDataset):
-    """
-    A dataset class for storing a set of FragmentDataset objects corresponding
-    to different whole-brain images.
-    """
-
-    def __init__(self, datasets, shuffle=True):
-        """
-        Instantiates a FragmentsDatasetCollection object.
-
-        Parameters
-        ----------
-        shuffle : bool, optional
-            True if examples should be shuffled each epoch. Default is True.
-        """
-        # Instance attributes
-        self.datasets = datasets
-        self.shuffle = shuffle
-
-    def __iter__(self):
-        """
-        Iterates over the datasets and yields model-ready inputs and targets.
-
-        Yields
-        ------
-        inputs : TensorDict
-            Heterogeneous graph data.
-        targets : torch.Tensor
-            Ground truth labels.
-        """
-        # Create prefetching data structures
-        samplers = self.init_samplers()
-        queue = Queue(maxsize=self.prefetch * len(samplers))
-        active_keys = set(samplers.keys())
-
-        # Launch one prefetch thread per dataset
-        with ThreadPoolExecutor(max_workers=len(samplers)) as executor:
-            for key, sampler in samplers.items():
-                executor.submit(self._worker, key, sampler, queue)
-
-            # Consume from queue until all datasets exhausted
-            while active_keys:
-                key, inputs, targets = queue.get()
-                if inputs is StopIteration:
-                    active_keys.discard(key)
-                    continue
-                if isinstance(inputs, Exception):
-                    raise inputs
-                yield inputs, targets
-
-    def _worker(self, key, sampler, queue):
-        """
-        Runs in a background thread, prefetches extracted features into queue.
-        """
-        try:
-            for subgraph in sampler:
-                features = self.datasets[key].feature_extractor(subgraph)
-                data = HeteroGraphData(features)
-                queue.put((key, data.get_inputs(), data.get_targets()))
-        except Exception as e:
-            queue.put((key, e, None))
-        finally:
-            queue.put((key, StopIteration, None))
-
-    def generate_proposals(self, proposal_config):
-        """
-        Generates proposals for each dataset.
-
-        Parameters
-        ----------
-        proposal_config : ProposalConfig
-            Configuration object containing parameters for generating
-            proposals.
-        """
-        for key in tqdm(self.datasets, desc="Generate Proposals"):
-            self.datasets[key].graph.generate_proposals(
-                proposal_config.search_radius,
-                **proposal_config,
-            )
-
-    # --- Helpers ---
-    def __len__(self):
-        """
-        Returns the number of datasets in self.
-
-        Returns
-        -------
-        float
-            Number of datasets in self.
-        """
-        return len(self.datasets)
-
-    def get_next_key(self, samplers):
-        """
-        Gets the next key to sample from a dictionary of samplers.
-
-        Parameters
-        ----------
-        samplers : Dict[Tuple[str], SubgraphSampler]
-            Mapping from keys to sampler objects.
-
-        Returns
-        -------
-        key : Tuple[str]
-            Selected key from "samplers".
-        """
-        if self.shuffle:
-            return util.sample_once(samplers.keys())
-        else:
-            keys = sorted(samplers.keys())
-            return keys[0]
-
-    def init_samplers(self):
-        """
-        Initializes subgraph samplers for each dataset.
-
-        Returns
-        -------
-        samplers : Dict[hashable, SubgraphSampler]
-            Subgraph samplers used to iterate over the datasets.
-        """
-        samplers = dict()
-        for key, dataset in self.datasets.items():
-            batch_size = dataset.config.ml.batch_size
-            samplers[key] = iter(
-                SubgraphSampler(dataset.graph, max_proposals=batch_size)
-            )
-        return samplers
-
-    def n_proposals(self):
-        """
-        Counts the number of proposals in the dataset.
-
-        Returns
-        -------
-        int
-            Number of proposals.
-        """
-        cnt = 0
-        for dataset in self.datasets.values():
-            cnt += dataset.graph.n_proposals()
-        return cnt
-
-    def p_accepts(self):
-        """
-        Computes the percentage of accepted proposals in ground truth.
-
-        Returns
-        -------
-        float
-            Percentage of accepted proposals in ground truth.
-        """
-        accepts_cnt = 0
-        for dataset in self.datasets.values():
-            accepts_cnt += len(dataset.graph.gt_accepts)
-        return 100 * accepts_cnt / (self.n_proposals() + 1e-5)
-
-    def save_examples_summary(self, path):
-        """
-        Saves a summary of examples in the dataset to the given path.
-
-        Parameters
-        ----------
-        path : str
-            Output path for the CSV file.
-        """
-        examples_summary = list()
-        for key in sorted(self.datasets.keys()):
-            n_proposals = self.datasets[key].graph.n_proposals()
-            examples_summary.extend([key] * n_proposals)
-        pd.DataFrame(examples_summary).to_csv(path)
-
-
-# -- Helpers --
-def generate_example_ids(ds_prefix):
-    """
-    Generates dataset example identifiers.
-
-    Parameters
-    ----------
-    ds_prefix : str
-        Root prefix under which datasets are organized.
-
-    Yields
-    -------
-    Tuple[str]
-        Dataset ID formatted as (brain_id, segmentation_id, block_id).
-    """
-    bucket_name, _ = util.parse_cloud_path(ds_prefix)
-    for brain_prefix in util.list_gcs_subprefixes(ds_prefix):
-        # Extract brain id
-        brain_id = brain_prefix.split("/")[-2]
-
-        # Iterate over segmentations
-        pred_prefix = os.path.join(bucket_name, brain_prefix, "pred_swcs")
-        for brain_seg_prefix in util.list_gcs_subprefixes(pred_prefix):
-            # Extract segmentation id
-            segmentation_id = brain_seg_prefix.split("/")[-2]
-
-            # Iterate over blocks
-            ex_prefix = os.path.join(bucket_name, brain_seg_prefix)
-            for block_prefix in util.list_gcs_subprefixes(ex_prefix):
-                # Extract block id
-                block_id = block_prefix.split("/")[-2]
-                yield brain_id, segmentation_id, block_id
diff --git a/src/neuron_proofreader/split_proofreading/split_feature_extraction.py b/src/neuron_proofreader/split_proofreading/split_feature_extraction.py
index 73302215..7c9e8c98 100644
--- a/src/neuron_proofreader/split_proofreading/split_feature_extraction.py
+++ b/src/neuron_proofreader/split_proofreading/split_feature_extraction.py
@@ -82,6 +82,18 @@ def __call__(self, subgraph):
             extractor(subgraph, features)
         return features
 
+    def close(self):
+        """Shuts down any persistent resources held by extractors."""
+        for extractor in self.extractors:
+            if hasattr(extractor, "close"):
+                extractor.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
 
 class SkeletonFeatureExtractor:
     """
@@ -156,7 +168,6 @@ def extract_edge_features(self, subgraph, features):
             edge_features[edge] = np.array(
                 [
                     np.mean(self.graph.node_radius[path]),
-                    min(self.graph.path_length(path), 5000) / 5000,
                 ],
             )
         features.set_features(edge_features, "edge")
@@ -246,6 +257,7 @@ def __init__(
             point of proposal for image patch extraction. Default is 40.
         """
         self.brightness_clip = brightness_clip
+        self.executor = ThreadPoolExecutor()
         self.graph = graph
         self.img = TensorStoreImage(img_path)
         self.patch_shape = patch_shape
@@ -262,21 +274,19 @@ def __call__(self, subgraph, features):
         features : FeatureSet
             Data structure that stores features.
         """
-        with ThreadPoolExecutor() as executor:
-            # Assign threads
-            pending = dict()
-            for proposal in subgraph.proposals:
-                thread = executor.submit(self.init_extractor, proposal)
-                pending[thread] = proposal
-
-            # Store results
-            patches, profiles = dict(), dict()
-            for thread in as_completed(pending.keys()):
-                proposal = pending.pop(thread)
-                extractor = thread.result()
+        # Assign threads
+        pending = {
+            self.executor.submit(self.init_extractor, proposal): proposal
+            for proposal in subgraph.proposals
+        }
 
-                profiles[proposal] = extractor.get_intensity_profile()
-                patches[proposal] = extractor.get_input_patch()
+        # Store results
+        patches, profiles = dict(), dict()
+        for thread in as_completed(pending.keys()):
+            proposal = pending.pop(thread)
+            extractor = thread.result()
+            profiles[proposal] = extractor.get_intensity_profile()
+            patches[proposal] = extractor.get_input_patch()
 
         # Update features
         features.set_features(patches, "proposal_patches")
@@ -312,6 +322,19 @@ def init_extractor(self, proposal):
         return extractor
 
     # --- Helpers ---
+    def close(self):
+        """
+        Shuts down the persistent thread pool. Call when this extractor is
+        no longer needed.
+        """
+        self.executor.shutdown(wait=True)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
     def create_segment_mask(self, proposal, shape, offset):
         # Find nearby nodes
         center = self.graph.proposal_midpoint(proposal)
@@ -582,7 +605,7 @@ class FeatureSet:
         "proposal": ("proposal_features", "proposal_index_mapping"),
         "proposal_patches": ("proposal_patches", "proposal_index_mapping"),
     }
-    n_branch_features = 2
+    n_branch_features = 1
     n_proposal_features = 70
 
     def __init__(self, graph):
@@ -963,7 +986,7 @@ def get_feature_dict():
         Dictionary that contains the number of features for branchs and
         proposals.
     """
-    return {"branch": 2, "proposal": 67}
+    return {"branch": 1, "proposal": 67}
 
 
 def resize_segmentation(mask, new_shape):
diff --git a/src/neuron_proofreader/split_proofreading/split_inference.py b/src/neuron_proofreader/split_proofreading/split_inference.py
index e9055d38..10503eee 100644
--- a/src/neuron_proofreader/split_proofreading/split_inference.py
+++ b/src/neuron_proofreader/split_proofreading/split_inference.py
@@ -31,7 +31,7 @@
 import os
 import torch
 
-from neuron_proofreader.split_proofreading.split_datasets import (
+from neuron_proofreader.split_proofreading.split_datamodules import (
     FragmentsDataset,
 )
 from neuron_proofreader.utils import ml_util, util
diff --git a/src/neuron_proofreader/utils/ml_util.py b/src/neuron_proofreader/utils/ml_util.py
index fa428978..562aa6e3 100644
--- a/src/neuron_proofreader/utils/ml_util.py
+++ b/src/neuron_proofreader/utils/ml_util.py
@@ -144,6 +144,8 @@ def update(self, pred, y, loss):
         self.loss += loss.item()
         self.n += y.numel()
 
+        del pred, y, loss
+
     def compute(self):
         precision = self.tp / (self.tp + self.fp + 1e-8)
         recall = self.tp / (self.tp + self.fn + 1e-8)
diff --git a/src/neuron_proofreader/utils/util.py b/src/neuron_proofreader/utils/util.py
index 31c65caf..1a7aecd8 100644
--- a/src/neuron_proofreader/utils/util.py
+++ b/src/neuron_proofreader/utils/util.py
@@ -217,8 +217,11 @@ def read_json(path):
     dict
         Contents of JSON file.
     """
-    with open(path, "r") as f:
-        return json.load(f)
+    if is_gcs_path(path):
+        return json.loads(read_gcs_txt(path))
+    else:
+        with open(path, "r") as f:
+            return json.load(f)
 
 
 def read_txt(path, client=None):
@@ -338,9 +341,9 @@ def get_google_swcs_prefix(root_prefix, brain_id, segmentation_id):
     # Determine old vs. new result
     prefix1 = os.path.join(root_prefix, brain_id, "whole_brain")
     prefix2 = os.path.join(root_prefix, "whole_brain", brain_id)
-    if check_gcs_exists(prefix1, is_prefix=True):
+    if check_gcs_prefix_exists(prefix1, is_prefix=True):
         prefix = prefix1
-    elif check_gcs_exists(prefix2, is_prefix=True):
+    elif check_gcs_prefix_exists(prefix2, is_prefix=True):
         prefix = prefix2
     else:
         raise Exception("Unable to find Google swcs result!")
@@ -411,16 +414,15 @@ def parse_cloud_path(path):
 
 
 # --- GCS Utils ---
-def check_gcs_exists(path, is_prefix=False):
+def check_gcs_file_exists(path, is_prefix=False):
     """
-    Checks if a file or prefix exists in GCS.
+    Checks if a file exists at the given GCS path.
+
     Parameters
     ----------
     path : str
         GCS path to check.
-    prefix : bool
-        If True, checks whether any object exists under the given prefix.
-        If False, checks whether the exact file exists.
+
     Returns
     -------
     bool
@@ -428,17 +430,13 @@ def check_gcs_exists(path, is_prefix=False):
     """
     bucket_name, key = parse_cloud_path(path)
     bucket = storage.Client().bucket(bucket_name)
-    if is_prefix:
-        key = key.rstrip("/") + "/"
-        return any(bucket.list_blobs(prefix=key, max_results=1))
-    else:
-        return bucket.blob(key).exists()
+    return bucket.blob(key).exists()
 
 
 def check_gcs_prefix_exists(path):
     bucket_name, prefix = parse_cloud_path(path)
-    prefix = prefix.rstrip("/") + "/"
     bucket = storage.Client().bucket(bucket_name)
+    prefix = prefix.rstrip("/") + "/"
     exists = any(bucket.list_blobs(prefix=prefix, max_results=1))
     return exists
 
@@ -483,7 +481,7 @@ def list_gcs_paths(bucket_name, prefix, extension=""):
     for name in [b.name for b in bucket.list_blobs(prefix=prefix)]:
         if extension in name:
             paths.append(os.path.join(f"gs://{bucket_name}", name))
-    return paths
+    return sorted(paths)
 
 
 def list_gcs_subprefixes(path):
@@ -511,13 +509,12 @@ def list_gcs_subprefixes(path):
 
     # Parse directory contents
     prefix_depth = len(prefix.split("/"))
-    subdirs = list()
+    subprefixes = list()
     for prefix in blobs.prefixes:
-        is_dir = prefix.endswith("/")
-        is_direct_subdir = len(prefix.split("/")) - 1 == prefix_depth
-        if is_dir and is_direct_subdir:
-            subdirs.append(prefix)
-    return subdirs
+        is_direct = len(prefix.split("/")) - 1 == prefix_depth
+        if prefix.endswith("/") and is_direct:
+            subprefixes.append(prefix)
+    return sorted(subprefixes)
 
 
 def read_gcs_txt(prefix, client=None):
@@ -526,7 +523,7 @@ def read_gcs_txt(prefix, client=None):
 
     Parameters
     ----------
-    path : str
+    prefix : str
         Path to txt file to be read.
 
     Returns
@@ -540,28 +537,6 @@ def read_gcs_txt(prefix, client=None):
     return bucket.blob(subprefix).download_as_text()
 
 
-def read_json_from_gcs(bucket_name, blob_path):
-    """
-    Reads JSON file stored in a GCS bucket.
-
-    Parameters
-    ----------
-    bucket_name : str
-        Name of the GCS bucket containing the JSON file.
-    blob_path : str
-        Path to the JSON file within the GCS bucket.
-
-    Returns
-    -------
-    dict
-        Parsed JSON content as a Python dictionary.
-    """
-    client = storage.Client()
-    bucket = client.bucket(bucket_name)
-    blob = bucket.blob(blob_path)
-    return json.loads(blob.download_as_text())
-
-
 # --- S3 Utils ---
 def is_s3_path(path):
     """

From a58207766fc3d92c42b91138a4ef551db31691d2 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Fri, 10 Jul 2026 22:33:38 +0000
Subject: [PATCH 3/3] reverted branch features

---
 .../split_proofreading/split_feature_extraction.py           | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/neuron_proofreader/split_proofreading/split_feature_extraction.py b/src/neuron_proofreader/split_proofreading/split_feature_extraction.py
index 7c9e8c98..e79b054c 100644
--- a/src/neuron_proofreader/split_proofreading/split_feature_extraction.py
+++ b/src/neuron_proofreader/split_proofreading/split_feature_extraction.py
@@ -168,6 +168,7 @@ def extract_edge_features(self, subgraph, features):
             edge_features[edge] = np.array(
                 [
                     np.mean(self.graph.node_radius[path]),
+                    min(self.graph.path_length(path), 5000) / 5000,
                 ],
             )
         features.set_features(edge_features, "edge")
@@ -605,7 +606,7 @@ class FeatureSet:
         "proposal": ("proposal_features", "proposal_index_mapping"),
         "proposal_patches": ("proposal_patches", "proposal_index_mapping"),
     }
-    n_branch_features = 1
+    n_branch_features = 2
     n_proposal_features = 70
 
     def __init__(self, graph):
@@ -986,7 +987,7 @@ def get_feature_dict():
         Dictionary that contains the number of features for branchs and
         proposals.
     """
-    return {"branch": 1, "proposal": 67}
+    return {"branch": 2, "proposal": 67}
 
 
 def resize_segmentation(mask, new_shape):