refactor: updated saved fragments

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_graph.py

@@ -39,8 +39,18 @@ class SkeletonGraph(nx.Graph):
         A 3D array that contains a voxel coordinate for each node.
 
     """
-
-    def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
+    colors = [
+        "# COLOR 1.0 0.0 1.0",  # pink
+        "# COLOR 0.0 1.0 1.0",  # cyan
+        "# COLOR 1.0 1.0 0.0",  # yellow
+        "# COLOR 0.0 0.5 1.0",  # blue
+        "# COLOR 1.0 0.5 0.0",  # orange
+        "# COLOR 0.5 0.0 1.0",  # purple
+        "# COLOR 0.0 0.8 0.8",  # teal
+        "# COLOR 0.6 0.0 0.6",  # plum
+    ]
+
+    def __init__(self, anisotropy=(1.0, 1.0, 1.0), is_groundtruth=False):
         """
         Initializes a SkeletonGraph, including setting the anisotropy and
         initializing the run length attributes.
@@ -50,6 +60,9 @@ def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
         anisotropy : ArrayLike, optional
             Image to physical coordinates scaling factors to account for the
             anisotropy of the microscope. The default is (1.0, 1.0, 1.0).
+        is_groundtruth : bool, optional
+            Indication of whether this graph corresponds to a ground truth
+            tracing. The default is False.
 
         Returns
         -------
@@ -62,6 +75,7 @@ def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
         # Instance attributes
         self.anisotropy = np.array(anisotropy)
         self.filename = None
+        self.is_groundtruth = is_groundtruth
         self.labels = None
         self.run_length = 0
         self.voxels = None
@@ -317,7 +331,7 @@ def upd_labels(self, nodes, label):
         for i in nodes:
             self.labels[i] = label
 
-    def to_zipped_swc(self, zip_writer, color=None):
+    def to_zipped_swc(self, zip_writer):
         """
         Writes the graph to an SWC file format, which is then stored in a ZIP
         archive.
@@ -337,12 +351,12 @@ def to_zipped_swc(self, zip_writer, color=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")
+            text_buffer.write(self.get_color())
+            text_buffer.write("\n" + "# id, type, z, y, x, r, pid")
 
             # Write entries
             node_to_idx = dict()
-            r = 6 if color else 3
+            r = 2 if self.is_groundtruth else 3
             for i, j in nx.dfs_edges(self):
                 # Special Case: Root
                 x, y, z = tuple(self.voxels[i] * self.anisotropy)
@@ -359,3 +373,10 @@ def to_zipped_swc(self, zip_writer, color=None):
 
             # Finish
             zip_writer.writestr(self.filename, text_buffer.getvalue())
+
+    def get_color(self):
+        if self.is_groundtruth:
+             return "# COLOR 1.0 1.0 1.0"
+        else:
+            return util.sample_once(SkeletonGraph.colors)
+            

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -174,6 +174,7 @@ def load_groundtruth(self, swc_pointer):
         print("\n(1) Load Ground Truth")
         graph_builder = gutil.GraphBuilder(
             anisotropy=self.anisotropy,
+            is_groundtruth=True,
             label_mask=self.label_mask,
             use_anisotropy=False,
         )
@@ -203,6 +204,7 @@ def load_fragments(self, swc_pointer):
         if swc_pointer:
             graph_builder = gutil.GraphBuilder(
                 anisotropy=self.anisotropy,
+                is_groundtruth=False,
                 selected_ids=self.get_all_node_labels(),
                 use_anisotropy=self.use_anisotropy,
             )
@@ -464,12 +466,13 @@ def detect_splits(self):
                 n_missing = n_before - n_after
                 p_omit = 100 * (n_missing + n_split_edges) / n_before
                 p_split = 100 * n_split_edges / n_before
+                gt_rl = graph.run_length
 
                 self.graphs[key] = graph
-                self.metrics.at[key, "% Omit"] = p_omit
+                self.metrics.at[key, "% Omit"] = round(p_omit, 2)
                 self.metrics.at[key, "# Splits"] = gutil.count_splits(graph)
-                self.metrics.loc[key, "% Split"] = p_split
-                self.metrics.loc[key, "GT Run Length"] = graph.run_length
+                self.metrics.loc[key, "% Split"] = round(p_split, 2)
+                self.metrics.loc[key, "GT Run Length"] = round(gt_rl, 2)
                 pbar.update(1)
 
     # -- Merge Detection --
@@ -571,7 +574,7 @@ def is_fragment_merge(self, key, label, kdtree):
                 for leaf in gutil.get_leafs(fragment_graph):
                     voxel = fragment_graph.voxels[leaf]
                     gt_voxel = util.kdtree_query(kdtree, voxel)
-                    if self.physical_dist(gt_voxel, voxel) > 50:
+                    if self.physical_dist(gt_voxel, voxel) > 60:
                         visited = self.find_merge_site(
                             key, kdtree, fragment_graph, leaf, visited
                         )
@@ -632,10 +635,13 @@ def process_merge_sites(self):
 
             # Save merge sites
             if self.save_merges:
+                row_names = list()
                 for i in range(len(self.merge_sites)):
                     filename = f"merge-{i + 1}.swc"
                     xyz = self.merge_sites.iloc[i]["World"]
                     swc_util.to_zipped_point(self.merge_writer, filename, xyz)
+                    row_names.append(filename)
+                self.merge_sites.index = row_names
                 self.merge_writer.close()
 
             # Update counter
@@ -645,7 +651,7 @@ def process_merge_sites(self):
 
             # Save results
             path = os.path.join(self.output_dir, "merge_sites.csv")
-            self.merge_sites.to_csv(path, index=False)
+            self.merge_sites.to_csv(path, index=True)
 
 
     def adjust_metrics(self, key):
@@ -757,7 +763,7 @@ def quantify_merges(self):
         """
         for key in self.graphs:
             p = self.n_merged_edges[key] / self.graphs[key].graph["n_edges"]
-            self.metrics.loc[key, "% Merged"] = 100 * p
+            self.metrics.loc[key, "% Merged"] = round(100 * p, 2)
 
     # -- Compute Metrics --
     def compute_edge_accuracy(self):
@@ -776,7 +782,8 @@ def compute_edge_accuracy(self):
         for key in self.graphs:
             p_omit = self.metrics.loc[key, "% Omit"]
             p_merged = self.metrics.loc[key, "% Merged"]
-            self.metrics.loc[key, "Edge Accuracy"] = 100 - p_omit - p_merged
+            edge_accuracy = round(100 - p_omit - p_merged, 2)
+            self.metrics.loc[key, "Edge Accuracy"] = edge_accuracy
 
     def compute_erl(self):
         """
@@ -799,8 +806,9 @@ def compute_erl(self):
             wgt = run_lengths / max(np.sum(run_lengths), 1)
 
             erl = np.sum(wgt * run_lengths)
-            self.metrics.loc[key, "ERL"] = erl
-            self.metrics.loc[key, "Normalized ERL"] = erl / max(run_length, 1)
+            n_erl = round(erl / max(run_length, 1), 4)
+            self.metrics.loc[key, "ERL"] = round(erl, 2)
+            self.metrics.loc[key, "Normalized ERL"] = n_erl
 
     def compute_weighted_avg(self, column_name):
         wgt = self.metrics["GT Run Length"]

src/segmentation_skeleton_metrics/utils/graph_util.py

@@ -27,6 +27,7 @@ class GraphBuilder:
     def __init__(
         self,
         anisotropy=(1.0, 1.0, 1.0),
+        is_groundtruth=False,
         label_mask=None,
         selected_ids=None,
         use_anisotropy=True,
@@ -39,6 +40,9 @@ def __init__(
         anisotropy : Tuple[int], optional
             Image to physical coordinates scaling factors to account for the
             anisotropy of the microscope. The default is [1.0, 1.0, 1.0].
+        is_groundtruth : bool, optional
+            Indication of whether this graph corresponds to a ground truth
+            tracing. The default is False.
         label_mask : ImageReader, optional
             Predicted segmentation mask.
         selected_ids : Set[int], optional
@@ -55,6 +59,7 @@ def __init__(
         """
         # Instance attributes
         self.anisotropy = anisotropy
+        self.is_groundtruth = is_groundtruth
         self.label_mask = label_mask
 
         # Reader
@@ -146,7 +151,9 @@ def to_graph(self, swc_dict):
 
         """
         # Initialize graph
-        graph = SkeletonGraph(anisotropy=self.anisotropy)
+        graph = SkeletonGraph(
+            anisotropy=self.anisotropy, is_groundtruth=self.is_groundtruth
+        )
         graph.init_voxels(swc_dict["voxel"])
         graph.set_filename(swc_dict["swc_id"] + ".swc")
         graph.set_nodes(len(swc_dict["id"]))

src/segmentation_skeleton_metrics/utils/swc_util.py

@@ -528,7 +528,7 @@ def to_zipped_point(zip_writer, filename, xyz):
     with StringIO() as text_buffer:
         # Preamble
         text_buffer.write("# COLOR 1.0 0.0 0.0")
-        text_buffer.write("# id, type, z, y, x, r, pid")
+        text_buffer.write("\n" + "# id, type, z, y, x, r, pid")
 
         # Write entry
         x, y, z = tuple(xyz)