refactor: improved new features (#102)

Co-authored-by: anna-grim <[email protected]>

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -27,6 +27,7 @@
 from segmentation_skeleton_metrics.utils import (
     graph_util as gutil,
     img_util,
+    swc_util,
     util,
 )
 
@@ -63,7 +64,7 @@ def __init__(
         anisotropy=(1.0, 1.0, 1.0),
         connections_path=None,
         fragments_pointer=None,
-        localize_merge=False,
+        localize_merges=False,
         preexisting_merges=None,
         save_merges=False,
         save_fragments=False,
@@ -94,7 +95,7 @@ def __init__(
             "swc_util.Reader" for documentation. Notes: (1) "anisotropy" is
             applied to these SWC files and (2) these SWC files are required
             for counting merges. The default is None.
-        localize_merge : bool, optional
+        localize_merges : bool, optional
             Indication of whether to search for the approximate location of a
             merge. The default is False.
         preexisting_merges : List[int], optional
@@ -119,7 +120,7 @@ def __init__(
         # Instance attributes
         self.anisotropy = anisotropy
         self.connections_path = connections_path
-        self.localize_merge = localize_merge
+        self.localize_merges = localize_merges
         self.output_dir = output_dir
         self.preexisting_merges = preexisting_merges
         self.save_merges = save_merges
@@ -136,12 +137,7 @@ def __init__(
         self.load_fragments(fragments_pointer)
 
         # Initialize writers
-        self.init_zip_writers()
-
-        # Initialize fragment projections directory
-        if self.save_projections:
-            self.projections_dir = os.path.join(output_dir, "projections")
-            util.mkdir(self.projections_dir)
+        self.init_writers()
 
     # --- Load Data ---
     def load_groundtruth(self, swc_pointer):
@@ -345,7 +341,7 @@ def get_node_labels(self, key, inverse_bool=False):
         else:
             return self.graphs[key].get_labels()
 
-    def init_zip_writers(self):
+    def init_writers(self):
         """
         Initializes "self.merge_writer" attribute by setting up a directory for
         output files and creating ZIP files for each graph in "self.graphs".
@@ -359,31 +355,45 @@ def init_zip_writers(self):
         None
 
         """
-        # Merged fragments zip writer
+        # Fragments writer
+        if self.save_fragments:
+            # Initialize direction
+            fragments_dir = os.path.join(self.output_dir, "fragments")
+            util.mkdir(fragments_dir, delete=True)
+
+            # 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])
+
+        # Merged fragments writer
         if self.save_merges:
             # Initialize directory
             merges_dir = os.path.join(self.output_dir, "merged_fragments")
-            util.mkdir(merged_fragments_dir)
+            util.mkdir(merged_fragments_dir, delete=True)
 
-            # Initialize zip writer
+            # 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)
+        # Merge sites
+        if self.localize_merges:
+            # Initialize directory
+            merges_dir = os.path.join(self.output_dir, "merge-sites")
+            util.mkdir(merges_dir, delete=True)
 
-            # 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])
+            # ZIP writer
+            zip_path = f"{merges_dir}/estimated-merge-sites.zip"
+            self.site_zip_writer = ZipFile(zip_path, "w")
+
+            # Txt writer
+            sites_path = os.path.join(merges_dir, "estimated-merge-sites.txt")
+            self.site_txt_writer = open(sites_path, "w", encoding="utf-8")
 
     # -- Main Routine --
     def run(self):
@@ -560,9 +570,6 @@ def count_merges(self, key, kdtree):
                 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):
         """
@@ -611,12 +618,12 @@ def is_fragment_merge(self, key, label, kdtree):
                     # Save merged fragment (if applicable)
                     if self.save_merges:
                         fragment_graph.to_zipped_swc(self.merge_writer[key])
-                    if self.localize_merge:
+                    if self.localize_merges:
                         self.find_merge_site(key, fragment_graph, kdtree) 
                     break
 
             # Save fragment (if applicable) 
-            if self.save_fragments and min_dist < 3: 
+            if self.save_fragments and min_dist < 3:
                 fragment_graph.to_zipped_swc(self.fragment_writer[key])
 
     def adjust_metrics(self, key):
@@ -729,7 +736,13 @@ def find_merge_site(self, key, fragment_graph, kdtree):
                     gt_voxel = util.kdtree_query(kdtree, voxel_j)
                     if self.physical_dist(gt_voxel, voxel_j) < 2:
                         hit = True
-                        print("Approximate Site:", img_util.to_physical(voxel_j, self.anisotropy))
+                        merge_cnt = np.sum(list(self.merge_cnt.values()))
+                        filename = f"{merge_cnt}.swc"
+                        xyz = img_util.to_physical(voxel_j, self.anisotropy)
+                        swc_util.to_zipped_point(
+                            self.site_zip_writer, filename, xyz
+                        )
+                        self.site_txt_writer.write(f"{tuple(xyz)}\n")
                         break
 
             # Check whether to continue

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
 
 

src/segmentation_skeleton_metrics/utils/swc_util.py

@@ -28,6 +28,7 @@
     ProcessPoolExecutor,
     ThreadPoolExecutor,
 )
+from io import StringIO
 from tqdm import tqdm
 from zipfile import ZipFile
 
@@ -380,3 +381,36 @@ def read_voxel(self, xyz_str, offset):
         """
         xyz = [float(xyz_str[i]) + offset[i] for i in range(3)]
         return img_util.to_voxels(xyz, self.anisotropy)
+
+
+# --- Write ---
+def to_zipped_point(zip_writer, filename, xyz):
+    """
+    Writes a point to an SWC file format, which is then stored in a ZIP
+    archive.
+
+    Parameters
+    ----------
+    zip_writer : zipfile.ZipFile
+        A ZipFile object that will store the generated SWC file.
+    filename : str
+        Filename of SWC file.
+    xyz : ArrayLike
+        Point to be written to SWC file.
+
+    Returns
+    -------
+    None
+
+    """
+    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")
+
+        # Write entry
+        x, y, z = tuple(xyz)
+        text_buffer.write("\n" + f"1 2 {x} {y} {z} 15 -1")
+
+        # Finish
+        zip_writer.writestr(filename, text_buffer.getvalue())