refactor: skeleton graph label attribute

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_graph.py

@@ -0,0 +1,165 @@
+from scipy.spatial import distance
+
+import networkx as nx
+import numpy as np
+
+from segmentation_skeleton_metrics.utils import util
+
+
+class SkeletonGraph(nx.Graph):
+
+    def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
+        # Call parent class
+        super(SkeletonGraph, self).__init__()
+
+        # Instance attributes
+        self.anisotropy = anisotropy
+        self.run_length = 0
+
+    def set_labels(self):
+        self.labels = np.zeros((self.number_of_nodes()), dtype=int)
+
+    def set_nodes(self):
+        num_nodes = len(self.voxels)
+        self.add_nodes_from(np.arange(num_nodes))
+
+    def set_voxels(self, voxels):
+        self.voxels = np.array(voxels, dtype=np.int32)
+
+    # --- Getters ---
+    def get_labels(self):
+        return np.unique(self.labels)
+
+    def nodes_with_label(self, label):
+        return np.where(self.labels == label)[0]
+
+    # --- Computation ---
+    def dist(self, i, j):
+        """
+        Computes the Euclidean distance between the voxel coordinates
+        cooresponding to the given nodes.
+
+        Parameters
+        ----------
+        i : int
+            Node ID.
+        j : int
+            Node ID.
+
+        Returns
+        -------
+        float
+            Distance between voxel coordinates of the given nodes.
+
+        """
+        return distance.euclidean(self.voxels[i], self.voxels[j])
+
+    def physical_dist(self, i, j):
+        """
+        Computes the Euclidean distance between the physical coordinates
+        cooresponding to the given nodes.
+
+        Parameters
+        ----------
+        i : int
+            Node ID.
+        j : int
+            Node ID.
+
+        Returns
+        -------
+        float
+            Distance between physical coordinates of the given nodes.
+
+        """
+        xyz_i = self.voxels[i] * self.anisotropy
+        xyz_j = self.voxels[j] * self.anisotropy
+        return distance.euclidean(xyz_i, xyz_j)
+
+    def get_bbox(self, nodes):
+        bbox_min = np.inf * np.ones(3)
+        bbox_max = np.zeros(3)
+        for i in nodes:
+            bbox_min = np.minimum(bbox_min, self.voxels[i])
+            bbox_max = np.maximum(bbox_max, self.voxels[i] + 1)
+        return {"min": bbox_min.astype(int), "max": bbox_max.astype(int)}
+
+    def remove_nodes_with_label(self, label):
+        """
+        Removes nodes with the given label
+
+        Parameters
+        ----------
+        label : int
+            Label to be deleted from graph.
+
+        Returns
+        -------
+        None
+
+        """
+        nodes = self.nodes_with_label(label)
+        self.remove_nodes_from(nodes)
+
+    def run_lengths(self):
+        """
+        Computes the path length of each connected component.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        numpy.ndarray
+            Array containing run lengths of each connected component.
+
+        """
+        run_lengths = []
+        if self.number_of_nodes() > 0:
+            for nodes in nx.connected_components(self):
+                root = util.sample_once(nodes)
+                run_lengths.append(self.run_length_from(root))
+        else:
+            run_lengths.append(0)
+        return np.array(run_lengths)
+
+    def run_length_from(self, root):
+        """
+        Computes the path length of the connected component that contains
+        "root".
+
+        Parameters
+        ----------
+        graph : networkx.Graph
+            Graph to be parsed.
+
+        Returns
+        -------
+        float
+            Path length.
+
+        """
+        run_length = 0
+        for i, j in nx.dfs_edges(self, source=root):
+            run_length += self.physical_dist(i, j)
+        return run_length
+
+    def upd_labels(self, nodes, label):
+        """
+        Updates the label of each node in "nodes" with "label".
+
+        Parameters
+        ----------
+        nodes : List[int]
+            Nodes to be updated.
+        label : int
+            Updated label.
+
+        Returns
+        -------
+        None
+
+        """
+        for i in nodes:
+            self.labels[i] = label

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -221,12 +221,11 @@ def label_graphs(self, key, batch_size=64):
             threads.append(executor.submit(self.get_patch_labels, key, batch))
 
             # Process results
-            n_nodes = self.graphs[key].number_of_nodes()
-            self.graphs[key].graph["label"] = np.zeros((n_nodes), dtype=int)
+            self.graphs[key].set_labels()
             for thread in as_completed(threads):
                 node_to_label = thread.result()
                 for i, label in node_to_label.items():
-                    self.graphs[key].graph["label"][i] = label
+                    self.graphs[key].labels[i] = label
 
     def get_patch_labels(self, key, nodes):
         bbox = self.graphs[key].get_bbox(nodes)
@@ -411,7 +410,7 @@ def detect_splits(self):
             self.split_percent = dict()
             for process in as_completed(processes):
                 key, graph, split_percent = process.result()
-                self.graphs[key] = gutil.remove_nodes(graph, 0)
+                self.graphs[key] = graph
                 self.split_percent[key] = split_percent
                 pbar.update(1)
 

src/segmentation_skeleton_metrics/split_detection.py

@@ -13,8 +13,6 @@
 
 import networkx as nx
 
-from segmentation_skeleton_metrics.utils import graph_util as gutil
-
 
 def run(process_id, graph):
     """
@@ -33,7 +31,7 @@ def run(process_id, graph):
     """
     # Initializations
     split_cnt = 0
-    source = gutil.sample_leaf(graph)
+    source = get_leaf(graph)
     dfs_edges = deque(list(nx.dfs_edges(graph, source=source)))
     visited_edges = set()
 
@@ -45,8 +43,8 @@ def run(process_id, graph):
             continue
 
         # Visit edge
-        label_i = int(graph.graph["label"][i])
-        label_j = int(graph.graph["label"][j])
+        label_i = int(graph.labels[i])
+        label_j = int(graph.labels[j])
         if is_split(label_i, label_j):
             graph.remove_edge(i, j)
             split_cnt += 1
@@ -56,6 +54,7 @@ def run(process_id, graph):
 
     # Finish
     split_percent = split_cnt / graph.graph["n_edges"]
+    graph.remove_nodes_with_label(0)
     return process_id, graph, split_percent
 
 
@@ -91,7 +90,7 @@ def check_misalignment(graph, visited_edges, nb, root):
     while len(queue) > 0:
         # Visit node
         j = queue.popleft()
-        label_j = int(graph.graph["label"][j])
+        label_j = int(graph.labels[j])
         if label_j != 0:
             label_collisions.add(label_j)
         visited.add(j)
@@ -107,7 +106,7 @@ def check_misalignment(graph, visited_edges, nb, root):
     # Upd zero nodes
     if len(label_collisions) == 1:
         label = label_collisions.pop()
-        upd_labels(graph, visited, label)
+        graph.upd_labels(visited, label)
 
 
 # -- Helpers --
@@ -131,24 +130,21 @@ def is_split(a, b):
     return (a > 0 and b > 0) and (a != b)
 
 
-def upd_labels(graph, nodes, label):
+def get_leaf(graph):
     """
-    Updates the label of each node in "nodes" with "label".
+    Gets a leaf node from "graph".
 
     Parameters
     ----------
     graph : networkx.Graph
-        Graph to be updated.
-    nodes : list
-        List of nodes to be updated.
-    label : int
-        New label of each node in "nodes".
+        Graph to be sampled from.
 
     Returns
     -------
-    networkx.Graph
-        Updated graph.
+    int
+        Leaf node of "graph"
 
     """
-    for i in nodes:
-        graph.graph["label"][i] = label
+    for i in graph.nodes:
+        if graph.degree[i] == 1:
+            return i