VisionLanguageGroup/MicroscopyMatching

1"""Regionprops features and its augmentations.2WindowedRegionFeatures (WRFeatures) is a class that holds regionprops features for a windowed track region.3Modified from Trackastra (https://github.com/weigertlab/trackastra)4"""5 6import itertools7import logging8from collections import OrderedDict9from collections.abc import Iterable #, Sequence10from functools import reduce11from typing import Literal12 13import joblib14import numpy as np15import pandas as pd16from edt import edt17from skimage.measure import regionprops, regionprops_table18from tqdm import tqdm19from typing import Tuple, Optional, Sequence, Union, List20import typing21import torch22logger = logging.getLogger(__name__)23 24_PROPERTIES = {25    "regionprops": (26        "area",27        "intensity_mean",28        "intensity_max",29        "intensity_min",30        "inertia_tensor",31    ),32    "regionprops2": (33        "equivalent_diameter_area",34        "intensity_mean",35        "inertia_tensor",36        "border_dist",37    ),38}39 40 41def _border_dist_fast(mask: np.ndarray, cutoff: float = 5):42    cutoff = int(cutoff)43    border = np.ones(mask.shape, dtype=np.float32)44    ndim = len(mask.shape)45 46    for axis, size in enumerate(mask.shape):47        # Create fade values for the band [0, cutoff)48        band_vals = np.arange(cutoff, dtype=np.float32) / cutoff49 50        # Build slices for the low border51        low_slices = [slice(None)] * ndim52        low_slices[axis] = slice(0, cutoff)53        border_low = border[tuple(low_slices)]54        border_low_vals = np.minimum(55            border_low, band_vals[(...,) + (None,) * (ndim - axis - 1)]56        )57        border[tuple(low_slices)] = border_low_vals58 59        # Build slices for the high border60        high_slices = [slice(None)] * ndim61        high_slices[axis] = slice(size - cutoff, size)62        band_vals_rev = band_vals[::-1]63        border_high = border[tuple(high_slices)]64        border_high_vals = np.minimum(65            border_high, band_vals_rev[(...,) + (None,) * (ndim - axis - 1)]66        )67        border[tuple(high_slices)] = border_high_vals68 69    dist = 1 - border70    return tuple(r.intensity_max for r in regionprops(mask, intensity_image=dist))71 72 73class WRFeatures:74    """regionprops features for a windowed track region."""75 76    def __init__(77        self,78        coords: np.ndarray,79        labels: np.ndarray,80        timepoints: np.ndarray,81        features: typing.OrderedDict[str, np.ndarray],82    ):83        self.ndim = coords.shape[-1]84        if self.ndim not in (2, 3):85            raise ValueError("Only 2D or 3D data is supported")86 87        self.coords = coords88        self.labels = labels89        self.features = features.copy()90        self.timepoints = timepoints91 92    def __repr__(self):93        s = (94            f"WindowRegionFeatures(ndim={self.ndim}, nregions={len(self.labels)},"95            f" ntimepoints={len(np.unique(self.timepoints))})\n\n"96        )97        for k, v in self.features.items():98            s += f"{k:>20} -> {v.shape}\n"99        return s100 101    @property102    def features_stacked(self):103        return np.concatenate([v for k, v in self.features.items()], axis=-1)104 105    def __len__(self):106        return len(self.labels)107 108    def __getitem__(self, key):109        if key in self.features:110            return self.features[key]111        else:112            raise KeyError(f"Key {key} not found in features")113 114    @classmethod115    def concat(cls, feats: Sequence["WRFeatures"]) -> "WRFeatures":116        """Concatenate multiple WRFeatures into a single one."""117        if len(feats) == 0:118            raise ValueError("Cannot concatenate empty list of features")119        return reduce(lambda x, y: x + y, feats)120 121    def __add__(self, other: "WRFeatures") -> "WRFeatures":122        """Concatenate two WRFeatures."""123        if self.ndim != other.ndim:124            raise ValueError("Cannot concatenate features of different dimensions")125        if self.features.keys() != other.features.keys():126            raise ValueError("Cannot concatenate features with different properties")127 128        coords = np.concatenate([self.coords, other.coords], axis=0)129        labels = np.concatenate([self.labels, other.labels], axis=0)130        timepoints = np.concatenate([self.timepoints, other.timepoints], axis=0)131 132        features = OrderedDict(133            (k, np.concatenate([v, other.features[k]], axis=0))134            for k, v in self.features.items()135        )136 137        return WRFeatures(138            coords=coords, labels=labels, timepoints=timepoints, features=features139        )140 141    @classmethod142    def from_mask_img(143        cls,144        mask: np.ndarray,145        img: np.ndarray,146        properties="regionprops2",147        t_start: int = 0,148    ):149        img = np.asarray(img)150        mask = np.asarray(mask)151 152        _ntime, ndim = mask.shape[0], mask.ndim - 1153        if ndim not in (2, 3):154            raise ValueError("Only 2D or 3D data is supported")155 156        properties = tuple(_PROPERTIES[properties])157        if "label" in properties or "centroid" in properties:158            raise ValueError(159                f"label and centroid should not be in properties {properties}"160            )161 162        if "border_dist" in properties:163            use_border_dist = True164            # remove border_dist from properties165            properties = tuple(p for p in properties if p != "border_dist")166        else:167            use_border_dist = False168 169        df_properties = ("label", "centroid", *properties)170        dfs = []171        for i, (y, x) in enumerate(zip(mask, img)):172            _df = pd.DataFrame(173                regionprops_table(y, intensity_image=x, properties=df_properties)174            )175            _df["timepoint"] = i + t_start176 177            if use_border_dist:178                _df["border_dist"] = _border_dist_fast(y)179 180            dfs.append(_df)181        df = pd.concat(dfs)182 183        if use_border_dist:184            properties = (*properties, "border_dist")185 186        timepoints = df["timepoint"].values.astype(np.int32)187        labels = df["label"].values.astype(np.int32)188        coords = df[[f"centroid-{i}" for i in range(ndim)]].values.astype(np.float32)189 190        features = OrderedDict(191            (192                p,193                np.stack(194                    [195                        df[c].values.astype(np.float32)196                        for c in df.columns197                        if c.startswith(p)198                    ],199                    axis=-1,200                ),201            )202            for p in properties203        )204 205        return cls(206            coords=coords, labels=labels, timepoints=timepoints, features=features207        )208 209 210def get_features(211    detections: np.ndarray,212    imgs: Optional[np.ndarray] = None,213    features: Literal["none", "wrfeat"] = "wrfeat",214    ndim: int = 2,215    n_workers=0,216    progbar_class=tqdm,217) -> List[WRFeatures]:218    detections = _check_dimensions(detections, ndim)219    imgs = _check_dimensions(imgs, ndim)220    logger.info(f"Extracting features from {len(detections)} detections")221    if n_workers > 0:222        logger.info(f"Using {n_workers} processes for feature extraction")223        features = joblib.Parallel(n_jobs=n_workers, backend="loky")(224            joblib.delayed(WRFeatures.from_mask_img)(225                # New axis for time component226                mask=mask[np.newaxis, ...].copy(),227                img=img[np.newaxis, ...].copy(),228                t_start=t,229            )230            for t, (mask, img) in progbar_class(231                enumerate(zip(detections, imgs)),232                total=len(imgs),233                desc="Extracting features",234            )235        )236    else:237        logger.info("Using single process for feature extraction")238        features = tuple(239            WRFeatures.from_mask_img(240                mask=mask[np.newaxis, ...],241                img=img[np.newaxis, ...],242                t_start=t,243            )244            for t, (mask, img) in progbar_class(245                enumerate(zip(detections, imgs)),246                total=len(imgs),247                desc="Extracting features",248            )249        )250 251    return features252 253 254def _check_dimensions(x: np.ndarray, ndim: int):255    if ndim == 2 and not x.ndim == 3:256        raise ValueError(f"Expected 2D data, got {x.ndim - 1}D data")257    elif ndim == 3:258        # if ndim=3 and data is two dimensional, it will be cast to 3D259        if x.ndim == 3:260            x = np.expand_dims(x, axis=1)261        elif x.ndim == 4:262            pass263        else:264            raise ValueError(f"Expected 3D data, got {x.ndim - 1}D data")265    return x266 267 268def build_windows_sd(269    features: List[WRFeatures], imgs_enc, imgs_stable, boxes, imgs, masks, window_size: int, progbar_class=tqdm270) -> List[dict]:271    windows = []272    for t1, t2 in progbar_class(273        zip(range(0, len(features)), range(window_size, len(features) + 1)),274        total=len(features) - window_size + 1,275        desc="Building windows",276    ):277        feat = WRFeatures.concat(features[t1:t2])278 279        labels = feat.labels280        timepoints = feat.timepoints281        coords = feat.coords282 283        if len(feat) == 0:284            coords = np.zeros((0, feat.ndim), dtype=int)285 286        w = dict(287            coords=coords,288            t1=t1,289            labels=labels,290            timepoints=timepoints,291            features=feat.features_stacked,292            img_enc=imgs_enc[t1:t2],293            image_stable=imgs_stable[t1:t2],294            boxes=boxes,295            img=imgs[t1:t2],296            mask=masks[t1:t2],297            coords_t=torch.tensor(coords, dtype=torch.float32),298            labels_t=torch.tensor(labels, dtype=torch.int32),299            timepoints_t=torch.tensor(timepoints, dtype=torch.int64),300            features_t=torch.tensor(feat.features_stacked, dtype=torch.float32),301            img_t=torch.tensor(imgs[t1:t2], dtype=torch.float32),302            mask_t=torch.tensor(masks[t1:t2], dtype=torch.int32),303        )304        windows.append(w)305 306    logger.debug(f"Built {len(windows)} track windows.\n")307    return windows308 309