Include marker extension in SGN detection

Author: schilling40

flamingo_tools/segmentation/marker_extension.py

@@ -1,29 +1,119 @@
-import multiprocessing as mp
+import multiprocessing
+import os
 import threading
 from concurrent import futures
-import os
-from typing import Optional, Tuple
+from threadpoolctl import threadpool_limits
+from typing import Optional, Tuple, Union
 
 import numpy as np
+from numpy.typing import ArrayLike
 import pandas as pd
+from scipy.ndimage import distance_transform_edt
+from skimage.segmentation import watershed
 import zarr
 
 from elf.io import open_file
 from elf.parallel.local_maxima import find_local_maxima
 from flamingo_tools.segmentation.unet_prediction import prediction_impl
 from tqdm import tqdm
 
+from elf.parallel.common import get_blocking
+
+
+def distance_based_marker_extension(
+    markers: np.ndarray,
+    output: ArrayLike,
+    extension_distance: float,
+    sampling: Union[float, Tuple[float, ...]],
+    block_shape: Tuple[int, ...],
+    n_threads: Optional[int] = None,
+    verbose: bool = False,
+    roi: Optional[Tuple[slice, ...]] = None,
+):
+    """
+    Extend SGN detection to emulate shape of SGNs for better visualization.
+
+    Args:
+        markers: Array of coordinates for seeding watershed.
+        output: Output for watershed.
+        extension_distance: Distance in micrometer for extension.
+        sampling: Resolution in micrometer.
+        block_shape:
+        n_threads:
+        verbose:
+        roi:
+    """
+    n_threads = multiprocessing.cpu_count() if n_threads is None else n_threads
+    blocking = get_blocking(output, block_shape, roi, n_threads)
+
+    lock = threading.Lock()
+
+    # determine the correct halo in pixels based on the sampling and the extension distance.
+    halo = [round(extension_distance / s) + 2 for s in sampling]
+
+    @threadpool_limits.wrap(limits=1)  # restrict the numpy threadpool to 1 to avoid oversubscription
+    def extend_block(block_id):
+        block = blocking.getBlockWithHalo(block_id, halo)
+        outer_block = block.outerBlock
+        inner_block = block.innerBlock
+
+        # TODO get the indices and coordinates of the markers in the INNER block
+        # markers_in_block_ids = [int(i) for i in np.unique(inner_block)[1:]]
+        mask = (
+            (inner_block.begin[0] <= markers[:, 0]) & (markers[:, 0] <= inner_block.end[0]) &
+            (inner_block.begin[1] <= markers[:, 1]) & (markers[:, 1] <= inner_block.end[1]) &
+            (inner_block.begin[2] <= markers[:, 2]) & (markers[:, 2] <= inner_block.end[2])
+        )
+        markers_in_block_ids = np.where(mask)[0]
+        markers_in_block_coords = markers[markers_in_block_ids]
+
+        # TODO offset the marker coordinates with respect to the OUTER block
+        markers_in_block_coords = [coord - outer_block.begin for coord in markers_in_block_coords]
+        markers_in_block_coords = [[round(c) for c in coord] for coord in markers_in_block_coords]
+        markers_in_block_coords = np.array(markers_in_block_coords, dtype=int)
+        z, y, x = markers_in_block_coords.T
+
+        # Shift index by one so that zero is reserved for background id
+        markers_in_block_ids += 1
+
+        # Create the seed volume.
+        outer_block_shape = tuple(end - begin for begin, end in zip(outer_block.begin, outer_block.end))
+        seeds = np.zeros(outer_block_shape, dtype="uint32")
+        seeds[z, y, x] = markers_in_block_ids
+
+        # Compute the distance map.
+        distance = distance_transform_edt(seeds == 0, sampling=sampling)
+
+        # And extend the seeds
+        mask = distance < extension_distance
+        segmentation = watershed(distance.max() - distance, markers=seeds, mask=mask)
+
+        # Write the segmentation. Note: we need to lock here because we write outside of our inner block
+        bb = tuple(slice(begin, end) for begin, end in zip(outer_block.begin, outer_block.end))
+        with lock:
+            this_output = output[bb]
+            this_output[mask] = segmentation[mask]
+            output[bb] = this_output
+
+    n_blocks = blocking.numberOfBlocks
+    with futures.ThreadPoolExecutor(n_threads) as tp:
+        list(tqdm(
+            tp.map(extend_block, range(n_blocks)), total=n_blocks, desc="Marker extension", disable=not verbose
+        ))
+
 
 def sgn_detection(
     input_path: str,
     input_key: str,
     output_folder: str,
     model_path: str,
+    extension_distance: float,
+    sampling: Union[float, Tuple[float, ...]],
     block_shape: Optional[Tuple[int, int, int]] = None,
     halo: Optional[Tuple[int, int, int]] = None,
-    spot_radius: int = 2,
+    n_threads: Optional[int] = None,
 ):
-    """Run prediction for sgn detection.
+    """Run prediction for SGN detection.
 
     Args:
         input_path: Input path to image channel for SGN detection.
@@ -70,22 +160,14 @@ def sgn_detection(
             chunks=chunks, compression="gzip"
         )
 
-        lock = threading.Lock()
-
-        def add_halo_segm(detection_index):
-            """Create a segmentation volume around all detected spots.
-            """
-            coord = list(detections[detection_index])
-            block_begin = [round(c) - spot_radius for c in coord]
-            block_end = [round(c) + spot_radius for c in coord]
-            volume_index = tuple(slice(beg, end) for beg, end in zip(block_begin, block_end))
-            with lock:
-                output_dataset[volume_index] = int(detection_index) + 1
-
-        # Limit the number of cores for parallelization.
-        n_threads = min(16, mp.cpu_count())
-        with futures.ThreadPoolExecutor(n_threads) as filter_pool:
-            list(tqdm(filter_pool.map(add_halo_segm, range(len(detections))), total=len(detections)))
+        distance_based_marker_extension(
+            markers=detections,
+            output=output_dataset,
+            extension_distance=extension_distance,
+            sampling=sampling,
+            block_shape=(128, 128, 128),
+            n_threads=n_threads,
+        )
 
         # Save the result in mobie compatible format.
         detections = np.concatenate(

flamingo_tools/segmentation/sgn_detection.py

@@ -14,6 +14,10 @@ def main():
     parser.add_argument("-k", "--input_key", default=None,
                         help="The key / internal path to image data.")
 
+    parser.add_argument("-d", "--extension_distance", type=float, default=12, help="Extension distance.")
+    parser.add_argument("-r", "--resolution", type=float, nargs="+", default=[3.0, 1.887779, 1.887779],
+                        help="Resolution of input in micrometer.")
+
     parser.add_argument("--s3", action="store_true", help="Use S3 bucket.")
     parser.add_argument("--s3_credentials", type=str, default=None,
                         help="Input file containing S3 credentials. "
@@ -28,6 +32,11 @@ def main():
     block_shape = (12, 128, 128)
     halo = (10, 64, 64)
 
+    if len(args.resolution) == 1:
+        resolution = tuple(args.resolution, args.resolution, args.resolution)
+    else:
+        resolution = tuple(args.resolution)
+
     if args.s3:
         input_path, fs = s3_utils.get_s3_path(args.input, bucket_name=args.s3_bucket_name,
                                               service_endpoint=args.s3_service_endpoint,
@@ -37,7 +46,8 @@ def main():
         input_path = args.input
 
     sgn_detection(input_path=input_path, input_key=args.input_key, output_folder=args.output_folder,
-                  model_path=args.model, block_shape=block_shape, halo=halo)
+                  model_path=args.model, block_shape=block_shape, halo=halo,
+                  extension_distance=args.extension_distance, sampling=resolution)
 
 
 if __name__ == "__main__":

scripts/sgn_detection/sgn_detection.py

@@ -8,7 +8,7 @@
 import scipy.ndimage as ndimage
 
 from flamingo_tools.s3_utils import get_s3_path
-from flamingo_tools.segmentation.marker_extension import distance_based_marker_extension
+from flamingo_tools.segmentation.sgn_detection import distance_based_marker_extension
 from flamingo_tools.file_utils import read_image_data