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