perf(postproc): rewrite stardist postproc

Author: okunator

cellseg_models_pytorch/postproc/__init__.py

@@ -13,7 +13,8 @@
 from .functional.drfns import post_proc_drfns
 from .functional.hovernet import post_proc_hovernet
 from .functional.omnipose import get_masks_omnipose, post_proc_omnipose
-from .functional.stardist import post_proc_stardist, post_proc_stardist_orig
+from .functional.stardist.nms import get_bboxes
+from .functional.stardist.stardist import post_proc_stardist, post_proc_stardist_orig
 
 POSTPROC_LOOKUP = {
     "stardist_orig": post_proc_stardist_orig,
@@ -44,4 +45,5 @@
     "post_proc_drfns",
     "post_proc_dcan",
     "post_proc_dran",
+    "get_bboxes",
 ]

cellseg_models_pytorch/postproc/functional/stardist/nms.py

@@ -0,0 +1,175 @@
+import math
+from typing import List, Sequence, Tuple
+
+import numpy as np
+from numba import njit, prange
+from scipy.spatial import KDTree
+
+from ....utils import intersection
+
+__all__ = ["get_bboxes", "nms_stardist"]
+
+
+@njit(parallel=True)
+def get_bboxes(
+    dist: np.ndarray, points: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
+    """Get bounding boxes from the non-zero pixels of the radial distance maps.
+
+    This is basically a translation from the stardist repo cpp code to python
+
+    NOTE: jit compiled and parallelized with numba.
+
+    Parameters
+    ----------
+        dist : np.ndarray
+            The non-zero values of the radial distance maps. Shape: (n_nonzero, n_rays).
+        points : np.ndarray
+            The yx-coordinates of the non-zero points. Shape (n_nonzero, 2).
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
+        Returns the x0, y0, x1, y1 bbox coordinates, bbox areas and the maximum
+        radial distance in the image.
+    """
+    n_polys = dist.shape[0]
+    n_rays = dist.shape[1]
+
+    bbox_x1 = np.zeros(n_polys)
+    bbox_x2 = np.zeros(n_polys)
+    bbox_y1 = np.zeros(n_polys)
+    bbox_y2 = np.zeros(n_polys)
+
+    areas = np.zeros(n_polys)
+    angle_pi = 2 * math.pi / n_rays
+    max_dist = 0
+
+    for i in prange(n_polys):
+        max_radius_outer = 0
+        py = points[i, 0]
+        px = points[i, 1]
+
+        for k in range(n_rays):
+            d = dist[i, k]
+            y = py + d * np.sin(angle_pi * k)
+            x = px + d * np.cos(angle_pi * k)
+
+            if k == 0:
+                bbox_x1[i] = x
+                bbox_x2[i] = x
+                bbox_y1[i] = y
+                bbox_y2[i] = y
+            else:
+                bbox_x1[i] = min(x, bbox_x1[i])
+                bbox_x2[i] = max(x, bbox_x2[i])
+                bbox_y1[i] = min(y, bbox_y1[i])
+                bbox_y2[i] = max(y, bbox_y2[i])
+
+            max_radius_outer = max(d, max_radius_outer)
+
+        areas[i] = (bbox_x2[i] - bbox_x1[i]) * (bbox_y2[i] - bbox_y1[i])
+        max_dist = max(max_dist, max_radius_outer)
+
+    return bbox_x1, bbox_y1, bbox_x2, bbox_y2, areas, max_dist
+
+
+@njit
+def _suppress_bbox(
+    query: Sequence[int],
+    current_idx: int,
+    boxes: np.ndarray,
+    areas: np.ndarray,
+    suppressed: List[bool],
+    iou_threshold: float = 0.5,
+) -> np.ndarray:
+    """Inner loop of the stardist nms algorithm where bboxes are suppressed.
+
+    NOTE: Numba compiled only for performance.
+          Parallelization had only a negative effect on run-time on.
+          12-core hyperthreaded Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz.
+    """
+    for i in range(len(query)):
+        query_idx = query[i]
+
+        if suppressed[query_idx]:
+            continue
+
+        overlap = intersection(boxes[current_idx], boxes[query_idx])
+        iou = overlap / min(areas[current_idx] + 1e-10, areas[query_idx] + 1e-10)
+        suppressed[query_idx] = iou > iou_threshold
+
+    return suppressed
+
+
+def nms_stardist(
+    boxes: np.ndarray,
+    points: np.ndarray,
+    scores: np.ndarray,
+    areas: np.ndarray,
+    max_dist: float,
+    score_threshold: float = 0.5,
+    iou_threshold: float = 0.5,
+) -> np.ndarray:
+    """Non maximum suppression for stardist bboxes.
+
+    NOTE: This implementation relies on `scipy.spatial` `KDTree`
+
+    NOTE: This version of nms is faster than the original one in stardist repo
+    and is fully written in python. The differenecs in the resulting instance
+    segmentation masks are neglible.
+
+    Parameters
+    ----------
+        boxes : np.ndarray
+            An array of bbox coords in pascal VOC format (x0, y0, x1, y1).
+            Shape: (n_points, 4). Dtype: float64.
+        points : np.ndarray
+            The yx-coordinates of the non-zero points. Shape (n_points, 2). Dtype: int64
+        scores : np.ndarray
+            The probability values at the point coordinates. Shape (n_points,).
+            Dtype: float32/float64.
+        areas : np.ndarray
+            The areas of the bounding boxes at the point coordinates. Shape (n_points,).
+            Dtype: float32/float64.
+        radius_outer : np.ndarray
+            The radial distances to background at each point. Shape (n_points, )
+        max_dist : float
+            The maximum radial distance of all the radial distances
+        score_threshold : float, default=0.5
+            Threshold for the probability distance map.
+        iou_threshold : float, default=0.5
+            Threshold for the IoU metric deciding whether to suppres a bbox.
+
+    Returns
+    -------
+        np.ndarray:
+            The indices of the bboxes that are not suppressed. Shape: (n_kept, ).
+    """
+    keep = []
+
+    if len(boxes) == 0:
+        return np.zeros(0, dtype=np.int64)
+
+    kdtree = KDTree(points, leafsize=16)
+
+    suppressed = np.full(len(boxes), False)
+    for current_idx in range(len(scores)):
+        # If already visited or discarded
+        if suppressed[current_idx]:
+            continue
+
+        # If score is already below threshold then break
+        if scores[current_idx] < score_threshold:
+            break
+
+        # Query the points
+        query = kdtree.query_ball_point(points[current_idx], max_dist)
+        suppressed = _suppress_bbox(
+            np.array(query), current_idx, boxes, areas, suppressed, iou_threshold
+        )
+
+        # Add the current box
+        keep.append(current_idx)
+
+    return np.array(keep)

cellseg_models_pytorch/postproc/functional/stardist.py renamed to cellseg_models_pytorch/postproc/functional/stardist/stardist.py

@@ -1,4 +1,4 @@
-"""Imported most of the stuff from stardist repo. Minor modifications.
+"""Copied the polygons to label utilities from stardist repo (with minor modifications).
 
 BSD 3-Clause License
 
@@ -34,15 +34,12 @@
 from typing import Tuple
 
 import numpy as np
-import scipy.ndimage as ndi
-from skimage import img_as_ubyte
 from skimage.draw import polygon
-from skimage.measure import regionprops
+from skimage.morphology import disk, erosion
 
-from ...utils import bounding_box, remap_label, remove_small_objects
-from .drfns import find_maxima, h_minima_reconstruction
+from .nms import get_bboxes, nms_stardist
 
-__all__ = ["post_proc_stardist", "post_proc_stardist_orig", "polygons_to_label"]
+__all__ = ["post_proc_stardist_orig", "polygons_to_label"]
 
 
 def polygons_to_label_coord(
@@ -191,42 +188,25 @@ def polygons_to_label(
     return polygons_to_label_coord(coord, shape=shape, labels=ind)
 
 
-def _clean_up(inst_map: np.ndarray, size: int = 150, **kwargs) -> np.ndarray:
-    """Clean up overlapping instances."""
-    mask = remap_label(inst_map.copy())
-    mask_connected = ndi.label(mask)[0]
-
-    labels_connected = np.unique(mask_connected)[1:]
-    for lab in labels_connected:
-        inst = np.array(mask_connected == lab, copy=True)
-        y1, y2, x1, x2 = bounding_box(inst)
-        y1 = y1 - 2 if y1 - 2 >= 0 else y1
-        x1 = x1 - 2 if x1 - 2 >= 0 else x1
-        x2 = x2 + 2 if x2 + 2 <= mask_connected.shape[1] - 1 else x2
-        y2 = y2 + 2 if y2 + 2 <= mask_connected.shape[0] - 1 else y2
-
-        box_insts = mask[y1:y2, x1:x2]
-        if len(np.unique(ndi.label(box_insts)[0])) <= 2:
-            real_labels, counts = np.unique(box_insts, return_counts=True)
-            real_labels = real_labels[1:]
-            counts = counts[1:]
-            max_pixels = np.max(counts)
-            max_label = real_labels[np.argmax(counts)]
-            for real_lab, count in list(zip(list(real_labels), list(counts))):
-                if count < max_pixels:
-                    if count < size:
-                        mask[mask == real_lab] = max_label
-
-    return mask
-
-
 def post_proc_stardist(
-    dist_map: np.ndarray, stardist_map: np.ndarray, thresh: float = 0.4, **kwargs
+    dist_map: np.ndarray,
+    stardist_map: np.ndarray,
+    score_thresh: float = 0.5,
+    iou_thresh: float = 0.5,
+    trim_bboxes: bool = True,
+    **kwargs,
 ) -> np.ndarray:
-    """Run post-processing for stardist.
+    """Run post-processing for stardist outputs.
+
+    NOTE: This is not the original cpp version.
+    This is a python re-implementation of the stardidst post-processing
+    pipeline that uses non-maximum-suppression. Here, critical parts of the
+    nms are accelerated with `numba` and `scipy.spatial.KDtree`.
 
-    NOTE: This is not the original version that uses NMS.
-    This is rather a workaround that is a little slower.
+    NOTE:
+    This implementaiton of the stardist post-processing is actually nearly twice
+    faster than the original version if `trim_bboxes` is set to True. The resulting
+    segmentation is not an exact match but the differences are mostly neglible.
 
     Parameters
     ----------
@@ -236,37 +216,75 @@ def post_proc_stardist(
             Predicted radial distances. Shape: (n_rays, H, W).
         thresh : float, default=0.4
             Threshold for the regressed distance transform.
+        trim_bboxes : bool, default=True
+            If True, The non-zero pixels are computed only from the cell contours
+            which prunes down the pixel search space drastically.
 
     Returns
     -------
         np.ndarray:
             Instance labelled mask. Shape: (H, W).
     """
-    stardist_map = stardist_map.transpose(1, 2, 0)
-    mask = _ind_prob_thresh(dist_map, thresh, b=2)
-
-    # invert distmap
-    inv_dist_map = 255 - img_as_ubyte(dist_map)
-
-    # find markers from minima erosion reconstructed maxima of inv dist map
-    reconstructed = h_minima_reconstruction(inv_dist_map)
-    markers = find_maxima(reconstructed, mask=mask)
-    markers = ndi.label(markers)[0]
-    markers = remove_small_objects(markers, min_size=5)
-    points = np.array(
-        tuple(np.array(r.centroid).astype(int) for r in regionprops(markers))
-    )
+    if (
+        not dist_map.ndim == 2
+        and not stardist_map.ndim == 3
+        and not dist_map.shape == stardist_map.shape[:2]
+    ):
+        raise ValueError(
+            "Illegal input shapes. Make sure that: "
+            f"`dist_map` has to have shape: (H, W). Got: {dist_map.shape} "
+            f"`stardist_map` has to have shape (H, W, nrays). Got: {stardist_map.shape}"
+        )
 
-    if len(points) == 0:
-        return np.zeros_like(mask)
+    dist = np.asarray(stardist_map).transpose(1, 2, 0)
+    prob = np.asarray(dist_map)
 
-    dist = stardist_map[tuple(points.T)]
-    scores = dist_map[tuple(points.T)]
+    # threshold the edt distance transform map
+    mask = _ind_prob_thresh(prob, score_thresh)
 
-    labels = polygons_to_label(
-        dist, points, prob=scores, shape=mask.shape, scale_dist=(1, 1)
+    # get only the mask contours to trim down bbox search space
+    if trim_bboxes:
+        fp = disk(2)
+        mask -= erosion(mask, fp)
+
+    points = np.stack(np.where(mask), axis=1)
+
+    # Get only non-zero pixels of the transforms
+    dist = dist[mask > 0]
+    scores = prob[mask > 0]
+
+    # sort descendingly
+    ind = np.argsort(scores)[::-1]
+    dist = dist[ind]
+    scores = scores[ind]
+    points = points[ind]
+
+    # get bounding boxes
+    x1, y1, x2, y2, areas, max_dist = get_bboxes(dist, points)
+    boxes = np.stack([x1, y1, x2, y2], axis=1)
+
+    # consider only boxes above score threshold
+    score_cond = scores >= score_thresh
+    boxes = boxes[score_cond]
+    scores = scores[score_cond]
+    areas = areas[score_cond]
+
+    # run nms
+    inds = nms_stardist(
+        boxes,
+        points,
+        scores,
+        areas,
+        max_dist,
+        score_threshold=score_thresh,
+        iou_threshold=iou_thresh,
     )
-    labels = _clean_up(labels, **kwargs)
+
+    # get the centroids
+    points = points[inds]
+    scores = scores[inds]
+    dist = dist[inds]
+    labels = polygons_to_label(dist, points, prob=scores, shape=dist_map.shape)
 
     return labels
 

cellseg_models_pytorch/utils/__init__.py

@@ -17,6 +17,7 @@
     get_inst_centroid,
     get_inst_types,
     get_type_instances,
+    intersection,
     label_semantic,
     majority_vote_parallel,
     majority_vote_sequential,
@@ -132,4 +133,5 @@
     "majority_vote_parallel",
     "med_filt_parallel",
     "med_filt_sequential",
+    "intersection",
 ]