Implement splitting of non-convex objects

Author: constantinpape

flamingo_tools/segmentation/postprocessing.py

@@ -1,7 +1,8 @@
 import math
 import multiprocessing as mp
+import threading
 from concurrent import futures
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import elf.parallel as parallel
 import numpy as np
@@ -15,6 +16,9 @@
 from scipy.spatial import distance
 from scipy.spatial import cKDTree, ConvexHull
 from skimage import measure
+from skimage.filters import gaussian
+from skimage.feature import peak_local_max
+from skimage.segmentation import find_boundaries, watershed
 from sklearn.neighbors import NearestNeighbors
 from tqdm import tqdm
 
@@ -734,3 +738,116 @@ def filter_cochlea_volume(
         combined_dilated[combined_dilated > 0] = 1
 
     return combined_dilated
+
+
+def split_nonconvex_objects(
+    segmentation: np.typing.ArrayLike,
+    output: np.typing.ArrayLike,
+    segmentation_table: pd.DataFrame,
+    min_size: int,
+    resolution: Union[float, Sequence[float]],
+    height_map: Optional[np.typing.ArrayLike] = None,
+    component_labels: Optional[List[int]] = None,
+    n_threads: Optional[int] = None,
+) -> Dict[int, List[int]]:
+    """Split noncovex objects into multiple parts inplace.
+
+    Args:
+        segmentation:
+        output:
+        segmentation_table:
+        min_size:
+        resolution:
+        height_map:
+        component_labels:
+        n_threads:
+    """
+    if isinstance(resolution, float):
+        resolution = [resolution] * 3
+    assert len(resolution) == 3
+    resolution = np.array(resolution)
+
+    lock = threading.Lock()
+    offset = len(segmentation_table)
+
+    def split_object(object_id):
+        nonlocal offset
+
+        row = segmentation_table[segmentation_table.label_id == object_id]
+        if min_size and row.n_pixels.values[0] < min_size:
+            return [object_id]
+
+        bb_min = np.array([
+            row.bb_min_z.values[0], row.bb_min_y.values[0], row.bb_min_x.values[0],
+        ]) / resolution
+        bb_max = np.array([
+            row.bb_max_z.values[0], row.bb_max_y.values[0], row.bb_max_x.values[0],
+        ]) / resolution
+
+        bb_min = np.maximum(bb_min.astype(int) - 1, np.array([0, 0, 0]))
+        bb_max = np.minimum(bb_max.astype(int) + 1, np.array(list(segmentation.shape)))
+        bb = tuple(slice(mi, ma) for mi, ma in zip(bb_min, bb_max))
+
+        seg = segmentation[bb]
+        mask = ~find_boundaries(seg)
+        dist = distance_transform_edt(mask, sampling=resolution)
+
+        seg_mask = seg == object_id
+        dist[~seg_mask] = 0
+        dist = gaussian(dist, (0.6, 1.2, 1.2))
+        maxima = peak_local_max(dist, min_distance=3, exclude_border=True)
+
+        if len(maxima) == 1:
+            return [object_id]
+
+        with lock:
+            old_offset = offset
+            offset += len(maxima)
+
+        seeds = np.zeros(seg.shape, dtype=int)
+        for i, pos in enumerate(maxima, 1):
+            seeds[tuple(pos)] = old_offset + i
+
+        if height_map is None:
+            hmap = dist.max() - dist
+        else:
+            hmap = height_map[bb]
+        new_seg = watershed(hmap, markers=seeds, mask=seg_mask)
+
+        seg_ids, sizes = np.unique(new_seg, return_counts=True)
+        seg_ids, sizes = seg_ids[1:], sizes[1:]
+
+        keep_ids = seg_ids[sizes > min_size]
+        if len(keep_ids) < 2:
+            return [object_id]
+
+        elif len(keep_ids) != len(seg_ids):
+            new_seg[~np.isin(new_seg, keep_ids)] = 0
+            new_seg = watershed(hmap, markers=new_seg, mask=seg_mask)
+
+        output[bb][seg_mask] = new_seg[seg_mask]
+        return seg_ids.tolist()
+
+        # import napari
+        # v = napari.Viewer()
+        # v.add_image(hmap)
+        # v.add_labels(seg)
+        # v.add_labels(new_seg)
+        # v.add_points(maxima)
+        # napari.run()
+
+    if component_labels is None:
+        object_ids = segmentation_table.label_id.values
+    else:
+        object_ids = segmentation_table[segmentation_table.isin(component_labels)].label_id.values
+
+    if n_threads is None:
+        n_threads = mp.cpu_count()
+
+    with futures.ThreadPoolExecutor(n_threads) as tp:
+        new_id_mapping = list(
+            tqdm(tp.map(split_object, object_ids), total=len(object_ids), desc="Split non-convex objects")
+        )
+
+    new_id_mapping = {object_id: mapped_ids for object_id, mapped_ids in zip(object_ids, new_id_mapping)}
+    return new_id_mapping

scripts/la-vision/fix_segmentation.py

@@ -0,0 +1,129 @@
+import os
+from glob import glob
+
+import napari
+import numpy as np
+import imageio.v3 as imageio
+import vigra
+
+from skimage.filters import gaussian
+from skimage.segmentation import find_boundaries, watershed
+from scipy.ndimage import distance_transform_edt
+from skimage.feature import peak_local_max
+from skimage.measure import regionprops, label
+
+
+def _size_filter(segmentation, heightmap, min_size):
+    ids, sizes = np.unique(segmentation, return_counts=True)
+    discard_ids = ids[sizes < min_size]
+    mask = segmentation > 0
+    segmentation[np.isin(segmentation, discard_ids)] = 0
+    return watershed(heightmap, markers=segmentation, mask=mask)
+
+
+def postproc(image, segmentation, view=False):
+    # First get rid of small objects.
+    min_size = 250
+    heightmap = vigra.filters.laplacianOfGaussian(image.astype("float32"), 3)
+
+    segmentation = _size_filter(segmentation, heightmap, min_size)
+
+    mask = ~find_boundaries(segmentation)
+    dist = distance_transform_edt(mask, sampling=(2, 1, 1))
+    dist[segmentation == 0] = 0
+    dist = gaussian(dist, (0.6, 1.2, 1.2))
+    maxima = peak_local_max(dist, min_distance=3, exclude_border=False)
+
+    maxima_image = np.zeros(segmentation.shape, dtype="uint8")
+    pos = tuple(maxima[:, i] for i in range(3))
+    maxima_image[pos] = 1
+    maxima_image = label(maxima_image)
+
+    def maxima_ids(seg, im):
+        ids = np.unique(im[seg])
+        return ids[1:]
+
+    seed_maxima_ids, keep_seg_ids, split_seg_ids = [], [], []
+    props = regionprops(segmentation, maxima_image, extra_properties=[maxima_ids])
+    for prop in props:
+        this_maxima_ids = prop.maxima_ids
+        if len(this_maxima_ids) == 1:
+            keep_seg_ids.append(prop.label)
+            continue
+        seed_maxima_ids.extend(this_maxima_ids.tolist())
+        split_seg_ids.append(prop.label)
+
+    split_mask = np.isin(segmentation, split_seg_ids)
+    # segmentation[split_mask] = 0
+
+    new_seeds = maxima_image.copy()
+    new_seeds[~np.isin(maxima_image, seed_maxima_ids)] = 0
+    new_seg = watershed(heightmap, markers=new_seeds, mask=split_mask)
+
+    segmentation[split_mask] = 0
+    offset = segmentation.max()
+    new_seg[new_seg != 0] += offset
+    segmentation[split_mask] = new_seg[split_mask]
+    segmentation = label(segmentation)
+    segmentation = _size_filter(segmentation, heightmap, min_size)
+
+    if view:
+        v = napari.Viewer()
+        v.add_image(image)
+        v.add_labels(segmentation)
+        # v.add_labels(new_seg)
+        # v.add_image(heightmap)
+        # v.add_image(dist)
+        # v.add_points(maxima)
+        # v.add_labels(split_mask)
+        napari.run()
+
+    return segmentation
+
+
+def postprocess_volume(im_path, seg_path, out_root):
+    image = imageio.imread(im_path)
+    segmentation = imageio.imread(seg_path)
+    segmentation = postproc(image, segmentation, view=True)
+
+    os.makedirs(out_root, exist_ok=True)
+    fname = os.path.basename(im_path)
+    imageio.imwrite(os.path.join(out_root, fname), segmentation, compression="zlib")
+
+
+def postprocess_volume_scalable(im_path, seg_path, out_root):
+    from flamingo_tools.segmentation.postprocessing import split_nonconvex_objects, compute_table_on_the_fly
+
+    image = imageio.imread(im_path)
+    segmentation = imageio.imread(seg_path)
+
+    # TODO aniso resolution
+    resolution = 0.38
+    table = compute_table_on_the_fly(segmentation, resolution)
+
+    out = np.zeros_like(segmentation)
+    id_mapping = split_nonconvex_objects(segmentation, out, table, n_threads=1, resolution=resolution, min_size=250)
+    n_prev = len(id_mapping)
+    n_after = sum([len(v) for v in id_mapping.values()])
+    print("Before splitting:", n_prev)
+    print("After splitting:", n_after)
+
+    v = napari.Viewer()
+    v.add_image(image)
+    v.add_labels(segmentation, visible=False)
+    v.add_labels(out)
+    napari.run()
+
+
+def main():
+    im_paths = sorted(glob("la-vision-sgn-new/images/*.tif"))
+    seg_paths = sorted(glob("la-vision-sgn-new/segmentation/*.tif"))
+    out_root = "la-vision-sgn-new/segmentation-postprocessed"
+    for im_path, seg_path in zip(im_paths, seg_paths):
+        # postprocess_volume(im_path, seg_path, out_root)
+        postprocess_volume_scalable(im_path, seg_path, out_root)
+        break
+
+
+if __name__ == "__main__":
+    main()