VisionLanguageGroup/MicroscopyMatching

1"""2Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.3"""4import logging5import io6from tqdm import tqdm, trange7import cv28from scipy.ndimage import find_objects9import numpy as np10import fastremap11import fill_voids12from models.seg_post_model import metrics13 14 15class TqdmToLogger(io.StringIO):16    """17        Output stream for TQDM which will output to logger module instead of18        the StdOut.19    """20    logger = None21    level = None22    buf = ""23 24    def __init__(self, logger, level=None):25        super(TqdmToLogger, self).__init__()26        self.logger = logger27        self.level = level or logging.INFO28 29    def write(self, buf):30        self.buf = buf.strip("\r\n\t ")31 32    def flush(self):33        self.logger.log(self.level, self.buf)34 35 36 37# def masks_to_outlines(masks):38#     """Get outlines of masks as a 0-1 array.39 40#     Args:41#         masks (int, 2D or 3D array): Size [Ly x Lx] or [Lz x Ly x Lx], where 0=NO masks and 1,2,...=mask labels.42 43#     Returns:44#         outlines (2D or 3D array): Size [Ly x Lx] or [Lz x Ly x Lx], where True pixels are outlines.45#     """46#     if masks.ndim > 3 or masks.ndim < 2:47#         raise ValueError("masks_to_outlines takes 2D or 3D array, not %dD array" %48#                          masks.ndim)49#     outlines = np.zeros(masks.shape, bool)50 51#     if masks.ndim == 3:52#         for i in range(masks.shape[0]):53#             outlines[i] = masks_to_outlines(masks[i])54#         return outlines55#     else:56#         slices = find_objects(masks.astype(int))57#         for i, si in enumerate(slices):58#             if si is not None:59#                 sr, sc = si60#                 mask = (masks[sr, sc] == (i + 1)).astype(np.uint8)61#                 contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,62#                                             cv2.CHAIN_APPROX_NONE)63#                 pvc, pvr = np.concatenate(contours[-2], axis=0).squeeze().T64#                 vr, vc = pvr + sr.start, pvc + sc.start65#                 outlines[vr, vc] = 166#         return outlines67 68 69def stitch3D(masks, stitch_threshold=0.25):70    """71    Stitch 2D masks into a 3D volume using a stitch_threshold on IOU.72 73    Args:74        masks (list or ndarray): List of 2D masks.75        stitch_threshold (float, optional): Threshold value for stitching. Defaults to 0.25.76 77    Returns:78        list: List of stitched 3D masks.79    """80    mmax = masks[0].max()81    empty = 082    for i in trange(len(masks) - 1):83        iou = metrics._intersection_over_union(masks[i + 1], masks[i])[1:, 1:]84        if not iou.size and empty == 0:85            masks[i + 1] = masks[i + 1]86            mmax = masks[i + 1].max()87        elif not iou.size and not empty == 0:88            icount = masks[i + 1].max()89            istitch = np.arange(mmax + 1, mmax + icount + 1, 1, masks.dtype)90            mmax += icount91            istitch = np.append(np.array(0), istitch)92            masks[i + 1] = istitch[masks[i + 1]]93        else:94            iou[iou < stitch_threshold] = 0.095            iou[iou < iou.max(axis=0)] = 0.096            istitch = iou.argmax(axis=1) + 197            ino = np.nonzero(iou.max(axis=1) == 0.0)[0]98            istitch[ino] = np.arange(mmax + 1, mmax + len(ino) + 1, 1, masks.dtype)99            mmax += len(ino)100            istitch = np.append(np.array(0), istitch)101            masks[i + 1] = istitch[masks[i + 1]]102            empty = 1103 104    return masks105 106 107# def diameters(masks):108#     """109#     Calculate the diameters of the objects in the given masks.110 111#     Parameters:112#     masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)113 114#     Returns:115#         tuple: A tuple containing the median diameter and an array of diameters for each object.116 117#     Examples:118#     >>> masks = np.array([[0, 1, 1], [1, 0, 0], [1, 1, 0]])119#     >>> diameters(masks)120#     (1.0, array([1.41421356, 1.0, 1.0]))121#     """122#     uniq, counts = fastremap.unique(masks.astype("int32"), return_counts=True)123#     counts = counts[1:]124#     md = np.median(counts**0.5)125#     if np.isnan(md):126#         md = 0127#     md /= (np.pi**0.5) / 2128#     return md, counts**0.5129 130 131# def radius_distribution(masks, bins):132#     """133#     Calculate the radius distribution of masks.134 135#     Args:136#         masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)137#         bins (int): Number of bins for the histogram.138 139#     Returns:140#         A tuple containing a normalized histogram of radii, median radius, array of radii.141 142#     """143#     unique, counts = np.unique(masks, return_counts=True)144#     counts = counts[unique != 0]145#     nb, _ = np.histogram((counts**0.5) * 0.5, bins)146#     nb = nb.astype(np.float32)147#     if nb.sum() > 0:148#         nb = nb / nb.sum()149#     md = np.median(counts**0.5) * 0.5150#     if np.isnan(md):151#         md = 0152#     md /= (np.pi**0.5) / 2153#     return nb, md, (counts**0.5) / 2154 155 156# def size_distribution(masks):157#     """158#     Calculates the size distribution of masks.159 160#     Args:161#         masks (ndarray): masks (0=no cells, 1=first cell, 2=second cell,...)162 163#     Returns:164#         float: The ratio of the 25th percentile of mask sizes to the 75th percentile of mask sizes.165#     """166#     counts = np.unique(masks, return_counts=True)[1][1:]167#     return np.percentile(counts, 25) / np.percentile(counts, 75)168 169 170def fill_holes_and_remove_small_masks(masks, min_size=15):171    """ Fills holes in masks (2D/3D) and discards masks smaller than min_size.172 173    This function fills holes in each mask using fill_voids.fill.174    It also removes masks that are smaller than the specified min_size.175 176    Parameters:177    masks (ndarray): Int, 2D or 3D array of labelled masks.178        0 represents no mask, while positive integers represent mask labels.179        The size can be [Ly x Lx] or [Lz x Ly x Lx].180    min_size (int, optional): Minimum number of pixels per mask.181        Masks smaller than min_size will be removed.182        Set to -1 to turn off this functionality. Default is 15.183 184    Returns:185        ndarray: Int, 2D or 3D array of masks with holes filled and small masks removed.186            0 represents no mask, while positive integers represent mask labels.187            The size is [Ly x Lx] or [Lz x Ly x Lx].188    """189 190    if masks.ndim > 3 or masks.ndim < 2:191        raise ValueError("masks_to_outlines takes 2D or 3D array, not %dD array" %192                         masks.ndim)193 194    # Filter small masks195    if min_size > 0:196        counts = fastremap.unique(masks, return_counts=True)[1][1:]197        masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)198        fastremap.renumber(masks, in_place=True)199        200    slices = find_objects(masks)201    j = 0202    for i, slc in enumerate(slices):203        if slc is not None:204            msk = masks[slc] == (i + 1)205            msk = fill_voids.fill(msk)206            masks[slc][msk] = (j + 1)207            j += 1208 209    if min_size > 0:210        counts = fastremap.unique(masks, return_counts=True)[1][1:]211        masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)212        fastremap.renumber(masks, in_place=True)213    214    return masks215