feat: graph builder class

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -23,7 +23,7 @@
 import numpy as np
 import os
 
-from segmentation_skeleton_metrics import graph_segmentation_alignment as gsa
+from segmentation_skeleton_metrics import split_detection
 from segmentation_skeleton_metrics.utils import (
     graph_util as gutil,
     img_util,
@@ -112,26 +112,20 @@ def __init__(
         self.output_dir = output_dir
         self.preexisting_merges = preexisting_merges
 
-        # Load ground truth
-        print("\n(1) Load Ground Truth")
-        assert type(valid_labels) is set if valid_labels else True
+        # Load data
+        assert isinstance(valid_labels, set) if valid_labels else True
         self.label_mask = pred_labels
         self.valid_labels = valid_labels
         self.init_label_map(connections_path)
-        self.init_graphs(gt_pointer)
-
-        print("\n(2) Load Prediction")
-        if fragments_pointer:
-            self.load_fragments(fragments_pointer)
-        else:
-            self.fragment_graphs = None
+        self.load_groundtruth(gt_pointer)
+        self.load_fragments(fragments_pointer)
 
         # Initialize writer
         self.save_projections = save_projections
         if self.save_projections:
             self.init_zip_writer()
 
-    # -- Initialize and Label Graphs --
+    # --- Load Data ---
     def init_label_map(self, path):
         """
         Initializes a dictionary that maps a label to its equivalent label in
@@ -157,7 +151,7 @@ def init_label_map(self, path):
             self.label_map = None
             self.inverse_label_map = None
 
-    def init_graphs(self, paths):
+    def load_groundtruth(self, swc_pointer):
         """
         Initializes "self.graphs" by iterating over "paths" which corresponds
         to neurons in the ground truth.
@@ -174,8 +168,13 @@ def init_graphs(self, paths):
 
         """
         # Build graphs
-        swc_dicts = swc_util.Reader().load(paths)
-        self.graphs = self.build_graphs(swc_dicts)
+        print("\n(1) Load Ground Truth")
+        graph_builder = gutil.GraphBuilder(
+            anisotropy=self.anisotropy,
+            label_mask=self.label_mask,
+            use_anisotropy=False,
+        )
+        self.graphs = graph_builder.run(swc_pointer)
 
         # Label nodes
         self.key_to_label_to_nodes = dict()  # {id: {label: nodes}}
@@ -185,23 +184,18 @@ def init_graphs(self, paths):
                 self.graphs[key]
             )
 
-    def build_graphs(self, swc_dicts):
-        graphs = dict()
-        with ProcessPoolExecutor() as executor:
-            # Assign processes
-            processes = list()
-            for swc_dict in swc_dicts:
-                processes.append(
-                    executor.submit(gutil.to_graph, swc_dict)
-                )
+    def load_fragments(self, swc_pointer):
+        print("\n(2) Load Fragments")
+        if swc_pointer:
+            graph_builder = gutil.GraphBuilder(
+                anisotropy=self.anisotropy,
+                selected_ids=self.get_all_node_labels(),
+                use_anisotropy=True,
+            )
+            self.fragment_graphs = graph_builder.run(swc_pointer)
+        else:
+            self.fragment_graphs = None
 
-            # Store results
-            pbar = tqdm(total=len(processes), desc="Build Graphs")
-            for process in as_completed(processes):
-                graphs.update(process.result())
-                pbar.update(1)
-        return graphs
-            
     def label_graphs(self, key, batch_size=128):
         """
         Iterates over nodes in "graph" and stores the corresponding label from
@@ -259,7 +253,7 @@ def get_patch_labels(self, key, nodes):
         # Get bounding box
         bbox = {"min": [np.inf, np.inf, np.inf], "max": [0, 0, 0]}
         for i in nodes:
-            voxel = deepcopy(self.graphs[key].graph["voxel"][i])
+            voxel = self.graphs[key].graph["voxel"][i]
             for idx in range(3):
                 if voxel[idx] < bbox["min"][idx]:
                     bbox["min"][idx] = voxel[idx]
@@ -363,39 +357,6 @@ def get_node_labels(self, key, inverse_bool=False):
         else:
             return set(self.key_to_label_to_nodes[key].keys())
 
-    # -- Load Fragments --
-    def load_fragments(self, fragments_pointer):
-        """
-        Loads and filters swc files from a local zip. These swc files are
-        assumed to be fragments from a predicted segmentation.
-
-        Parameters
-        ----------
-        zip_path : str
-            Path to the local zip file containing the fragments
-
-        Returns
-        -------
-        dict
-            Dictionary that maps an swc id to the fragment graph.
-
-        """
-        # Read SWC files
-        reader = swc_util.Reader(anisotropy=self.anisotropy)
-        swc_dicts = deque(reader.load(fragments_pointer))
-
-        # Filter SWC files
-        filtered_swc_dicts = list()
-        labels = self.get_all_node_labels()
-        while len(swc_dicts) > 0:
-            swc_dict = swc_dicts.popleft()
-            swc_id = int(swc_dict["swc_id"])
-            if swc_id in labels:
-                swc_dict["swc_id"] = swc_id
-                filtered_swc_dicts.append(swc_dict)
-        self.fragment_graphs = self.build_graphs(filtered_swc_dicts)
-        print("# Fragments:", len(self.fragment_graphs))
-
     def init_zip_writer(self):
         """
         Initializes "self.zip_writer" attribute by setting up a directory for
@@ -513,7 +474,7 @@ def detect_splits(self):
             for key, graph in self.graphs.items():
                 processes.append(
                        executor.submit(
-                           gsa.correct_graph_misalignments,
+                           split_detection.run,
                            key,
                            graph,
                        )

…_metrics/graph_segmentation_alignment.py …tion_skeleton_metrics/split_detection.pysrc/segmentation_skeleton_metrics/graph_segmentation_alignment.py renamed to src/segmentation_skeleton_metrics/split_detection.py src/segmentation_skeleton_metrics/graph_segmentation_alignment.py renamed to src/segmentation_skeleton_metrics/split_detection.py

@@ -16,7 +16,7 @@
 from segmentation_skeleton_metrics.utils import graph_util as gutil
 
 
-def correct_graph_misalignments(process_id, graph):
+def run(process_id, graph):
     """
     Adjusts misalignments between ground truth graph and segmentation mask.
 

src/segmentation_skeleton_metrics/utils/graph_util.py

@@ -9,62 +9,111 @@
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, as_completed
 from random import sample
+from tqdm import tqdm
 
 import networkx as nx
 import numpy as np
 from scipy.spatial import distance
 
-from segmentation_skeleton_metrics.utils import img_util
+from segmentation_skeleton_metrics.utils import img_util, swc_util
 
 ANISOTROPY = np.array([0.748, 0.748, 1.0])
 
 
-def to_graph(swc_dict):
-    """
-    Builds a graph from a dictionary that contains the contents of an SWC
-    file.
-
-    Parameters
-    ----------
-    swc_dict : dict
-        ...
-
-    Returns
-    -------
-    networkx.Graph
-        Graph built from an SWC file.
-
-    """
-    # Initializations
-    old_to_new = dict()
-    run_length = 0
-    voxels = np.zeros((len(swc_dict["id"]), 3), dtype=np.int32)
-
-    # Build graph
-    graph = nx.Graph()
-    for i in range(len(swc_dict["id"])):            
-        # Get node id
-        old_id = swc_dict["id"][i]
-        old_to_new[old_id] = i
-
-        # Update graph
-        voxels[i] = swc_dict["voxel"][i]
-        if swc_dict["pid"][i] != -1:
-            # Add edge
-            parent = old_to_new[swc_dict["pid"][i]]
-            graph.add_edge(i, parent)
-
-            # Update run length
-            xyz_i = voxels[i] * ANISOTROPY
-            xyz_p = voxels[parent] * ANISOTROPY
-            run_length += distance.euclidean(xyz_i, xyz_p)
-
-    # Set graph-level attributes
-    graph.graph["n_edges"] = graph.number_of_edges()
-    graph.graph["run_length"] = run_length
-    graph.graph["voxel"] = voxels
-    return {swc_dict["swc_id"]: graph}
-
+class GraphBuilder:
+    
+    def __init__(
+        self,
+        anisotropy=(1.0, 1.0, 1.0),
+        label_mask=None,
+        selected_ids=None,
+        use_anisotropy=True,
+    ):
+        # Instance attributes
+        self.anisotropy = anisotropy
+        self.label_mask = label_mask
+        self.selected_ids = selected_ids
+
+        # Reader
+        anisotropy = anisotropy if use_anisotropy else (1.0, 1.0, 1.0)
+        self.swc_reader = swc_util.Reader(anisotropy)
+
+    def run(self, swc_pointer):
+        self._build_graphs_from_swcs(swc_pointer)
+        self._label_graphs_with_segmentation()
+        return self.graphs
+
+    # --- Build Graphs ---
+    def _build_graphs_from_swcs(self, swc_pointer):
+        with ProcessPoolExecutor() as executor:
+            # Assign processes
+            processes = list()
+            for swc_dict in self.swc_reader.load(swc_pointer):
+                if self._process_swc_dict(swc_dict["swc_id"]):
+                    processes.append(executor.submit(self.to_graph, swc_dict))
+
+            # Store results
+            self.graphs = dict()
+            pbar = tqdm(total=len(processes), desc="Build Graphs")
+            for process in as_completed(processes):
+                self.graphs.update(process.result())
+                pbar.update(1)
+
+    def _process_swc_dict(self, swc_id):
+        if self.selected_ids:
+            segment_id = get_segment_id(swc_id)
+            if segment_id not in self.selected_ids:
+                return False
+        return True
+
+    def to_graph(self, swc_dict):
+        """
+        Builds a graph from a dictionary that contains the contents of an SWC
+        file.
+
+        Parameters
+        ----------
+        swc_dict : dict
+            ...
+
+        Returns
+        -------
+        networkx.Graph
+            Graph built from an SWC file.
+
+        """
+        # Extract data from swc_dict
+        ids = swc_dict["id"]
+        voxels = np.array(swc_dict["voxel"], dtype=np.int32)
+
+        # Build graph
+        graph = nx.Graph()
+        id_lookup = dict()
+        run_length = 0
+        for i in range(len(swc_dict["id"])):
+            id_lookup[ids[i]] = i
+            if swc_dict["pid"][i] != -1:
+                # Add edge
+                parent = id_lookup[swc_dict["pid"][i]]
+                graph.add_edge(i, parent)
+
+                # Update run length
+                xyz_i = voxels[i] * self.anisotropy
+                xyz_p = voxels[parent] * self.anisotropy
+                run_length += distance.euclidean(xyz_i, xyz_p)
+
+        # Set graph-level attributes
+        graph.graph["n_edges"] = graph.number_of_edges()
+        graph.graph["run_length"] = run_length
+        graph.graph["voxel"] = voxels
+        return {swc_dict["swc_id"]: graph}
+
+    # --- Label Graphs ---
+    def _label_graphs_with_segmentation(self):
+        pass
+
+    def _label_graph(self, key):
+        pass
 
 # --- Update graph ---
 def delete_nodes(graph, target_label):
@@ -206,7 +255,11 @@ def compute_run_length(graph):
     return path_length
 
 
-# -- miscellaneous --
+# -- Miscellaneous --
+def get_segment_id(swc_id):
+    return int(swc_id.split(".")[0])
+
+    
 def sample_leaf(graph):
     """
     Samples leaf node from "graph".