bug: evaluate corrected segmentation

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_graph.py

@@ -9,6 +9,7 @@
 
 """
 
+from io import StringIO
 from scipy.spatial import distance
 
 import networkx as nx
@@ -59,7 +60,10 @@ def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
 
         # Instance attributes
         self.anisotropy = np.array(anisotropy)
+        self.filename = None
+        self.labels = None
         self.run_length = 0
+        self.voxels = None
 
     def init_labels(self):
         """
@@ -91,6 +95,23 @@ def init_voxels(self, voxels):
         """
         self.voxels = np.array(voxels, dtype=np.int32)
 
+    def set_filename(self, filename):
+        """
+        Sets the filename attribute which corresponds to the SWC file that the
+        graph is built from.
+
+        Parameters
+        ----------
+        filename : str
+            Name of SWC file that graph is built from.
+
+        Returns
+        -------
+        None
+
+        """
+        self.filename = filename
+
     def set_nodes(self):
         """
         Adds nodes to the graph. The nodes are assigned indices from 0 to the
@@ -111,7 +132,7 @@ def set_nodes(self):
     # --- Getters ---
     def get_labels(self):
         """
-        Gets the unique label values in the "labels" attribute.
+        Gets the unique non-zero label values in the "labels" attribute.
 
         Parameters
         ----------
@@ -120,10 +141,13 @@ def get_labels(self):
         Returns
         -------
         numpy.ndarray
-            A 1D array of unique labels assigned to nodes in the graph.
+            A 1D array of unique non-zero labels assigned to nodes in the
+            graph.
 
         """
-        return np.unique(self.labels)
+        labels = set(np.unique(self.labels))
+        labels.discard(0)
+        return labels
 
     def nodes_with_label(self, label):
         """
@@ -289,3 +313,47 @@ def upd_labels(self, nodes, label):
         """
         for i in nodes:
             self.labels[i] = label
+
+    def to_zipped_swc(self, zip_writer, color=None):
+        """
+        Writes a graph to an SWC file that is to be stored in a zip.
+
+        Parameters
+        ----------
+        zip_writer : zipfile.ZipFile
+            ...
+        color : str, optional
+            ...
+
+        Returns
+        -------
+        None
+
+        """
+        with StringIO() as text_buffer:
+            # Preamble
+            text_buffer.write("# COLOR " + color) if color else None
+            text_buffer.write("# id, type, z, y, x, r, pid")
+
+            # Write entries
+            n_entries = 0
+            node_to_idx = dict()
+            r = 5 if color else 3
+            for i, j in nx.dfs_edges(self):
+                # Special Case: Root
+                x, y, z = tuple(self.voxels[i] * self.anisotropy)
+                if len(node_to_idx) == 0:
+                    parent = -1
+                    node_to_idx[i] = 1
+                    text_buffer.write("\n" + f"1 2 {x} {y} {z} {r} {parent}")
+                    n_entries += 1
+
+                # General Case
+                node = n_entries + 1
+                parent = node_to_idx[i]
+                node_to_idx[j] = n_entries + 1
+                text_buffer.write("\n" + f"{node} 2 {x} {y} {z} {r} {parent}")
+                n_entries += 1
+
+            # Finish
+            zip_writer.writestr(self.filename, text_buffer.getvalue())

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -23,12 +23,9 @@
 from segmentation_skeleton_metrics.utils import (
     graph_util as gutil,
     img_util,
-    swc_util,
     util
 )
 
-MIN_CNT = 40
-
 
 class SkeletonMetric:
     """
@@ -155,10 +152,8 @@ def load_groundtruth(self, swc_pointer):
     def load_fragments(self, swc_pointer):
         print("\n(2) Load Fragments")
         if swc_pointer:
-            coords_only = False  #not self.save_projections
             graph_builder = gutil.GraphBuilder(
                 anisotropy=self.anisotropy,
-                coords_only=coords_only,
                 selected_ids=self.get_all_node_labels(),
                 use_anisotropy=True,
             )
@@ -172,7 +167,7 @@ def set_fragment_ids(self):
         for key in self.fragment_graphs:
             self.fragment_ids.add(util.get_segment_id(key))
 
-    def label_graphs(self, key, batch_size=64):
+    def label_graphs(self, key, batch_size=128):
         """
         Iterates over nodes in "graph" and stores the corresponding label from
         predicted segmentation mask (i.e. "self.label_mask") as a node-level
@@ -201,7 +196,7 @@ def label_graphs(self, key, batch_size=64):
                     visited.add(i)
 
                 # Check whether to submit batch
-                is_node_far = self.graphs[key].dist(root, j) > 128
+                is_node_far = self.graphs[key].dist(root, j) > batch_size
                 is_batch_full = len(batch) >= batch_size
                 if is_node_far or is_batch_full:
                     threads.append(
@@ -306,9 +301,7 @@ def init_zip_writer(self):
         self.zip_writer = dict()
         for key in self.graphs.keys():
             self.zip_writer[key] = ZipFile(f"{output_dir}/{key}.zip", "w")
-            swc_util.to_zipped_swc(
-                self.zip_writer[key], self.graphs[key],
-            )
+            self.graphs[key].to_zipped_swc(self.zip_writer[key])
 
     # -- Main Routine --
     def run(self):
@@ -331,11 +324,6 @@ def run(self):
         self.detect_splits()
         self.quantify_splits()
 
-        # Check for prexisting merges
-        if self.preexisting_merges:
-            for key in self.graphs:
-                self.adjust_metrics(key)
-
         # Merge evaluation
         self.detect_merges()
         self.quantify_merges()
@@ -344,39 +332,6 @@ def run(self):
         full_results, avg_results = self.compile_results()
         return full_results, avg_results
 
-    def adjust_metrics(self, key):
-        """
-        Adjusts the metrics of the graph associated with the given key by
-        removing nodes corresponding to known merges and their corresponding
-        subgraphs. Updates the total number of edges and run lengths in the
-        graph.
-
-        Parameters
-        ----------
-        key : str
-            Identifier for the graph to adjust.
-
-        Returns
-        -------
-        None
-
-        """
-        for label in self.preexisting_merges:
-            label = self.label_map[label] if self.label_map else label
-            if label in self.graphs[key].get_labels():
-                # Extract subgraph
-                nodes = self.graphs[key].nodes_with_label(label)
-                subgraph = self.graphs[key].subgraph(nodes)
-
-                # Adjust metrics
-                n_edges = subgraph.number_of_edges()
-                rls = gutil.compute_run_lengths(subgraph)
-                self.graphs[key].graph["run_length"] -= np.sum(rls)
-                self.graphs[key].graph["n_edges"] -= n_edges
-
-                # Update graph
-                self.graphs[key].remove_nodes_from(nodes)
-
     # -- Split Detection --
     def detect_splits(self):
         """
@@ -393,7 +348,7 @@ def detect_splits(self):
 
         """
         pbar = tqdm(total=len(self.graphs), desc="Split Detection")
-        with ProcessPoolExecutor() as executor:
+        with ProcessPoolExecutor(max_workers=8) as executor:
             # Assign processes
             processes = list()
             for key, graph in self.graphs.items():
@@ -470,7 +425,12 @@ def detect_merges(self):
                     self.count_merges(key, kdtree)
                 pbar.update(1)
 
-        # Process merges
+        # Adjust metrics (if applicable)
+        if self.preexisting_merges:
+            for key in self.graphs:
+                self.adjust_metrics(key)
+
+        # Find graphs with common node labels
         for (key_1, key_2), label in self.find_label_intersections():
             self.process_merge(key_1, label, -1)
             self.process_merge(key_2, label, -1)
@@ -502,7 +462,7 @@ def count_merges(self, key, kdtree):
         """
         for label in self.get_node_labels(key):
             nodes = self.graphs[key].nodes_with_label(label)
-            if len(nodes) > MIN_CNT:
+            if len(nodes) > 50:
                 for label in self.label_handler.get_class(label):
                     if label in self.fragment_ids:
                         self.is_fragment_merge(key, label, kdtree)
@@ -539,16 +499,45 @@ def is_fragment_merge(self, key, label, kdtree):
                 self.merged_labels.add((key, equiv_label, tuple(xyz)))
 
                 # Save merged fragment (if applicable)
-                if self.save_projections and label in self.fragment_graphs:
-                    swc_util.to_zipped_swc(
-                        self.zip_writer[key], self.fragment_graphs[label]
-                    )
+                if self.save_projections:
+                    fragment_graph.to_zipped_swc(self.zip_writer[key])
                 break
 
-    def find_graph_from_label(self, label):
-        for key in self.fragment_graphs:
-            if label == util.get_segment_id(key):
-                return self.fragment_graphs[key]
+    def adjust_metrics(self, key):
+        """
+        Adjusts the metrics of the graph associated with the given key by
+        removing nodes corresponding to known merges and their corresponding
+        subgraphs. Updates the total number of edges and run lengths in the
+        graph.
+
+        Parameters
+        ----------
+        key : str
+            Identifier for the graph to adjust.
+
+        Returns
+        -------
+        None
+
+        """
+        visited = set()
+        for label in self.preexisting_merges:
+            label = self.label_handler.mapping[label]
+            if label in self.graphs[key].get_labels():
+                if label not in visited and label != 0:
+                    # Get component with label
+                    nodes = self.graphs[key].nodes_with_label(label)
+                    root = util.sample_once(list(nodes))
+
+                    # Adjust metrics
+                    rl = self.graphs[key].run_length_from(root)
+                    self.graphs[key].run_length -= np.sum(rl)
+                    self.graphs[key].graph["n_edges"] -= len(nodes) - 1
+
+                    # Update graph
+                    self.graphs[key].remove_nodes_from(nodes)
+                    visited.add(label)
+                    print("# nodes deleted:", len(nodes))
 
     def find_label_intersections(self):
         """
@@ -673,7 +662,7 @@ def get_merged_label(self, label):
         for l in self.label_handler.get_class(label):
             if l in self.fragment_graphs.keys():
                 return l
-        return self.inverse_label_map[label]
+        return self.label_handler.inverse_mapping[label]
 
     # -- Compute Metrics --
     def compile_results(self):
@@ -866,7 +855,13 @@ def list_metrics(self):
         ]
         return metrics
 
-    # -- util --
+    # -- Helpers --
+    def find_graph_from_label(self, label):
+        for key in self.fragment_graphs:
+            if label == util.get_segment_id(key):
+                return self.fragment_graphs[key]
+        return None
+
     def physical_dist(self, voxel_1, voxel_2):
         xyz_1 = img_util.to_physical(voxel_1, self.anisotropy)
         xyz_2 = img_util.to_physical(voxel_2, self.anisotropy)
@@ -896,40 +891,6 @@ def to_local_voxels(self, key, i, offset):
 
 
 # -- util --
-def find_sites(graphs, get_labels):
-    """
-    Detects merges between ground truth graphs which are considered to be
-    potential merge sites.
-
-    Parameters
-    ----------
-    graphs : dict
-        Dictionary where the keys are graph ids and values are graphs.
-    get_labels : func
-        Gets the label of a node in "graphs".
-
-    Returns
-    -------
-    merge_ids : set[tuple]
-        Set of tuples containing a tuple of graph ids and common label between
-        the graphs.
-
-    """
-    merge_ids = set()
-    visited = set()
-    for key_1 in graphs:
-        for key_2 in graphs:
-            keys = frozenset((key_1, key_2))
-            if key_1 != key_2 and keys not in visited:
-                visited.add(keys)
-                intersection = get_labels(key_1).intersection(
-                    get_labels(key_2)
-                )
-                for label in intersection:
-                    merge_ids.add((keys, label))
-    return merge_ids
-
-
 def generate_result(keys, stats):
     """
     Reorders items in "stats" with respect to the order defined by "keys".