refactor: adjustments for split eval

src/segmentation_skeleton_metrics/data_handling/graph_classes.py

@@ -10,7 +10,7 @@
 
 from collections import defaultdict, deque
 from io import StringIO
-from scipy.spatial import distance, KDTree
+from scipy.spatial import KDTree
 
 import networkx as nx
 import numpy as np
@@ -150,7 +150,7 @@ def dist(self, i, j):
         float
             Distance between voxel coordinates of the given nodes.
         """
-        return distance.euclidean(self.node_voxel[i], self.node_voxel[j])
+        return np.linalg.norm(self.node_voxel[i] - self.node_voxel[j])
 
     def physical_dist(self, i, j):
         """
@@ -170,7 +170,7 @@ def physical_dist(self, i, j):
             Euclidean distance between physical coordinates of the given
             nodes.
         """
-        return distance.euclidean(self.node_xyz(i), self.node_xyz(j))
+        return np.linalg.norm(self.node_xyz(i) - self.node_xyz(j))
 
     def leafs(self):
         """
@@ -562,7 +562,7 @@ def __init__(
         name=None,
         label=None,
         segment_id=None,
-        swc_color=""
+        swc_color="",
     ):
         """
         Instantiates a FragmentGraph object.
@@ -585,9 +585,7 @@ def __init__(
             Color used in SWC file of the graph. Default is an empty string.
         """
         # Call parent class
-        super().__init__(
-            anisotropy=anisotropy, name=name, swc_color=swc_color
-        )
+        super().__init__(anisotropy=anisotropy, name=name, swc_color=swc_color)
 
         # Instance attributes
         self.label = label

src/segmentation_skeleton_metrics/data_handling/graph_loading.py

@@ -291,24 +291,30 @@ def to_graph(self, swc_dict):
         Dict[str, SkeletonGraph]
             Graph built from an SWC file.
         """
-        # Initialize graph
+        # Initializations
         graph = self._init_graph(swc_dict)
-        key = graph.name if self.is_groundtruth else graph.label
-
-        # Build graph structure
-        id_lookup = dict()
-        for i, id_i in enumerate(swc_dict["id"]):
-            id_lookup[id_i] = i
-            if swc_dict["pid"][i] != -1:
-                parent = id_lookup[swc_dict["pid"][i]]
-                graph.add_edge(i, parent)
-                graph.run_length += graph.dist(i, parent)
-
-        # Apply voxel coordinate conversion (if applicable)
+        ids = swc_dict["id"]
+        pids = swc_dict["pid"]
+        voxels = swc_dict["voxel"]
+
+        # Build id->index lookup
+        id_lookup = {id_i: i for i, id_i in enumerate(ids)}
+
+        # Vectorized edge building
+        children = np.where(pids != -1)[0]
+        parents = np.array([id_lookup[pids[i]] for i in children], dtype=int)
+        edges = list(zip(children.tolist(), parents.tolist()))
+        graph.add_edges_from(edges)
+
+        # Vectorized run length
+        graph.run_length = np.linalg.norm(
+            voxels[children] - voxels[parents], axis=1
+        ).sum()
+
         if self.use_anisotropy:
             graph.node_voxel = (graph.node_voxel / self.anisotropy).astype(int)
             graph.node_voxel[:, [0, 2]] = graph.node_voxel[:, [2, 0]]
-        return {key: graph}
+        return {graph.name: graph}
 
     def _init_graph(self, swc_dict):
         """

src/segmentation_skeleton_metrics/data_handling/swc_loading.py

@@ -100,7 +100,7 @@ def __call__(self, swc_pointer):
             # Local SWC files
             paths = util.read_paths(swc_pointer, extension=".swc")
             if len(paths) > 0:
-                return self.read_swcs(paths, self.read_swc)
+                return self.read_swcs(paths)
 
             raise Exception("Directory is Invalid!")
 
@@ -144,14 +144,12 @@ def read_swc(self, path):
         else:
             return None
 
-    def read_swcs(self, swc_paths, read_fn):
+    def read_swcs(self, swc_paths):
         """
         Reads SWC files stored in a GCS or S3 bucket.
 
         Parameters
         ----------
-        bucket_name : str
-            Name of bucket containing SWC files.
         swc_paths : List[str]
             List of paths to SWC files to be read.
 
@@ -166,7 +164,7 @@ def read_swcs(self, swc_paths, read_fn):
             threads = set()
             for path in swc_paths:
                 if self.confirm_read(os.path.basename(path)):
-                    threads.add(executor.submit(read_fn, path))
+                    threads.add(executor.submit(self.read_swc, path))
 
             # Store results
             swc_dicts = deque()
@@ -283,22 +281,22 @@ def read_from_cloud(self, path):
             Dictionaries whose keys and values are the attribute names and
             values from an SWC file.
         """
-        # Extact info
+        # Extract info
         assert util.is_s3_path(path) or util.is_gcs_path(path)
         use_s3 = util.is_s3_path(path)
 
-        # List filenames
+        # List paths
         swc_paths = util.list_cloud_paths(path, ".swc")
         zip_paths = util.list_cloud_paths(path, ".zip")
 
         # Call reader
         if swc_paths:
-            return self.read_swcs(swc_paths, self.read_swc)
+            return self.read_swcs(swc_paths)
         elif zip_paths:
             read_fn = self.read_s3_zip if use_s3 else self.read_gcs_zip
             return self.read_zips(zip_paths, read_fn)
-
-        raise Exception(f"SWC Pointer is invalid {path}")
+        else:
+            return list()
 
     def read_gcs_swc(self, path):
         """
@@ -341,33 +339,14 @@ def read_gcs_zip(self, path):
             Dictionaries whose keys and values are the attribute names and
             values from an SWC file.
         """
-        # Download ZIP
-        bucket_name, path = util.parse_cloud_path(path)
+        bucket_name, key = util.parse_cloud_path(path)
         bucket = storage.Client().bucket(bucket_name)
         try:
-            zip_content = bucket.blob(path).download_as_bytes()
+            zip_content = bucket.blob(key).download_as_bytes()
         except TransportError:
-            print(f"Failed to read {zip_path}!")
+            print(f"Failed to read {path}!")
             return deque()
-
-        # Parse ZIP contents
-        swc_dicts = deque()
-        with ZipFile(BytesIO(zip_content), "r") as zf:
-            with ThreadPoolExecutor() as executor:
-                # Assign threads
-                threads = set()
-                for name in zf.namelist():
-                    if self.confirm_read(name):
-                        threads.add(
-                            executor.submit(self.read_zipped_swc, zf, name)
-                        )
-
-                # Process results
-                for thread in as_completed(threads):
-                    result = thread.result()
-                    if result:
-                        swc_dicts.append(result)
-        return swc_dicts
+        return self._parse_zip_bytes(zip_content)
 
     def read_s3_zip(self, path):
         """
@@ -385,29 +364,10 @@ def read_s3_zip(self, path):
             Dictionaries whose keys and values are the attribute names and
             values from an SWC file.
         """
-        # Initialize cloud reader
         bucket, key = util.parse_cloud_path(path)
         s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED))
         zip_content = s3.get_object(Bucket=bucket, Key=key)["Body"].read()
-
-        # Parse ZIP
-        with ZipFile(BytesIO(zip_content), "r") as zf:
-            with ThreadPoolExecutor() as executor:
-                # Assign threads
-                threads = set()
-                for name in zf.namelist():
-                    if self.confirm_read(name):
-                        threads.add(
-                            executor.submit(self.read_zipped_swc, zf, name)
-                        )
-
-                # Store results
-                swc_dicts = deque()
-                for thread in as_completed(threads):
-                    result = thread.result()
-                    if result:
-                        swc_dicts.append(result)
-        return swc_dicts
+        return self._parse_zip_bytes(zip_content)
 
     def confirm_read(self, path):
         """
@@ -427,23 +387,19 @@ def confirm_read(self, path):
         name = os.path.splitext(os.path.basename(path))[0]
         return name in self.swc_names if self.swc_names else True
 
-    # -- Process Text ---
-    def iterator(self, iterator):
-        """
-        Gets an iterator that optionally displays a progress bar.
-
-        Parameters
-        ----------
-        iterator : iterable
-            Object to be iterated over.
-
-        Returns
-        -------
-        tqdm.tqdm
-            Iterator that is optionally wrapped in a progress bar.
-        """
-        return tqdm(iterator, desc="Read SWCs") if self.verbose else iterator
+    def _parse_zip_bytes(self, zip_content):
+        with ZipFile(BytesIO(zip_content), "r") as zf:
+            names = [f for f in zf.namelist() if f.endswith(".swc")]
+            with ThreadPoolExecutor() as executor:
+                threads = {
+                    executor.submit(self.read_zipped_swc, zf, name)
+                    for name in names
+                }
+                return deque(
+                    t.result() for t in as_completed(threads) if t.result()
+                )
 
+    # -- Process Text ---
     def manual_progress_bar(self, total):
         """
         Gets progress bar that needs to be updated manually.
@@ -521,14 +477,15 @@ def process_content(self, content):
                 offset = self.read_coordinate(parts[2:5])
             if not line.startswith("#") and len(line.strip()) > 0:
                 return content[i:], offset
+        return [], offset
 
-    def read_coordinate(self, xyz_str, offset=(0, 0, 0)):
+    def read_coordinate(self, coord_str, offset=(0, 0, 0)):
         """
         Reads a coordinate from a string and converts it to voxel coordinates.
 
         Parameters
         ----------
-        xyz_str : str
+        coord_str : str
             Coordinate stored as a string.
         offset : Tuple[int]
             Offset of coordinates in SWC file. Default is (0, 0, 0).
@@ -538,4 +495,4 @@ def read_coordinate(self, xyz_str, offset=(0, 0, 0)):
         Tuple[int]
             xyz coordinates of an entry from an SWC file.
         """
-        return [float(xyz_str[i]) + offset[i] for i in range(3)]
+        return [float(coord_str[i]) + offset[i] for i in range(3)]