diff --git a/src/segmentation_skeleton_metrics/data_handling/graph_loading.py b/src/segmentation_skeleton_metrics/data_handling/graph_loading.py index 8781bb3..f3591f7 100644 --- a/src/segmentation_skeleton_metrics/data_handling/graph_loading.py +++ b/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} diff --git a/src/segmentation_skeleton_metrics/skeleton_metrics.py b/src/segmentation_skeleton_metrics/skeleton_metrics.py index b84db24..4560890 100644 --- a/src/segmentation_skeleton_metrics/skeleton_metrics.py +++ b/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