dynamics.py

import logging

import cv2
import fastremap
import numpy as np
import scipy.ndimage
import torch
from numba import njit
from scipy.ndimage.filters import maximum_filter1d

import metrics
import transforms
import utils

dynamics_logger = logging.getLogger(__name__)

TORCH_ENABLED = True
torch_GPU = torch.device('cuda')
torch_CPU = torch.device('cpu')


@njit('(float64[:], int32[:], int32[:], int32, int32, int32, int32)',
      nogil=True,
      cache=True)
def _extend_centers(T, y, x, ymed, xmed, Lx, niter):
    """ run diffusion from center of mask (ymed, xmed) on mask pixels (y, x)
    Parameters
    --------------
    T: float64, array
        _ x Lx array that diffusion is run in
    y: int32, array
        pixels in y inside mask
    x: int32, array
        pixels in x inside mask
    ymed: int32
        center of mask in y
    xmed: int32
        center of mask in x
    Lx: int32
        size of x-dimension of masks
    niter: int32
        number of iterations to run diffusion
    Returns
    ---------------
    T: float64, array
        amount of diffused particles at each pixel
    """

    for t in range(niter):
        T[ymed * Lx + xmed] += 1
        T[y * Lx +
          x] = 1 / 9. * (T[y * Lx + x] + T[(y - 1) * Lx + x] + T[
              (y + 1) * Lx + x] + T[y * Lx + x - 1] + T[y * Lx + x + 1] +
                         T[(y - 1) * Lx + x - 1] + T[(y - 1) * Lx + x + 1] +
                         T[(y + 1) * Lx + x - 1] + T[(y + 1) * Lx + x + 1])
    return T


def _extend_centers_gpu(neighbors: np.ndarray,
                        centers: np.ndarray,
                        is_neighbor: np.ndarray,
                        height: int,
                        width: int,
                        n_iter: int = 200,
                        device=None):
    """ runs diffusion on GPU to generate flows for training images or quality control

    neighbors is 9 x pixels in masks,
    centers are mask centers,
    isneighbor is valid neighbor boolean 9 x pixels

    """
    if device is not None:
        device = torch_GPU

    n_img = neighbors.shape[0] // 9
    pt = torch.from_numpy(neighbors).to(device)

    T = torch.zeros((n_img, height, width), dtype=torch.double, device=device)
    meds = torch.from_numpy(centers.astype(int)).to(device).long()
    isneigh = torch.from_numpy(is_neighbor).to(device)
    for i in range(n_iter):
        T[:, meds[:, 0], meds[:, 1]] += 1
        Tneigh = T[:, pt[:, :, 0], pt[:, :, 1]]
        Tneigh *= isneigh
        T[:, pt[0, :, 0], pt[0, :, 1]] = Tneigh.mean(axis=1)
    del meds, isneigh, Tneigh
    T = torch.log(1. + T)
    # gradient positions
    grads = T[:, pt[[2, 1, 4, 3], :, 0], pt[[2, 1, 4, 3], :, 1]]
    del pt
    dy = grads[:, 0] - grads[:, 1]
    dx = grads[:, 2] - grads[:, 3]
    del grads
    mu_torch = np.stack((dy.cpu().squeeze(), dx.cpu().squeeze()), axis=-2)
    return mu_torch


def masks_to_flows_gpu(mask: np.ndarray, device=None):
    """ convert masks to flows using diffusion from center pixel
    Center of masks where diffusion starts is defined using COM
    Parameters
    -------------
    mask: int, 2D or 3D array
        labelled masks 0=NO masks; 1,2,...=mask labels
    Returns
    -------------
    mu: float, 3D or 4D array
        flows in Y = mu[-2], flows in X = mu[-1].
        if masks are 3D, flows in Z = mu[0].
    mu_c: float, 2D or 3D array
        for each pixel, the distance to the center of the mask
        in which it resides
    """
    if device is None:
        device = torch_GPU

    height, width = mask.shape
    height_padded, width_padded = height + 2, width + 2

    masks_padded = np.zeros((height_padded, width_padded), np.int64)
    masks_padded[1:-1, 1:-1] = mask

    # get mask pixel neighbors
    y, x = np.nonzero(masks_padded)
    neighbors_y = np.stack((y, y - 1, y + 1, y, y, y - 1, y - 1, y + 1, y + 1),
                           axis=0)
    neighbors_x = np.stack((x, x, x, x - 1, x + 1, x - 1, x + 1, x - 1, x + 1),
                           axis=0)
    neighbors = np.stack((neighbors_y, neighbors_x), axis=-1)

    # get mask centers
    slices = scipy.ndimage.find_objects(mask)

    centers = np.zeros((mask.max(), 2), 'int')
    for i, si in enumerate(slices):
        if si is not None:
            sr, sc = si
            y_i, x_i = np.nonzero(mask[sr, sc] == (i + 1))
            # add padding
            y_i, x_i = y_i.astype(np.int32) + 1, x_i.astype(np.int32) + 1
            y_med, x_med = np.median(y_i), np.median(x_i)
            i_min = np.argmin((x_i - x_med)**2 + (y_i - y_med)**2)
            x_med, y_med = x_i[i_min], y_i[i_min]
            centers[i, 0], centers[i, 1] = y_med + sr.start, x_med + sc.start

    # get neighbor validator (not all neighbors are in same mask)
    neighbor_masks = masks_padded[neighbors[:, :, 0], neighbors[:, :, 1]]
    is_neighbor = neighbor_masks == neighbor_masks[0]
    ext = np.array(
        [[sr.stop - sr.start + 1, sc.stop - sc.start + 1] for sr, sc in slices])
    n_iter = 2 * (ext.sum(axis=1)).max()
    # run diffusion
    mu = _extend_centers_gpu(neighbors,
                             centers,
                             is_neighbor,
                             height_padded,
                             width_padded,
                             n_iter=n_iter,
                             device=device)

    # normalize
    mu /= (1e-20 + (mu**2).sum(axis=0)**0.5)

    # put into original image
    mu0 = np.zeros((2, height, width))
    mu0[:, y - 1, x - 1] = mu
    mu_c = np.zeros_like(mu0)
    return mu0, mu_c


def _masks_to_flows_cpu(masks, **kwargs):
    """ convert masks to flows using diffusion from center pixel
    Center of masks where diffusion starts is defined to be the
    closest pixel to the median of all pixels that is inside the
    mask. Result of diffusion is converted into flows by computing
    the gradients of the diffusion density map.
    Parameters
    -------------
    masks: int, 2D array
        labelled masks 0=NO masks; 1,2,...=mask labels
    Returns
    -------------
    mu: float, 3D array
        flows in Y = mu[-2], flows in X = mu[-1].
        if masks are 3D, flows in Z = mu[0].
    mu_c: float, 2D array
        for each pixel, the distance to the center of the mask
        in which it resides
    """

    height, width = masks.shape
    mu = np.zeros((2, height, width), np.float64)
    mu_c = np.zeros((height, width), np.float64)

    slices = scipy.ndimage.find_objects(masks)
    dia = utils.diameters(masks)[0]
    s2 = (.15 * dia)**2
    for i, si in enumerate(slices):
        if si is not None:
            sr, sc = si
            ly, lx = sr.stop - sr.start + 1, sc.stop - sc.start + 1
            y, x = np.nonzero(masks[sr, sc] == (i + 1))
            y, x = y.astype(np.int32) + 1, x.astype(np.int32) + 1
            y_med, x_med = np.median(y), np.median(x)
            i_min = np.argmin((x - x_med)**2 + (y - y_med)**2)
            x_med, y_med = x[i_min], y[i_min]

            d2 = (x - x_med)**2 + (y - y_med)**2
            mu_c[sr.start + y - 1, sc.start + x - 1] = np.exp(-d2 / s2)

            niter = 2 * np.int32(np.ptp(x) + np.ptp(y))
            t = np.zeros((ly + 2) * (lx + 2), np.float64)
            t = _extend_centers(t, y, x, y_med, x_med, np.int32(lx),
                                np.int32(niter))
            t[(y + 1) * lx + x + 1] = np.log(1. + t[(y + 1) * lx + x + 1])

            dy = t[(y + 1) * lx + x] - t[(y - 1) * lx + x]
            dx = t[y * lx + x + 1] - t[y * lx + x - 1]
            mu[:, sr.start + y - 1, sc.start + x - 1] = np.stack((dy, dx))

    mu /= (1e-20 + (mu**2).sum(axis=0)**0.5)

    return mu, mu_c


def _masks_to_flows(mask: np.ndarray, use_gpu=False, device=None):
    """ convert masks to flows using diffusion from center pixel

    Center of masks where diffusion starts is defined to be the
    closest pixel to the median of all pixels that is inside the
    mask. Result of diffusion is converted into flows by computing
    the gradients of the diffusion density map.

    Parameters
    -------------

    mask: int, 2D or 3D array
        labelled masks 0=NO masks; 1,2,...=mask labels

    Returns
    -------------

    mu: float, 3D or 4D array
        flows in Y = mu[-2], flows in X = mu[-1].
        if masks are 3D, flows in Z = mu[0].

    mu_c: float, 2D or 3D array
        for each pixel, the distance to the center of the mask
        in which it resides

    """
    if mask.max() == 0:
        dynamics_logger.warning('empty masks!')
        return np.zeros((2, *mask.shape), 'float32')

    if use_gpu:
        device = torch_GPU
        masks_to_flows_device = masks_to_flows_gpu
    else:
        masks_to_flows_device = _masks_to_flows_cpu

    mu, mu_c = masks_to_flows_device(mask, device=device)
    return mu


@njit([
    '(int16[:,:,:], float32[:], float32[:], float32[:,:])',
    '(float32[:,:,:], float32[:], float32[:], float32[:,:])'
],
      cache=True)
def _map_coordinates(I, yc, xc, Y):
    """
    bilinear interpolation of image 'I' in-place with ycoordinates yc and xcoordinates xc to Y

    Parameters
    -------------
    I : C x Ly x Lx
    yc : ni
        new y coordinates
    xc : ni
        new x coordinates
    Y : C x ni
        I sampled at (yc,xc)
    """
    C, Ly, Lx = I.shape
    yc_floor = yc.astype(np.int32)
    xc_floor = xc.astype(np.int32)
    yc = yc - yc_floor
    xc = xc - xc_floor
    for i in range(yc_floor.shape[0]):
        yf = min(Ly - 1, max(0, yc_floor[i]))
        xf = min(Lx - 1, max(0, xc_floor[i]))
        yf1 = min(Ly - 1, yf + 1)
        xf1 = min(Lx - 1, xf + 1)
        y = yc[i]
        x = xc[i]
        for c in range(C):
            Y[c, i] = (np.float32(I[c, yf, xf]) * (1 - y) * (1 - x) +
                       np.float32(I[c, yf, xf1]) * (1 - y) * x +
                       np.float32(I[c, yf1, xf]) * y * (1 - x) +
                       np.float32(I[c, yf1, xf1]) * y * x)


def _steps_2d_interpolation(pixel_locations: np.ndarray,
                            gradients: np.ndarray,
                            n_iter: int,
                            use_gpu: bool = True,
                            device=None):
    shape = gradients.shape[1:]
    if use_gpu:
        if device is None:
            device = torch_GPU
        shape = np.array(shape)[[1, 0]].astype(
            'float') - 1  # Y and X dimensions (dP is 2.Ly.Lx), flipped X-1, Y-1
        pt = torch.from_numpy(pixel_locations[[
            1, 0
        ]].T).float().to(device).unsqueeze(0).unsqueeze(
            0)  # p is n_points by 2, so pt is [1 1 2 n_points]
        im = torch.from_numpy(gradients[[1, 0]]).float().to(device).unsqueeze(
            0)  # covert flow numpy array to tensor on GPU, add dimension
        # normalize pt between  0 and  1, normalize the flow
        for k in range(2):
            im[:, k, :, :] *= 2. / shape[k]
            pt[:, :, :, k] /= shape[k]

        # normalize to between -1 and 1
        pt = pt * 2 - 1

        # here is where the stepping happens
        for t in range(n_iter):
            # align_corners default is False, just added to suppress warning
            dPt = torch.nn.functional.grid_sample(im, pt, align_corners=False)

            for k in range(2):  # clamp the final pixel locations
                pt[:, :, :, k] = torch.clamp(pt[:, :, :, k] + dPt[:, k, :, :],
                                             -1., 1.)

        # undo the normalization from before, reverse order of operations
        pt = (pt + 1) * 0.5
        for k in range(2):
            pt[:, :, :, k] *= shape[k]

        pixel_locations = pt[:, :, :, [1, 0]].cpu().numpy().squeeze().T
        return pixel_locations

    else:
        dPt = np.zeros(pixel_locations.shape, np.float32)

        for t in range(n_iter):
            _map_coordinates(gradients.astype(np.float32), pixel_locations[0],
                             pixel_locations[1], dPt)
            for k in range(len(pixel_locations)):
                pixel_locations[k] = np.minimum(
                    shape[k] - 1, np.maximum(0, pixel_locations[k] + dPt[k]))
        return pixel_locations


@njit('(float32[:,:,:], float32[:,:,:], int32[:,:], int32)', nogil=True)
def _steps_2d(pixel_locations: np.ndarray, gradients: np.ndarray,
              pixel_indexes: np.ndarray, n_iter: int):
    """ run dynamics of pixels to recover masks in 2D

    Euler integration of dynamics gradients for n_iter steps

    Parameters
    ----------------

    pixel_locations: float32, 3D array
        pixel locations [axis x Ly x Lx] (start at initial meshgrid)

    gradients: float32, 3D array
        flows [axis x Ly x Lx]

    pixel_indexes: int32, 2D array
        non-zero pixels to run dynamics on [npixels x 2]

    n_iter: int32
        number of iterations of dynamics to run

    Returns
    ---------------

    p: float32, 3D array
        final locations of each pixel after dynamics

    """
    shape = pixel_locations.shape[1:]
    for t in range(n_iter):
        for j in range(pixel_indexes.shape[0]):
            # starting coordinates
            y, x = pixel_indexes[j, 0], pixel_indexes[j, 1]
            p0, p1 = int(pixel_locations[0, y, x]), int(pixel_locations[1, y,
                                                                        x])
            step = gradients[:, p0, p1]
            for k in range(pixel_locations.shape[0]):
                pixel_locations[k, y, x] = min(
                    shape[k] - 1, max(0, pixel_locations[k, y, x] + step[k]))
    return pixel_locations


def _follow_flows(gradients: np.ndarray,
                  n_iter: float = 200,
                  interp: bool = True,
                  use_gpu: bool = True,
                  device=None):
    """ define pixels and run dynamics to recover masks in 2D

    Pixels are meshgrid. Only pixels with non-zero cell-probability
    are used (as defined by pixel_indexes)

    Parameters
    ----------------

    gradients: float32, 3D or 4D array
        flows [axis x Ly x Lx] or [axis x Lz x Ly x Lx]


    n_iter: int (optional, default 200)
        number of iterations of dynamics to run

    interp: bool (optional, default True)
        interpolate during 2D dynamics (not available in 3D)
        (in previous versions + paper it was False)

    use_gpu: bool (optional, default False)
        use GPU to run interpolated dynamics (faster than CPU)


    Returns
    ---------------

    pixel_locations: float32, 3D or 4D array
        final locations of each pixel after dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx]

    pixel_indexes: int32, 3D or 4D array
        indices of pixels used for dynamics; [axis x Ly x Lx] or [axis x Lz x Ly x Lx]

    """
    shape = np.array(gradients.shape[1:]).astype(np.int32)
    n_iter = np.uint32(n_iter)

    pixel_locations = np.meshgrid(np.arange(shape[0]),
                                  np.arange(shape[1]),
                                  indexing='ij')
    pixel_locations = np.array(pixel_locations, dtype=np.float32)

    pixel_indexes = np.array(np.nonzero(np.abs(gradients[0]) > 1e-3),
                             dtype=np.int32).T

    if pixel_indexes.ndim < 2 or pixel_indexes.shape[0] < 5:
        # dynamics_logger.warning('WARNING: no mask pixels found')
        return pixel_locations, None

    if interp:
        p_interp = _steps_2d_interpolation(pixel_locations[:, pixel_indexes[:,
                                                                            0],
                                                           pixel_indexes[:, 1]],
                                           gradients,
                                           n_iter,
                                           use_gpu=use_gpu,
                                           device=device)
        pixel_locations[:, pixel_indexes[:, 0], pixel_indexes[:, 1]] = p_interp
    else:
        pixel_locations = _steps_2d(pixel_locations, gradients, pixel_indexes,
                                    n_iter)

    return pixel_locations, pixel_indexes


def _remove_bad_flow_masks(mask: np.ndarray,
                           flows,
                           threshold=0.4,
                           use_gpu=False,
                           device=None):
    """ remove masks which have inconsistent flows

    Uses metrics.flow_error to compute flows from predicted masks
    and compare flows to predicted flows from network. Discards
    masks with flow errors greater than the threshold.

    Parameters
    ----------------

    mask: int, 2D or 3D array
        labelled masks, 0=NO masks; 1,2,...=mask labels,
        size [Ly x Lx] or [Lz x Ly x Lx]

    flows: float, 3D or 4D array
        flows [axis x Ly x Lx] or [axis x Lz x Ly x Lx]

    threshold: float (optional, default 0.4)
        masks with flow error greater than threshold are discarded.

    Returns
    ---------------

    masks: int, 2D or 3D array
        masks with inconsistent flow masks removed,
        0=NO masks; 1,2,...=mask labels,
        size [Ly x Lx] or [Lz x Ly x Lx]

    """
    mask_errors, _ = metrics.flow_error(mask, flows, use_gpu, device)
    bad_indexes = 1 + (mask_errors > threshold).nonzero()[0]
    mask[np.isin(mask, bad_indexes)] = 0
    return mask


def _get_masks(pixel_locations, cell_probability_mask=None, r_pad=20):
    """ create masks using pixel convergence after running dynamics

    Makes a histogram of final pixel locations p, initializes masks
    at peaks of histogram and extends the masks from the peaks so that
    they include all pixels with more than 2 final pixels p. Discards
    masks with flow errors greater than the threshold.
    Parameters
    ----------------
    pixel_locations: float32, 3D or 4D array
        final locations of each pixel after dynamics,
        size [axis x Ly x Lx] or [axis x Lz x Ly x Lx].
    cell_probability_mask: bool, 2D or 3D array
        if iscell is not None, set pixels that are
        iscell False to stay in their original location.
    r_pad: int (optional, default 20)
        histogram edge padding

    Returns
    ---------------
    m0: int, 2D or 3D array
        masks with inconsistent flow masks removed,
        0=NO masks; 1,2,...=mask labels,
        size [Ly x Lx] or [Lz x Ly x Lx]

    """

    p_flows, edges = [], []
    shape0 = pixel_locations.shape[1:]
    dims = len(pixel_locations)
    if cell_probability_mask is not None:
        indexes = np.meshgrid(np.arange(shape0[0]),
                              np.arange(shape0[1]),
                              indexing='ij')

        for i in range(dims):
            pixel_locations[
                i, ~cell_probability_mask] = indexes[i][~cell_probability_mask]

    for i in range(dims):
        p_flows.append(pixel_locations[i].flatten().astype('int32'))
        edges.append(np.arange(-.5 - r_pad, shape0[i] + .5 + r_pad, 1))

    histogram, _ = np.histogramdd(tuple(p_flows), bins=edges)
    histogram_max = histogram.copy()
    for i in range(dims):
        histogram_max = maximum_filter1d(histogram_max, 5, axis=i)

    seeds = np.nonzero(
        np.logical_and(histogram - histogram_max > -1e-6, histogram > 10))
    pix = list(np.array(seeds).T)

    shape = histogram.shape
    expand = np.nonzero(np.ones((3, 3)))

    for n_iter in range(5):
        for k in range(len(pix)):
            if n_iter == 0:
                pix[k] = list(pix[k])

            new_pix, iin = [], []
            for i, e in enumerate(expand):
                epix = e[:, np.newaxis] + np.expand_dims(pix[k][i], 0) - 1
                epix = epix.flatten()

                iin.append(np.logical_and(epix >= 0, epix < shape[i]))
                new_pix.append(epix)

            new_pix = tuple(new_pix)
            i_good = histogram[new_pix] > 2
            for i in range(dims):
                pix[k][i] = new_pix[i][i_good]
            if n_iter == 4:
                pix[k] = tuple(pix[k])

    m = np.zeros(histogram.shape, np.uint32)
    for k in range(len(pix)):
        m[pix[k]] = 1 + k

    for i in range(dims):
        p_flows[i] = p_flows[i] + r_pad
    m0 = m[tuple(p_flows)]

    # remove big masks
    uniq, counts = fastremap.unique(m0, return_counts=True)
    big = np.prod(shape0) * 0.4
    bigc = uniq[counts > big]
    if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):
        m0 = fastremap.mask(m0, bigc)
    fastremap.renumber(
        m0, in_place=True)  # convenient to guarantee non-skipped labels
    return np.reshape(m0, shape0)


def compute_masks(gradients: np.ndarray,
                  cell_probability: np.ndarray,
                  n_iter: float = 200,
                  cell_probability_threshold: float = 0.0,
                  flow_threshold: float = 0.4,
                  interp: bool = True,
                  min_size: int = 15,
                  resize=None,
                  device=None,
                  use_gpu: bool = True):
    cell_probability_mask = cell_probability > cell_probability_threshold

    if not np.any(cell_probability_mask):
        shape = resize if resize is not None else cell_probability.shape
        mask = np.zeros(shape, np.uint16)
        pixel_locations = np.zeros((len(shape), *shape), np.uint16)
        return mask, pixel_locations

    # follow flows
    pixel_locations, pixel_indexes = \
        _follow_flows(gradients * cell_probability_mask / 5., n_iter=n_iter, interp=interp, device=device)

    if pixel_indexes is None:
        # dynamics_logger.info('No cell pixels found.')
        shape = resize if resize is not None else cell_probability.shape
        mask = np.zeros(shape, np.uint16)
        pixel_locations = np.zeros((len(shape), *shape), np.uint16)
        return mask, pixel_locations

    # calculate masks
    mask = _get_masks(pixel_locations,
                      cell_probability_mask=cell_probability_mask)

    # flow thresholding factored out of get_masks
    if mask.max() > 0 and flow_threshold is not None and flow_threshold > 0:
        # make sure labels are unique at output of get_masks
        mask = _remove_bad_flow_masks(mask,
                                      gradients,
                                      threshold=flow_threshold,
                                      use_gpu=use_gpu,
                                      device=device)

    if resize is not None:
        if mask.max() > 2**16 - 1:
            recast = True
            mask = mask.astype(np.float32)
        else:
            recast = False
            mask = mask.astype(np.uint16)
        mask = transforms.resize_image(mask,
                                       resize[0],
                                       resize[1],
                                       interpolation=cv2.INTER_NEAREST)
        if recast:
            mask = mask.astype(np.uint32)

    elif mask.max() < 2**16:
        mask = mask.astype(np.uint16)

    # moving the cleanup to the end helps avoid some bugs arising from scaling...
    # maybe better would be to rescale the min_size and hole_size parameters to do the
    # cleanup at the prediction scale, or switch depending on which one is bigger...
    mask = utils.fill_holes_and_remove_small_masks(mask, min_size=min_size)

    if mask.dtype == np.uint32:
        dynamics_logger.warning(
            'more than 65535 masks in image, masks returned as np.uint32')

    return mask, pixel_locations