Optimizations for finding decision points on sparse volumes

- If the number of voxels is below a certain threshold do not count it as a segment to expand.
- If number of unique segments in a subvolume is <2 do not expand.

PiperOrigin-RevId: 888191240

Author: timblakely

ffn/utils/decision_point.py

@@ -24,11 +24,14 @@
 from scipy import ndimage
 
 
-def find_decision_points(seg: np.ndarray,
-                         voxel_size: Sequence[float],
-                         max_distance: Optional[float] = None,
-                         subvol_box: Optional[bounding_box.BoundingBox] = None
-                        ) -> dict[tuple[int, int], tuple[float, np.ndarray]]:
+def find_decision_points(
+    seg: np.ndarray,
+    voxel_size: Sequence[float],
+    max_distance: Optional[float] = None,
+    subvol_box: Optional[bounding_box.BoundingBox] = None,
+    optimize_sparse: bool = False,
+    sparse_noise_threshold: int = 0,
+) -> dict[tuple[int, int], tuple[float, np.ndarray]]:
   """Identifies decision points in a segmentation subvolume.
 
   Args:
@@ -39,12 +42,34 @@ def find_decision_points(seg: np.ndarray,
     subvol_box: selector for a subvolume within `seg` within which
       to search for decision points; the whole subvolume is always used
       to compute the distance transform
+    optimize_sparse: if True, first counts the number of segments in `seg`
+      and returns early if there are fewer than 2.
+    sparse_noise_threshold: if > 0 and `optimize_sparse` is True, ignores
+      components with voxel counts <= this threshold when counting segments.
 
   Returns:
     dict from segment ID pairs to tuples of:
       approximate physical distance from the segment to the decision point
       (x, y, z) decision point
   """
+  if optimize_sparse:
+    # Use pandas to quickly get the unique labels and their counts (excluding 0)
+    counts = pd.Series(seg.ravel()).value_counts(sort=False)
+    if 0 in counts:
+      counts = counts.drop(0)
+
+    if sparse_noise_threshold > 0:
+      small_segs = counts[counts <= sparse_noise_threshold].index.to_numpy()
+      if len(small_segs) > 0:
+        # Zero out small segments and update counts
+        seg[np.isin(seg, small_segs)] = 0
+        counts = counts.drop(small_segs)
+
+    if len(counts) <= 1:
+      # If there are 0 or 1 unique segments (excluding background),
+      # they cannot possibly touch another segment.
+      return {}
+
   # EDT is the Euclidean Distance Transform, specifying how far voxels added
   # in 'expanded_seg' are from the seeds in 'seg'.
   expanded_seg, edt = labels.watershed_expand(seg, voxel_size, max_distance)

ffn/utils/tests/decision_point_test.py

@@ -55,6 +55,39 @@ def test_find_decision_point(self):
     self.assertIn((1, 2), points)
     self.assertLen(points, 1)
 
+  def test_find_decision_point_optimize_sparse(self):
+    # 2 segments, but one is very small and should be filtered out
+    seg = np.zeros((100, 80, 60), dtype=np.uint64)
+    seg[:40, :, :] = 1
+    # 2 voxels of label 2
+    seg[60, 0, 0] = 2
+    seg[61, 0, 0] = 2
+
+    # Without optimization, 1 and 2 might connect if they are grown far enough.
+    points = decision_point.find_decision_points(seg, (1, 1, 1))
+    self.assertIn((1, 2), points)
+
+    # With optimization but threshold 0, they still connect (no size filtering)
+    points = decision_point.find_decision_points(
+        seg, (1, 1, 1), optimize_sparse=True, sparse_noise_threshold=0
+    )
+    self.assertIn((1, 2), points)
+
+    # With optimization and threshold >= 2, label 2 is zeroed.
+    # We're left with label 1 (size > 2), so 1 segment -> returns empty
+    points = decision_point.find_decision_points(
+        seg, (1, 1, 1), optimize_sparse=True, sparse_noise_threshold=2
+    )
+    self.assertEmpty(points)
+
+    # With just 1 segment from the start
+    seg = np.zeros((100, 80, 60), dtype=np.uint64)
+    seg[:40, :, :] = 1
+    points = decision_point.find_decision_points(
+        seg, (1, 1, 1), optimize_sparse=True, sparse_noise_threshold=0
+    )
+    self.assertEmpty(points)
+
 
 if __name__ == '__main__':
   absltest.main()