refactor: improved tensorstore reader

Author: anna-grim

pyproject.toml

@@ -16,6 +16,7 @@ dynamic = ["version"]
 
 dependencies = [
     "boto3",
+    'cloud-volume[crackle,seg-codecs]',
     "google-cloud-storage",
     "networkx",
     "numpy",

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -364,7 +364,7 @@ def compute_split_metrics(self, key, n_split_edges):
         gt_rl = self.graphs[key].run_length
         split_rate = rl / n_splits if n_splits > 0 else np.nan
 
-        # Update metrics
+        # Record metrics
         self.metrics.at[key, "# Splits"] = n_splits
         self.metrics.at[key, "Split Rate"] = split_rate
         self.metrics.at[key, "% Split Edges"] = round(p_split, 2)
@@ -404,7 +404,7 @@ def count_merges(self, key, kdtree):
         """
         Counts the number of label merges for a given graph key based on
         whether the fragment graph corresponding to a label has a node that is
-        more that 200ums away from the nearest point in
+        more that 40ums away from the nearest point in
         "kdtree".
 
         Parameters
@@ -414,20 +414,19 @@ def count_merges(self, key, kdtree):
         kdtree : scipy.spatial.KDTree
             A KD-tree built from voxels in graph corresponding to "key".
         """
-        # Iterate over fragments that intersect with GT skeleton
+        # Iterate over fragments that intersect with GT graphs
         for label in self.get_node_labels(key):
             nodes = self.graphs[key].nodes_with_label(label)
-            if len(nodes) > 70:
+            if len(nodes) > 60:
                 for label in self.label_handler.get_class(label):
                     if label in self.fragment_ids:
-                        self.is_fragment_merge(key, label, kdtree)
+                        self.check_fragment_for_merges(key, label, kdtree)
 
-    def is_fragment_merge(self, key, label, kdtree):
+    def check_fragment_for_merges(self, key, label, kdtree):
         """
-        Determines whether fragment corresponding to "label" is falsely merged
-        to graph corresponding to "key". A fragment is said to be merged if
-        there is a node in the fragment more than 200ums away from the nearest
-        point in "kdtree".
+        Checks whether the fragment corresponding to "label" has a merge
+        mistake. A fragment has a merge mistake if it has a leaf node more
+        than 40μm away from the ground-truth graph corresponding to "key".
 
         Parameters
         ----------
@@ -501,7 +500,7 @@ def find_merge_site(self, key, kdtree, fragment_graph, source, visited):
                             gutil.write_graph(
                                 self.gt_graphs[key], self.merge_writer
                             )
-                        return
+                        return None
 
     def is_valid_merge(self, graph, kdtree, root):
         n_hits = 0

src/segmentation_skeleton_metrics/utils/img_util.py

@@ -10,6 +10,7 @@
 """
 
 from abc import ABC, abstractmethod
+from cloudvolume import CloudVolume
 from tifffile import imread
 
 import io
@@ -18,6 +19,8 @@
 import tensorstore as ts
 import zipfile
 
+from segmentation_skeleton_metrics.utils import util
+
 
 class ImageReader(ABC):
     """
@@ -119,20 +122,43 @@ def __init__(self, img_path, driver):
         driver : str
             Storage driver needed to read the image.
         """
-        self.driver = driver
+        self.driver = self.get_driver(img_path)
         super().__init__(img_path)
 
+    def get_driver(self, img_path):
+        """
+        Gets the storage driver needed to read the image.
+
+        Returns
+        -------
+        str
+            Storage driver needed to read the image.
+        """
+        if ".zarr" in img_path:
+            return "zarr"
+        elif ".n5" in img_path:
+            return "n5"
+        elif is_neuroglancer_precomputed(img_path):
+            return "neuroglancer_precomputed"
+        else:
+            raise ValueError(f"Unsupported image format: {img_path}")
+
     def _load_image(self):
         """
         Loads image using the TensorStore library.
         """
+        # Extract metadata
+        bucket_name, path = util.parse_cloud_path(self.img_path)
+        storage_driver = get_storage_driver(self.img_path)
+
+        # Load image
         self.img = ts.open(
             {
                 "driver": self.driver,
                 "kvstore": {
-                    "driver": "gcs",
-                    "bucket": "allen-nd-goog",
-                    "path": self.img_path,
+                    "driver": storage_driver,
+                    "bucket": bucket_name,
+                    "path": path,
                 },
                 "context": {
                     "cache_pool": {"total_bytes_limit": 1000000000},
@@ -142,8 +168,6 @@ def _load_image(self):
                 "recheck_cached_data": "open",
             }
         ).result()
-        if self.driver == "neuroglancer_precomputed":
-            return self.img[ts.d["channel"][0]]
 
     def read(self, voxel, shape):
         """
@@ -230,7 +254,51 @@ def _load_zipped_image(self):
                 self.img = imread(io.BytesIO(f.read()))
 
 
-# --- Miscellaneous ---
+# --- Helpers ---
+def get_storage_driver(img_path):
+    """
+    Gets the storage driver needed to read the image.
+
+    Parameters
+    ----------
+    img_path : str
+        Image path to be checked.
+
+    Returns
+    -------
+    str
+        Storage driver needed to read the image.
+    """
+    if util.is_s3_path(img_path):
+        return "s3"
+    elif util.is_gcs_path(img_path):
+        return "gcs"
+    else:
+        raise ValueError(f"Unsupported path type: {img_path}")
+
+
+def is_neuroglancer_precomputed(path):
+    """
+    Checks if the path points to a neuroglancer precomputed dataset.
+
+    Parameters
+    ----------
+    path : str
+        Path to be checked.
+
+    Returns
+    -------
+    bool
+        Indication of whether the path points to a neuroglancer precomputed
+        dataset.
+    """
+    try:
+        vol = CloudVolume(path)
+        return all(k in vol.info for k in ["data_type", "scales", "type"])
+    except Exception:
+        return False
+
+
 def to_physical(voxel, anisotropy):
     """
     Converts a voxel coordinate to a physical coordinate by applying the