bugs: path parser, write (#150)

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

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_graph.py

@@ -8,6 +8,7 @@
 
 """
 
+from collections import defaultdict
 from io import StringIO
 from scipy.spatial import distance
 
@@ -292,27 +293,28 @@ def to_zipped_swc(self, zip_writer):
         zip_writer : zipfile.ZipFile
             A ZipFile object that will store the generated SWC file.
         """
+        # Subroutines
+        def write_entry(node, parent):
+            x, y, z = tuple(self.voxels[i] * self.anisotropy)
+            r = self.node_radius[node] if preserve_radius else 2
+            node_id = len(node_to_idx) + 1
+            node_to_idx[node] = node_id
+            text_buffer.write("\n" + f"{node_id} 2 {x} {y} {z} {r} {parent}")
+
+        # Main
         with StringIO() as text_buffer:
             # Preamble
-            text_buffer.write(self.get_color())
             text_buffer.write("\n" + "# id, type, z, y, x, r, pid")
 
             # Write entries
-            node_to_idx = dict()
-            r = 2 if self.is_groundtruth else 3
+            node_to_idx = defaultdict(lambda: -1)
             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}")
-
-                # General Case
-                node = len(node_to_idx) + 1
-                parent = node_to_idx[i]
-                node_to_idx[j] = node
-                text_buffer.write("\n" + f"{node} 2 {x} {y} {z} {r} {parent}")
+                    write_entry(i, node_to_idx[i])
+
+                # General Case: Non-Root
+                write_entry(j, node_to_idx[i])
 
             # Finish
             zip_writer.writestr(self.filename, text_buffer.getvalue())

src/segmentation_skeleton_metrics/utils/swc_util.py

@@ -293,7 +293,7 @@ def read_from_gcs(self, gcs_path):
             names and values from an SWC file.
         """
         # List filenames
-        bucket_name, prefix = parse_cloud_path(gcs_path)
+        bucket_name, prefix = util.parse_cloud_path(gcs_path)
         swc_paths = util.list_gcs_filenames(bucket_name, prefix, ".swc")
         zip_paths = util.list_gcs_filenames(bucket_name, prefix, ".zip")
 
@@ -430,7 +430,7 @@ def read_from_gcs_zip(self, bucket_name, path):
 
     def read_from_s3(self, s3_path):
         # List filenames
-        bucket_name, prefix = parse_cloud_path(s3_path)
+        bucket_name, prefix = util.parse_cloud_path(s3_path)
         swc_paths = util.list_s3_paths(bucket_name, prefix, extension=".swc")
 
         # Parse SWC files
@@ -548,37 +548,6 @@ def read_voxel(self, xyz_str, offset):
 
 
 # --- Helpers ---
-def parse_cloud_path(path):
-    """
-    Parses a cloud storage path into its bucket name and key/prefix. Supports
-    paths of the form: "{scheme}://bucket_name/prefix" or without a scheme.
-
-    Parameters
-    ----------
-    path : str
-        Path to be parsed.
-
-    Returns
-    -------
-    bucket_name : str
-        Name of the bucket.
-    prefix : str
-        Cloud prefix.
-    """
-    # Remove s3:// if present
-    if path.startswith("s3://"):
-        path = path[len("s3://"):]
-
-    # Remove gs:// if present
-    if path.startswith("gs://"):
-        path = path[len("gs://"):]
-
-    parts = path.split("/", 1)
-    bucket_name = parts[0]
-    prefix = parts[1] if len(parts) > 1 else ""
-    return bucket_name, prefix
-
-
 def to_zipped_point(zip_writer, filename, xyz):
     """
     Writes a point to an SWC file format, which is then stored in a ZIP

src/segmentation_skeleton_metrics/utils/util.py

@@ -450,6 +450,38 @@ def kdtree_query(kdtree, xyz):
     return tuple(kdtree.data[idx])
 
 
+def parse_cloud_path(path):
+    """
+    Parses a cloud storage path into its bucket name and key/prefix. Supports
+    paths of the form: "{scheme}://bucket_name/prefix" or without a scheme.
+
+    Parameters
+    ----------
+    path : str
+        Path to be parsed.
+
+    Returns
+    -------
+    bucket_name : str
+        Name of the bucket.
+    prefix : str
+        Cloud prefix.
+    """
+    # Remove s3:// if present
+    if is_s3_path(path):
+        path = path[len("s3://"):]
+
+    # Remove gs:// if present
+    if is_gcs_path(path):
+        path = path[len("gs://"):]
+
+    # Extract bucket and prefix
+    parts = path.split("/", 1)
+    bucket_name = parts[0]
+    prefix = parts[1] if len(parts) > 1 else ""
+    return bucket_name, prefix
+
+
 def sample_once(my_container):
     """
     Samples a single element from "my_container".