bug: save fragments

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -66,7 +66,7 @@ def __init__(
         localize_merge=False,
         preexisting_merges=None,
         save_merges=False,
-        save_projections=False,
+        save_fragments=False,
         valid_labels=None,
     ):
         """
@@ -103,7 +103,7 @@ def __init__(
         save_merges: bool, optional
             Indication of whether to save fragments with a merge mistake. The
             default is None.
-        save_projections : bool, optional
+        save_fragments : bool, optional
             Indication of whether to save fragments that project onto each
             ground truth skeleton. The default is False.
         valid_labels : set[int], optional
@@ -123,7 +123,7 @@ def __init__(
         self.output_dir = output_dir
         self.preexisting_merges = preexisting_merges
         self.save_merges = save_merges
-        self.save_projections = save_projections
+        self.save_fragments = save_fragments
 
         # Label handler
         self.label_handler = gutil.LabelHandler(
@@ -135,14 +135,8 @@ def __init__(
         self.load_groundtruth(gt_pointer)
         self.load_fragments(fragments_pointer)
 
-        # Initialize writer
-        if self.save_merges:
-            self.init_zip_writer()
-
-        # Initialize fragment projections directory
-        if self.save_projections:
-            self.projections_dir = os.path.join(output_dir, "projections")
-            util.mkdir(self.projections_dir)
+        # Initialize writers
+        self.init_zip_writers()
 
     # --- Load Data ---
     def load_groundtruth(self, swc_pointer):
@@ -346,9 +340,9 @@ def get_node_labels(self, key, inverse_bool=False):
         else:
             return self.graphs[key].get_labels()
 
-    def init_zip_writer(self):
+    def init_zip_writers(self):
         """
-        Initializes "self.zip_writer" attribute by setting up a directory for
+        Initializes "self.merge_writer" attribute by setting up a directory for
         output files and creating ZIP files for each graph in "self.graphs".
 
         Parameters
@@ -360,16 +354,31 @@ def init_zip_writer(self):
         None
 
         """
-        # Initialize output directory
-        merged_fragments_dir = os.path.join(self.output_dir, "merged_fragments")
-        util.mkdir(merged_fragments_dir)
-
-        # Save intial graphs
-        self.zip_writer = dict()
-        for key in self.graphs.keys():
-            zip_path = f"{merged_fragments_dir}/{key}.zip"
-            self.zip_writer[key] = ZipFile(zip_path, "w")
-            self.graphs[key].to_zipped_swc(self.zip_writer[key])
+        # Merged fragments zip writer
+        if self.save_merges:
+            # Initialize directory
+            merges_dir = os.path.join(self.output_dir, "merged_fragments")
+            util.mkdir(merged_fragments_dir)
+
+            # Initialize zip writer
+            self.merge_writer = dict()
+            for key in self.graphs.keys():
+                zip_path = f"{merged_fragments_dir}/{key}.zip"
+                self.merge_writer[key] = ZipFile(zip_path, "w")
+                self.graphs[key].to_zipped_swc(self.merge_writer[key])
+
+        # Fragments zip writer
+        if self.save_fragments:
+            # Initialize direction
+            fragments_dir = os.path.join(self.output_dir, "fragments")
+            util.mkdir(fragments_dir)
+
+            # Initialize zip writer
+            self.fragment_writer = dict()
+            for key in self.graphs.keys():
+                zip_path = f"{fragments_dir}/{key}.zip"
+                self.fragment_writer[key] = ZipFile(zip_path, "w")
+                self.graphs[key].to_zipped_swc(self.fragment_writer[key])
 
     # -- Main Routine --
     def run(self):
@@ -539,22 +548,13 @@ def count_merges(self, key, kdtree):
         None
 
         """
-        # Initialize zip writer
-        if self.save_projections:
-            zip_path = os.path.join(self.projections_dir, key + ".zip")
-            zip_writer = ZipFile(zip_path, "w")
-            #self.graphs[key].to_zipped_swc(zip_writer)
-
         # Iterate over fragments that intersect with GT skeleton
         for label in self.get_node_labels(key):
             nodes = self.graphs[key].nodes_with_label(label)
             if len(nodes) > 40:
                 for label in self.label_handler.get_class(label):
                     if label in self.fragment_ids:
                         self.is_fragment_merge(key, label, kdtree)
-                        if self.save_projections:
-                            fragment_graph = self.find_graph_from_label(label)[0]
-                            fragment_graph.to_zipped_swc(zip_writer)
 
     def is_fragment_merge(self, key, label, kdtree):
         """
@@ -580,6 +580,7 @@ def is_fragment_merge(self, key, label, kdtree):
         """
         # Search graphs
         for fragment_graph in self.find_graph_from_label(label):
+            # Search for merge
             max_dist = 0
             min_dist = np.inf
             for voxel in fragment_graph.voxels:
@@ -601,11 +602,15 @@ def is_fragment_merge(self, key, label, kdtree):
 
                     # Save merged fragment (if applicable)
                     if self.save_merges:
-                        fragment_graph.to_zipped_swc(self.zip_writer[key])
+                        fragment_graph.to_zipped_swc(self.merge_writer[key])
                     if self.localize_merge:
                         self.find_merge_site(key, fragment_graph, kdtree) 
                     break
 
+            # Save fragment (if applicable) 
+            if self.save_fragments and min_dist < 3: 
+                fragment_graph.to_zipped_swc(self.fragment_writer[key])
+
     def adjust_metrics(self, key):
         """
         Adjusts the metrics of the graph associated with the given key by

src/segmentation_skeleton_metrics/split_detection.py

@@ -46,15 +46,15 @@ def run(process_id, graph):
         label_i = int(graph.labels[i])
         label_j = int(graph.labels[j])
         if is_split(label_i, label_j):
-            graph.remove_edge(i, j) temp
+            graph.remove_edge(i, j)
             split_cnt += 1
         elif label_j == 0:
             check_misalignment(graph, visited_edges, i, j)
         visited_edges.add(frozenset({i, j}))
 
     # Finish
     split_percent = split_cnt / graph.graph["n_edges"]
-    graph.remove_nodes_with_label(0) temp
+    graph.remove_nodes_with_label(0)
     return process_id, graph, split_percent