From fb9089bf914513ca61e869ff14ebe15772ab0451 Mon Sep 17 00:00:00 2001 From: ericup Date: Sat, 28 Sep 2024 16:16:13 +0200 Subject: [PATCH] Update cpn data ops --- celldetection/data/cpn.py | 240 +++++++++++++++++++++++++++++++++++--- 1 file changed, 221 insertions(+), 19 deletions(-) diff --git a/celldetection/data/cpn.py b/celldetection/data/cpn.py index 8a877b6..38a438c 100644 --- a/celldetection/data/cpn.py +++ b/celldetection/data/cpn.py @@ -1,12 +1,17 @@ +import warnings import numpy as np import cv2 import torch from skimage.measure import regionprops from skimage import img_as_ubyte from collections import OrderedDict +from scipy.ndimage import distance_transform_edt +from multiprocessing import Pool, shared_memory +from gc import collect from .segmentation import filter_instances_ from .misc import labels2properties, resample_contours -from ..util.util import asnumpy +from ..util.util import asnumpy, cpu_count +from ..visualization.cmaps import random_colors_hsv __all__ = [ 'CPNTargetGenerator', 'contours2labels', 'render_contour', 'clip_contour_', 'masks2labels', @@ -15,7 +20,7 @@ ] -def efd(contour, order=10, epsilon=1e-6): +def efd(contour, order=10, epsilon=1e-6, autoclose=True): """Elliptic fourier descriptor. Computes elliptic fourier descriptors from contour data. @@ -29,6 +34,8 @@ def efd(contour, order=10, epsilon=1e-6): order: Order of resulting descriptor. The higher the order, the more detail can be preserved. An order of 1 produces ellipses. epsilon: Epsilon value. Used to avoid division by zero. + autoclose: Whether to automatically close contours (first point equals last point in contour). + If `False` and non-closed contours are found, an AssertionError is raised. Notes: Locations may contain NaN if `contour` only contains a single point. @@ -44,6 +51,13 @@ def efd(contour, order=10, epsilon=1e-6): for i in range(len(res)): res[i].append(r_[i]) return tuple(map(np.array, res)) + + if autoclose and (not np.allclose(contour[..., 0, :], contour[..., -1, :])): + contour = np.concatenate((contour, contour[..., :1, :]), axis=-2) + else: + assert np.allclose(contour[..., 0, :], contour[..., -1, :]), ( + 'Please make sure that contours are explicitly closed (first point must be equal to last point).') + dxy = np.diff(contour, axis=-2) # shape: (..., p, d) dt = np.sqrt(np.sum(np.square(dxy), axis=-1)) + epsilon # shape: (..., p) cumsum = np.cumsum(dt, axis=-1) # shape: (..., p) @@ -630,21 +644,209 @@ def sampled_sizes(self): return self._sampled_sizes -def contours2overlay(contours, size, hue_range=(0, 180), saturation_range=(60, 133), value_range=(180, 256), - rounded=True, clip=True, intermediate_dtype='float16'): - from ..visualization.cmaps import random_colors_hsv - overlay = np.zeros(tuple(size) + (4,), dtype=intermediate_dtype) - counter = np.zeros(tuple(size), dtype=intermediate_dtype) - for idx, contour in enumerate(contours): - if rounded: - contour = np.round(contour) - if clip: - clip_contour_(contour, np.array(size) - 1) - a, (xmin, xmax), (ymin, ymax) = render_contour(contour, val=1, dtype='uint8') - c = random_colors_hsv(1, hue_range=hue_range, saturation_range=saturation_range, value_range=value_range) - counter[ymin:ymin + a.shape[0], xmin:xmin + a.shape[1]] += a - overlay[ymin:ymin + a.shape[0], xmin:xmin + a.shape[1]] += (list(c[0]) + [1.]) * a[..., None] - m = counter > 1 - overlay[m] /= counter[m][..., None] - overlay = img_as_ubyte(overlay) +def _add_contour_to_overlay_(counter, overlay, contour, size, hue_range, saturation_range, value_range, rounded, clip): + + dtype = overlay.dtype + is_int = np.issubdtype(dtype, np.integer) + if rounded: + contour = np.round(contour) + if clip: + clip_contour_(contour, np.array(size) - 1) + a, (xmin, xmax), (ymin, ymax) = render_contour(contour, val=1, dtype='uint8') + a = a.astype(dtype) + c, = random_colors_hsv(1, hue_range=hue_range, saturation_range=saturation_range, value_range=value_range, + ubyte=is_int) + c = np.array(list(c) + [255 if is_int else 1.], dtype=dtype) + counter[ymin:ymin + a.shape[0], xmin:xmin + a.shape[1]] += a + overlay[ymin:ymin + a.shape[0], xmin:xmin + a.shape[1]] += c * a[..., None] + + +def _add_contour_to_overlay_mp_(args): + """ + + Worker function for multiprocessing that adds a contour to shared overlay and counter arrays. + + Args: + args: + + Returns: + + """ + (contour, overlay_shm_name, overlay_shape, counter_shm_name, counter_shape, dtype_char, + size, hue_range, saturation_range, value_range, rounded, clip, seed) = args + + dtype = np.dtype(dtype_char) + + # Access the shared memory blocks + existing_overlay_shm = shared_memory.SharedMemory(name=overlay_shm_name) + existing_counter_shm = shared_memory.SharedMemory(name=counter_shm_name) + try: + # Reconstruct the shared arrays + overlay = np.ndarray(overlay_shape, dtype=dtype, buffer=existing_overlay_shm.buf) + counter = np.ndarray(counter_shape, dtype=dtype, buffer=existing_counter_shm.buf) + + # Add contour to overlay and counter + np.random.seed(seed) # specific seed is about 100x faster than np.random.seed(None) + _add_contour_to_overlay_(counter, overlay, contour, size, hue_range, saturation_range, value_range, + rounded, clip) + finally: + # Close shared memory (do not unlink, only the creator should unlink) + existing_overlay_shm.close() + existing_counter_shm.close() + + +def _postprocess_overlay(overlay, counter, normalize=True): + """ + + Post-process the overlay by normalizing overlapping areas and converting to uint8. + + Args: + overlay: + counter: + normalize: + + Returns: + + """ + dtype = overlay.dtype + if normalize: + m = counter > 1 + if np.issubdtype(dtype, np.integer): + overlay[m] //= counter[m][..., None] + else: + overlay[m] /= counter[m][..., None] + if np.issubdtype(dtype, np.integer): + overlay = overlay.astype('uint8') + else: + overlay = img_as_ubyte(overlay) + return overlay + + +def _contours2overlay_mp(contours, size, hue_range=(0, 360), saturation_range=(60, 133), + value_range=(180, 256), rounded=True, clip=True, + intermediate_dtype='uint16', processes=None) -> np.ndarray: + """ + + Multiprocessing version of contours to overlay. + + Args: + contours: + size: + hue_range: + saturation_range: + value_range: + rounded: + clip: + intermediate_dtype: + processes: + + Returns: + + """ + dtype = np.dtype(intermediate_dtype) + dtype_char = dtype.char + overlay_shape = tuple(size) + (4,) + counter_shape = tuple(size) + + # Create shared memory blocks + overlay_size = np.prod(overlay_shape) * dtype.itemsize + counter_size = np.prod(counter_shape) * dtype.itemsize + overlay_shm = shared_memory.SharedMemory(create=True, size=overlay_size) + counter_shm = shared_memory.SharedMemory(create=True, size=counter_size) + + # Create numpy arrays backed by shared memory + overlay = np.ndarray(overlay_shape, dtype=dtype, buffer=overlay_shm.buf) + counter = np.ndarray(counter_shape, dtype=dtype, buffer=counter_shm.buf) + + # Initialize arrays + overlay.fill(0) + counter.fill(0) + + try: + if processes is None or processes <= 0: + processes = min(len(contours), cpu_count() or 1) + + # Prepare arguments for each contour + args = [ + ( + contour, + overlay_shm.name, + overlay_shape, + counter_shm.name, + counter_shape, + dtype_char, + size, + hue_range, + saturation_range, + value_range, + rounded, + clip, + seed # multiprocessing requires seeds + ) + for contour, seed in zip( + contours, + np.random.randint(0, np.iinfo(np.uint32).max, size=len(contours), dtype=np.uint32) + ) + ] + + # Process contours in parallel + with Pool(processes) as pool: + pool.map(_add_contour_to_overlay_mp_, args) + finally: + # Post-process the overlay before closing shared memory + overlay = _postprocess_overlay(overlay.copy(), counter.copy(), normalize=contours is not None) + + # Close and unlink shared memory + overlay_shm.close() + overlay_shm.unlink() + counter_shm.close() + counter_shm.unlink() + collect() + return overlay + + +def contours2overlay(contours, size, hue_range=(0, 180), saturation_range=(60, 133), value_range=(180, 256), + rounded=True, clip=True, intermediate_dtype='uint16', processes=None): + """Contours to overlay. + + Creates an overlay image from contours. + Overlay images are RGBA images that can be used to visualize the contours with a transparent background. + To visualize overlapping contour areas, colors in those regions are interpolated. + + Notes: + - 1min 44s for approx. 5.6M contours, size of (63348, 50638), uint16 as dtype and 128 processes. + + Args: + contours: Contours as Array[num_contours, num_points, 2] or List[Array[num_points, 2]]. + size: Image size. + hue_range: Hue range for HSV colors. + saturation_range: Saturation range for HSV colors. + value_range: Value range for HSV colors. + rounded: Whether to round contour values. + clip: Whether to clip contours to `size`. + intermediate_dtype: Intermediate dtype. Small dtypes reduce memory consumption and speed up processing. + The dtype also limits the maximum number of overlaps. The sum of all overlapping + colors, as well as the number of overlaps, must remain below the dtypes upper limit. + processes: Number of processes for multiprocessing. `None` disables multiprocessing. + Negative values cause number of processes to be inferred from system CPU count and + environment (e.g. MPI, Slurm). + + Returns: + Overlay image. + """ + if processes is not None and processes < 0: + processes = min(len(contours), cpu_count(adjust=True)) + + if processes is not None and processes not in (1, 0): + return _contours2overlay_mp(contours, size, hue_range, saturation_range, value_range, rounded, clip, + intermediate_dtype, processes) + else: + overlay = np.zeros(tuple(size) + (4,), dtype=intermediate_dtype) + counter = np.zeros(tuple(size), dtype=intermediate_dtype) + if contours is not None: + for contour in contours: + _add_contour_to_overlay_(counter, overlay, contour, size, hue_range, saturation_range, value_range, + rounded, clip) + overlay = _postprocess_overlay(overlay, counter, normalize=contours is not None) + return overlay