feat: enhance Hausdorff distance computation by utilizing statistics for truth and prediction labels

Author: rhoadesScholar

src/cellmap_segmentation_challenge/evaluate.py

@@ -1,4 +1,5 @@
 import argparse
+import gc
 import json
 import os
 import shutil
@@ -8,12 +9,13 @@
 import numpy as np
 import zarr
 from scipy.spatial.distance import dice
+from scipy.ndimage import distance_transform_edt
 from fastremap import remap, unique
 import cc3d
 from cc3d.types import StatisticsDict, StatisticsSlicesDict
 
 from zarr.errors import PathNotFoundError
-from sklearn.metrics import accuracy_score, jaccard_score
+from sklearn.metrics import jaccard_score
 from tqdm import tqdm
 from upath import UPath
 
@@ -257,12 +259,16 @@ def optimized_hausdorff_distances(
         return np.empty((0,), dtype=np.float32)
 
     voxel_size = np.asarray(voxel_size, dtype=np.float64)
+    truth_stats = cc3d.statistics(truth_label)
+    pred_stats = cc3d.statistics(pred_label)
 
     def get_distance(i: int):
         tid = int(truth_ids[i])
         h_dist = compute_hausdorff_distance_roi(
             truth_label,
+            truth_stats,
             pred_label,
+            pred_stats,
             tid,
             voxel_size,
             hausdorff_distance_max,
@@ -285,62 +291,59 @@ def get_distance(i: int):
     return dists
 
 
-def bbox_for_label_cc3d(
+def bbox_for_label(
     stats: StatisticsDict | StatisticsSlicesDict,
-    label_vol: np.ndarray,
+    ndim: int,
     label_id: int,
 ):
     """
     Try to get bbox without allocating a full boolean mask using cc3d statistics.
     Falls back to mask-based bbox if cc3d doesn't provide expected fields.
     Returns (mins, maxs) inclusive-exclusive in voxel indices, or None if missing.
     """
-    try:
-        # stats = cc3d.statistics(label_vol)
-        # cc3d.statistics usually returns dict-like with keys per label id.
-        # There are multiple API variants; try common patterns.
-        if "bounding_boxes" in stats:
-            bb = stats["bounding_boxes"].get(label_id, None)
-            if bb is None:
-                return None
-            # bb might be (z0,z1,y0,y1,x0,x1) with end exclusive
-            ndim = label_vol.ndim
+    # stats = cc3d.statistics(label_vol)
+    # cc3d.statistics usually returns dict-like with keys per label id.
+    # There are multiple API variants; try common patterns.
+    if "bounding_boxes" in stats:
+        # bounding_boxes is a list where index corresponds to label_id
+        bounding_boxes = stats["bounding_boxes"]
+        if label_id >= len(bounding_boxes):
+            return None
+        bb = bounding_boxes[label_id]
+        if bb is None:
+            return None
+        # bb is a tuple of slices, convert to (mins, maxs)
+        if isinstance(bb, tuple) and all(isinstance(s, slice) for s in bb):
+            mins = [s.start for s in bb]
+            maxs = [s.stop for s in bb]
+            return mins, maxs
+        # bb might be (z0,z1,y0,y1,x0,x1) with end exclusive
+        mins = [bb[2 * k] for k in range(ndim)]
+        maxs = [bb[2 * k + 1] for k in range(ndim)]
+        return mins, maxs
+
+    if label_id in stats:
+        s = stats[label_id]
+        if "bounding_box" in s:
+            bb = s["bounding_box"]
             mins = [bb[2 * k] for k in range(ndim)]
             maxs = [bb[2 * k + 1] for k in range(ndim)]
             return mins, maxs
 
-        if label_id in stats:
-            s = stats[label_id]
-            if "bounding_box" in s:
-                bb = s["bounding_box"]
-                ndim = label_vol.ndim
-                mins = [bb[2 * k] for k in range(ndim)]
-                maxs = [bb[2 * k + 1] for k in range(ndim)]
-                return mins, maxs
-    except Exception:
-        pass
-
-    # Fallback: mask-based (allocates a temporary bool)
-    coords = np.where(label_vol == label_id)
-    if len(coords) == 0 or coords[0].size == 0:
-        return None
-    mins = [int(c.min()) for c in coords]
-    maxs = [int(c.max()) + 1 for c in coords]
-    return mins, maxs
-
 
 def roi_slices_for_pair(
-    truth_label: np.ndarray,
-    pred_label: np.ndarray,
+    truth_stats: StatisticsDict | StatisticsSlicesDict,
+    pred_stats: StatisticsDict | StatisticsSlicesDict,
     tid: int,
     voxel_size,
+    ndim: int,
+    shape: tuple[int, ...],
     max_distance: float,
 ):
     """
     ROI = union(bbox(truth==tid), bbox(pred==tid)) padded by P derived from max_distance.
     Returns tuple of slices suitable for numpy indexing.
     """
-    ndim = truth_label.ndim
     vs = np.asarray(voxel_size, dtype=float)
     if vs.size != ndim:
         # tolerate vs longer (e.g. includes channel), take last ndim
@@ -349,14 +352,12 @@ def roi_slices_for_pair(
     # padding per axis in voxels
     pad = np.ceil(max_distance / vs).astype(int) + 2
 
-    truth_stats = cc3d.statistics(truth_label)
-    tb = bbox_for_label_cc3d(truth_stats, truth_label, tid)
+    tb = bbox_for_label(truth_stats, ndim, tid)
     if tb is None:
         return None  # should not happen for tid from truth_ids
 
     tmins, tmaxs = tb
-    pred_stats = cc3d.statistics(pred_label)
-    pb = bbox_for_label_cc3d(pred_stats, pred_label, tid)
+    pb = bbox_for_label(pred_stats, ndim, tid)
     if pb is None:
         pmins, pmaxs = tmins, tmaxs
     else:
@@ -366,7 +367,6 @@ def roi_slices_for_pair(
     maxs = [max(tmaxs[d], pmaxs[d]) for d in range(ndim)]
 
     # expand and clamp
-    shape = truth_label.shape
     out_slices = []
     for d in range(ndim):
         a = max(0, mins[d] - int(pad[d]))
@@ -377,7 +377,9 @@ def roi_slices_for_pair(
 
 def compute_hausdorff_distance_roi(
     truth_label: np.ndarray,
+    truth_stats: StatisticsDict | StatisticsSlicesDict,
     pred_label: np.ndarray,
+    pred_stats: StatisticsDict | StatisticsSlicesDict,
     tid: int,
     voxel_size,
     max_distance: float,
@@ -388,9 +390,17 @@ def compute_hausdorff_distance_roi(
     Same metric as compute_hausdorff_distance(), but operates on an ROI slice
     and builds masks only inside ROI.
     """
-    from scipy.ndimage import distance_transform_edt
+    ndim = truth_label.ndim
 
-    roi = roi_slices_for_pair(truth_label, pred_label, tid, voxel_size, max_distance)
+    roi = roi_slices_for_pair(
+        truth_stats,
+        pred_stats,
+        tid,
+        voxel_size,
+        ndim,
+        truth_label.shape,
+        max_distance,
+    )
     if roi is None:
         return float(max_distance)
 
@@ -407,7 +417,6 @@ def compute_hausdorff_distance_roi(
     if a_n == 0 or b_n == 0:
         return float(max_distance)
 
-    ndim = truth_label.ndim
     vs = np.asarray(voxel_size, dtype=np.float64)
     if vs.size != ndim:
         vs = vs[-ndim:]
@@ -497,7 +506,7 @@ def score_instance(
 
     # Compute the scores
     logging.info("Computing accuracy score...")
-    accuracy = accuracy_score(truth_label.ravel(), pred_label.ravel())
+    accuracy = float((truth_label == pred_label).mean())
     # When there are no Hausdorff distances, use np.inf so that
     # normalize_distance(hausdorff_dist, voxel_size) returns 0.0. This encodes
     # the absence of matched instances and ensures the combined_score is 0.0.
@@ -612,6 +621,7 @@ def score_label(
         mask = zarr.open(mask_path.path, mode="r")[:]
         pred_label = pred_label * mask
         truth_label = truth_label * mask
+        del mask
 
     # Compute the scores
     if label_name in instance_classes:
@@ -628,6 +638,9 @@ def score_label(
     results["num_voxels"] = int(np.prod(truth_label.shape))
     results["voxel_size"] = crop.voxel_size
     results["is_missing"] = False
+    # drop big arrays before returning
+    del truth_label, pred_label, truth_label_ds
+    gc.collect()
     return crop_name, label_name, results
 
 

tests/test_evaluate_metrics.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+import cc3d
 import numpy as np
 import zarr
 from fastremap import unique
@@ -114,8 +115,18 @@ def test_roi_hausdorff_identical_instance_is_zero():
     truth[2, 2] = tid
     pred[2, 2] = tid
 
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
+
     d = ev.compute_hausdorff_distance_roi(
-        truth, pred, tid, voxel_size=(1.0, 1.0), max_distance=10.0, method="standard"
+        truth,
+        truth_stats,
+        pred,
+        pred_stats,
+        tid,
+        voxel_size=(1.0, 1.0),
+        max_distance=10.0,
+        method="standard",
     )
     assert np.isclose(d, 0.0)
 
@@ -133,9 +144,14 @@ def test_roi_hausdorff_matches_full_reference_standard_and_modified():
     voxel_size = (1.0, 1.0)
     max_distance = 100.0
 
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
+
     d_roi = ev.compute_hausdorff_distance_roi(
         truth,
+        truth_stats,
         pred,
+        pred_stats,
         tid,
         voxel_size=voxel_size,
         max_distance=max_distance,
@@ -154,7 +170,9 @@ def test_roi_hausdorff_matches_full_reference_standard_and_modified():
 
     d_roi_mod = ev.compute_hausdorff_distance_roi(
         truth,
+        truth_stats,
         pred,
+        pred_stats,
         tid,
         voxel_size=voxel_size,
         max_distance=max_distance,
@@ -188,9 +206,14 @@ def test_roi_hausdorff_percentile_matches_full_reference():
     voxel_size = (1.0, 1.0)
     max_distance = 100.0
 
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
+
     d_roi = ev.compute_hausdorff_distance_roi(
         truth,
+        truth_stats,
         pred,
+        pred_stats,
         tid,
         voxel_size=voxel_size,
         max_distance=max_distance,
@@ -221,16 +244,32 @@ def test_roi_hausdorff_empty_sets_and_missing_instance():
 
     # present only in truth -> max_distance
     truth[0, 0] = tid
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
     d1 = ev.compute_hausdorff_distance_roi(
-        truth, pred, tid, voxel_size=voxel_size, max_distance=max_distance
+        truth,
+        truth_stats,
+        pred,
+        pred_stats,
+        tid,
+        voxel_size=voxel_size,
+        max_distance=max_distance,
     )
     assert np.isclose(d1, max_distance)
 
     # present only in pred -> max_distance
     truth[0, 0] = 0
     pred[0, 0] = tid
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
     d2 = ev.compute_hausdorff_distance_roi(
-        truth, pred, tid, voxel_size=voxel_size, max_distance=max_distance
+        truth,
+        truth_stats,
+        pred,
+        pred_stats,
+        tid,
+        voxel_size=voxel_size,
+        max_distance=max_distance,
     )
     assert np.isclose(d2, max_distance)
 
@@ -248,9 +287,14 @@ def test_roi_hausdorff_clips_to_max_distance_matches_reference():
     voxel_size = (1.0, 1.0)
     max_distance = 3.0
 
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
+
     d_roi = ev.compute_hausdorff_distance_roi(
         truth,
+        truth_stats,
         pred,
+        pred_stats,
         tid,
         voxel_size=voxel_size,
         max_distance=max_distance,
@@ -282,9 +326,14 @@ def test_roi_hausdorff_anisotropic_voxel_size_matches_reference():
     voxel_size = (2.0, 0.5)  # physical distance = 2 * 0.5 = 1.0
     max_distance = 100.0
 
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
+
     d_roi = ev.compute_hausdorff_distance_roi(
         truth,
+        truth_stats,
         pred,
+        pred_stats,
         tid,
         voxel_size=voxel_size,
         max_distance=max_distance,
@@ -323,8 +372,17 @@ def _fake_roi(*args, **kwargs):
     # Make tid present in exactly one volume so we don't trigger the "both absent -> 0" shortcut
     truth[2, 2] = tid  # present in truth only
 
+    truth_stats = cc3d.statistics(truth)
+    pred_stats = cc3d.statistics(pred)
+
     d = ev.compute_hausdorff_distance_roi(
-        truth, pred, tid, voxel_size=(1.0, 1.0), max_distance=7.0
+        truth,
+        truth_stats,
+        pred,
+        pred_stats,
+        tid,
+        voxel_size=(1.0, 1.0),
+        max_distance=7.0,
     )
     assert np.isclose(d, 7.0)