"""
Copyright © 2022 Howard Hughes Medical Institute,
Authored by Carsen Stringer and Marius Pachitariu.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
3. Neither the name of HHMI nor the names of its contributors may be used to
endorse or promote products derived from this software without specific
prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
--------------------------------------------------------------------------
MEDIAR Prediction uses CellPose's Gradient Flow Tracking.
This code is adapted from the following codes:
[1] https://github.com/MouseLand/cellpose/blob/main/cellpose/utils.py
[2] https://github.com/MouseLand/cellpose/blob/main/cellpose/dynamics.py
[3] https://github.com/MouseLand/cellpose/blob/main/cellpose/metrics.py
"""
import torch
from torch.nn.functional import grid_sample
import numpy as np
import fastremap
from skimage import morphology
from scipy.ndimage import mean, find_objects
from scipy.ndimage.filters import maximum_filter1d
[docs]
torch_GPU = torch.device("cuda")
[docs]
torch_CPU = torch.device("cpu")
[docs]
def labels_to_flows(labels, use_gpu=False, device=None, redo_flows=False):
"""
Convert labels (list of masks or flows) to flows for training model
"""
# Labels b x 1 x h x w
labels = labels.cpu().numpy().astype(np.int16)
nimg = len(labels)
if labels[0].ndim < 3:
labels = [labels[n][np.newaxis, :, :] for n in range(nimg)]
# Flows need to be recomputed
if labels[0].shape[0] == 1 or labels[0].ndim < 3 or redo_flows:
# compute flows; labels are fixed here to be unique, so they need to be passed back
# make sure labels are unique!
labels = [fastremap.renumber(label, in_place=True)[0] for label in labels]
veci = [
masks_to_flows(labels[n][0], use_gpu=use_gpu, device=device)
for n in range(nimg)
]
# concatenate labels, distance transform, vector flows, heat (boundary and mask are computed in augmentations)
flows = [
np.concatenate((labels[n], labels[n] > 0.5, veci[n]), axis=0).astype(
np.float32
)
for n in range(nimg)
]
return np.array(flows)
[docs]
def compute_masks(
dP,
cellprob,
p=None,
niter=200,
cellprob_threshold=0.4,
flow_threshold=0.4,
interp=True,
resize=None,
use_gpu=False,
device=None,
):
"""compute masks using dynamics from dP, cellprob, and boundary"""
cp_mask = cellprob > cellprob_threshold
cp_mask = morphology.remove_small_holes(cp_mask, area_threshold=16)
cp_mask = morphology.remove_small_objects(cp_mask, min_size=16)
if np.any(cp_mask): # mask at this point is a cell cluster binary map, not labels
# follow flows
if p is None:
p, inds = follow_flows(
dP * cp_mask / 5.0,
niter=niter,
interp=interp,
use_gpu=use_gpu,
device=device,
)
if inds is None:
shape = resize if resize is not None else cellprob.shape
mask = np.zeros(shape, np.uint16)
p = np.zeros((len(shape), *shape), np.uint16)
return mask, p
# calculate masks
mask = get_masks(p, iscell=cp_mask)
# flow thresholding factored out of get_masks
shape0 = p.shape[1:]
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, dP, threshold=flow_threshold, use_gpu=use_gpu, device=device
)
else: # nothing to compute, just make it compatible
shape = resize if resize is not None else cellprob.shape
mask = np.zeros(shape, np.uint16)
p = np.zeros((len(shape), *shape), np.uint16)
return mask, p
def _extend_centers_gpu(
neighbors, centers, isneighbor, Ly, Lx, n_iter=200, device=torch.device("cuda")
):
if device is not None:
device = device
nimg = neighbors.shape[0] // 9
pt = torch.from_numpy(neighbors).to(device)
T = torch.zeros((nimg, Ly, Lx), dtype=torch.double, device=device)
meds = torch.from_numpy(centers.astype(int)).to(device).long()
isneigh = torch.from_numpy(isneighbor).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.0 + 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
[docs]
def diameters(masks):
_, counts = np.unique(np.int32(masks), return_counts=True)
counts = counts[1:]
md = np.median(counts ** 0.5)
if np.isnan(md):
md = 0
md /= (np.pi ** 0.5) / 2
return md, counts ** 0.5
[docs]
def masks_to_flows_gpu(masks, device=None):
if device is None:
device = torch.device("cuda")
Ly0, Lx0 = masks.shape
Ly, Lx = Ly0 + 2, Lx0 + 2
masks_padded = np.zeros((Ly, Lx), np.int64)
masks_padded[1:-1, 1:-1] = masks
# get mask pixel neighbors
y, x = np.nonzero(masks_padded)
neighborsY = np.stack((y, y - 1, y + 1, y, y, y - 1, y - 1, y + 1, y + 1), axis=0)
neighborsX = np.stack((x, x, x, x - 1, x + 1, x - 1, x + 1, x - 1, x + 1), axis=0)
neighbors = np.stack((neighborsY, neighborsX), axis=-1)
# get mask centers
slices = find_objects(masks)
centers = np.zeros((masks.max(), 2), "int")
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
yi, xi = np.nonzero(masks[sr, sc] == (i + 1))
yi = yi.astype(np.int32) + 1 # add padding
xi = xi.astype(np.int32) + 1 # add padding
ymed = np.median(yi)
xmed = np.median(xi)
imin = np.argmin((xi - xmed) ** 2 + (yi - ymed) ** 2)
xmed = xi[imin]
ymed = yi[imin]
centers[i, 0] = ymed + sr.start
centers[i, 1] = xmed + sc.start
# get neighbor validator (not all neighbors are in same mask)
neighbor_masks = masks_padded[neighbors[:, :, 0], neighbors[:, :, 1]]
isneighbor = 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, isneighbor, Ly, Lx, 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, Ly0, Lx0))
mu0[:, y - 1, x - 1] = mu
mu_c = np.zeros_like(mu0)
return mu0, mu_c
[docs]
def masks_to_flows(masks, use_gpu=False, device=None):
if masks.max() == 0 or (masks != 0).sum() == 1:
# dynamics_logger.warning('empty masks!')
return np.zeros((2, *masks.shape), "float32")
if use_gpu:
if use_gpu and device is None:
device = torch_GPU
elif device is None:
device = torch_CPU
masks_to_flows_device = masks_to_flows_gpu
if masks.ndim == 3:
Lz, Ly, Lx = masks.shape
mu = np.zeros((3, Lz, Ly, Lx), np.float32)
for z in range(Lz):
mu0 = masks_to_flows_device(masks[z], device=device)[0]
mu[[1, 2], z] += mu0
for y in range(Ly):
mu0 = masks_to_flows_device(masks[:, y], device=device)[0]
mu[[0, 2], :, y] += mu0
for x in range(Lx):
mu0 = masks_to_flows_device(masks[:, :, x], device=device)[0]
mu[[0, 1], :, :, x] += mu0
return mu
elif masks.ndim == 2:
mu, mu_c = masks_to_flows_device(masks, device=device)
return mu
else:
raise ValueError("masks_to_flows only takes 2D or 3D arrays")
[docs]
def steps2D_interp(p, dP, niter, use_gpu=False, device=None):
shape = dP.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(p[[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(dP[[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.0 / 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(niter):
# align_corners default is False, just added to suppress warning
dPt = 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.0, 1.0
)
# undo the normalization from before, reverse order of operations
pt = (pt + 1) * 0.5
for k in range(2):
pt[:, :, :, k] *= shape[k]
p = pt[:, :, :, [1, 0]].cpu().numpy().squeeze().T
return p
else:
assert print("ho")
[docs]
def follow_flows(dP, mask=None, niter=200, interp=True, use_gpu=True, device=None):
shape = np.array(dP.shape[1:]).astype(np.int32)
niter = np.uint32(niter)
p = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
p = np.array(p).astype(np.float32)
inds = np.array(np.nonzero(np.abs(dP[0]) > 1e-3)).astype(np.int32).T
if inds.ndim < 2 or inds.shape[0] < 5:
return p, None
if not interp:
assert print("woo")
else:
p_interp = steps2D_interp(
p[:, inds[:, 0], inds[:, 1]], dP, niter, use_gpu=use_gpu, device=device
)
p[:, inds[:, 0], inds[:, 1]] = p_interp
return p, inds
[docs]
def flow_error(maski, dP_net, use_gpu=False, device=None):
if dP_net.shape[1:] != maski.shape:
print("ERROR: net flow is not same size as predicted masks")
return
# flows predicted from estimated masks
dP_masks = masks_to_flows(maski, use_gpu=use_gpu, device=device)
# difference between predicted flows vs mask flows
flow_errors = np.zeros(maski.max())
for i in range(dP_masks.shape[0]):
flow_errors += mean(
(dP_masks[i] - dP_net[i] / 5.0) ** 2,
maski,
index=np.arange(1, maski.max() + 1),
)
return flow_errors, dP_masks
[docs]
def remove_bad_flow_masks(masks, flows, threshold=0.4, use_gpu=False, device=None):
merrors, _ = flow_error(masks, flows, use_gpu, device)
badi = 1 + (merrors > threshold).nonzero()[0]
masks[np.isin(masks, badi)] = 0
return masks
[docs]
def get_masks(p, iscell=None, rpad=20):
pflows = []
edges = []
shape0 = p.shape[1:]
dims = len(p)
for i in range(dims):
pflows.append(p[i].flatten().astype("int32"))
edges.append(np.arange(-0.5 - rpad, shape0[i] + 0.5 + rpad, 1))
h, _ = np.histogramdd(tuple(pflows), bins=edges)
hmax = h.copy()
for i in range(dims):
hmax = maximum_filter1d(hmax, 5, axis=i)
seeds = np.nonzero(np.logical_and(h - hmax > -1e-6, h > 10))
Nmax = h[seeds]
isort = np.argsort(Nmax)[::-1]
for s in seeds:
s = s[isort]
pix = list(np.array(seeds).T)
shape = h.shape
if dims == 3:
expand = np.nonzero(np.ones((3, 3, 3)))
else:
expand = np.nonzero(np.ones((3, 3)))
for e in expand:
e = np.expand_dims(e, 1)
for iter in range(5):
for k in range(len(pix)):
if iter == 0:
pix[k] = list(pix[k])
newpix = []
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]))
newpix.append(epix)
iin = np.all(tuple(iin), axis=0)
for p in newpix:
p = p[iin]
newpix = tuple(newpix)
igood = h[newpix] > 2
for i in range(dims):
pix[k][i] = newpix[i][igood]
if iter == 4:
pix[k] = tuple(pix[k])
M = np.zeros(h.shape, np.uint32)
for k in range(len(pix)):
M[pix[k]] = 1 + k
for i in range(dims):
pflows[i] = pflows[i] + rpad
M0 = M[tuple(pflows)]
# remove big masks
uniq, counts = fastremap.unique(M0, return_counts=True)
big = np.prod(shape0) * 0.9
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
M0 = np.reshape(M0, shape0)
return M0