refactor: label handler

src/segmentation_skeleton_metrics/data_handling/graph_loading.py

@@ -523,13 +523,13 @@ def __init__(
         self.inverse_mapping = dict()  # Maps class id to list of labels
 
         # Set label mapping
-        if use_segment_mapping:
-            self.set_segment_mappings()
+        if label_pairs:
+            self.set_equivalence_mapping(label_pairs)
         else:
-            self.set_mappings(label_pairs)
+            self.set_segment_mapping()
 
     # --- Constructor Helpers ---
-    def set_mappings(self, label_pairs):
+    def set_equivalence_mapping(self, label_pairs):
         """
         Stores dictionaries that map between segment IDs and equivalence class
         IDS.
@@ -548,18 +548,21 @@ def set_mappings(self, label_pairs):
                 self.mapping[label] = class_id
                 self.inverse_mapping[class_id].add(label)
 
-    def set_segment_mappings(self):
+    def set_segment_mapping(self):
         """
         Stores dictionaries that map between segment IDs and labels.
         """
-        assert self.labels
-        self.mapping = {"0": "0"}
-        self.inverse_mapping = defaultdict(set)
-        self.inverse_mapping["0"] = {"0"}
-        for label in map(str, self.labels):
-            segment_id = util.get_segment_id(label)
-            self.mapping[label] = segment_id
-            self.inverse_mapping[segment_id].add(label)
+        if self.labels:
+            self.mapping = {"0": "0"}
+            self.inverse_mapping = defaultdict(set)
+            self.inverse_mapping["0"] = {"0"}
+            for label in map(str, self.labels):
+                segment_id = util.get_segment_id(label)
+                self.mapping[label] = segment_id
+                self.inverse_mapping[segment_id].add(label)
+        else:
+            self.mapping = LazyMapping()
+            self.inverse_mapping = InverseLazyMapping()
 
     def label_equiv_classes(self, label_pairs):
         """
@@ -593,7 +596,7 @@ def get(self, label):
 
         Returns
         -------
-        int
+        str
             Class ID corresponding to the label.
         """
         return self.mapping.get(label, "0") if self.labels else str(label)
@@ -617,3 +620,49 @@ def node_labels(self, graph):
         if self.labels:
             labels = set().union(*(self.inverse_mapping[u] for u in labels))
         return labels
+
+
+# --- Helpers ---
+class LazyMapping(dict):
+    """
+    Dictionary-like mapping that lazily converts labels to segment IDs.
+    """
+
+    def __missing__(self, label):
+        """
+        Gets the segment ID corresponding to the given label.
+
+        Parameters
+        ----------
+        label : hashable
+            Label to be converted into a segment ID.
+
+        Returns
+        -------
+        str
+            Segment ID corresponding to the given label.
+        """
+        return util.get_segment_id(str(label))
+
+
+class InverseLazyMapping(dict):
+    """
+    Dictionary-like inverse mapping from segment IDs to label sets.
+    """
+
+    def __missing__(self, label):
+        """
+        Gets the segment ID corresponding to the given label.
+
+        Parameters
+        ----------
+        label : hashable
+            Label to be returned as singleton set.
+
+        Returns
+        -------
+        Set[str]
+            Singleton containing the segment ID corresponding to the given
+            label.
+        """
+        return {"0"} if label == "0" else {label}

src/segmentation_skeleton_metrics/skeleton_metrics.py

@@ -342,7 +342,7 @@ class MergeCountMetric(SkeletonMetric):
     A skeleton metric subclass that counts the number merges.
     """
 
-    merge_dist_threshold = 50
+    dist_away_threshold = 50
 
     def __init__(self, verbose=True):
         """
@@ -428,7 +428,7 @@ def search_for_merges(self, gt_graph, fragment_graph):
             dist, _ = gt_graph.kdtree.query(xyz)
 
             # Check if distance to ground truth flags a merge mistake
-            if dist > MergeCountMetric.merge_dist_threshold:
+            if dist > MergeCountMetric.dist_away_threshold:
                 self.find_merge_site(gt_graph, fragment_graph, leaf, visited)
 
     def find_merge_site(self, gt_graph, fragment_graph, source, visited):
@@ -916,7 +916,7 @@ def compute_added_length(self, gt_graph, fragment_graph):
         # Remove nodes close to ground truth
         xyz_arr = fragment_graph.node_voxel * fragment_graph.anisotropy
         dists, _ = gt_graph.kdtree.query(xyz_arr)
-        max_dist = MergeCountMetric.merge_dist_threshold
+        max_dist = MergeCountMetric.dist_away_threshold
         fragment_graph.remove_nodes_from(np.where(dists < max_dist)[0])
 
         # Compute cable length