VisionLanguageGroup/MicroscopyMatching
0
1"""2Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.3"""4import os5from scipy.ndimage import find_objects, center_of_mass, mean6import torch7import numpy as np8import tifffile9from tqdm import trange10import fastremap11 12import logging13 14dynamics_logger = logging.getLogger(__name__)15 16from . import utils17 18import torch19import torch.nn.functional as F20 21def _extend_centers_gpu(neighbors, meds, isneighbor, shape, n_iter=200, 22 device=torch.device("cpu")):23 """Runs diffusion on GPU to generate flows for training images or quality control.24 25 Args:26 neighbors (torch.Tensor): 9 x pixels in masks.27 meds (torch.Tensor): Mask centers.28 isneighbor (torch.Tensor): Valid neighbor boolean 9 x pixels.29 shape (tuple): Shape of the tensor.30 n_iter (int, optional): Number of iterations. Defaults to 200.31 device (torch.device, optional): Device to run the computation on. Defaults to torch.device("cpu").32 33 Returns:34 torch.Tensor: Generated flows.35 36 """37 if torch.prod(torch.tensor(shape)) > 4e7 or device.type == "mps":38 T = torch.zeros(shape, dtype=torch.float, device=device)39 else:40 T = torch.zeros(shape, dtype=torch.double, device=device)41 42 for i in range(n_iter):43 T[tuple(meds.T)] += 144 Tneigh = T[tuple(neighbors)]45 Tneigh *= isneighbor46 T[tuple(neighbors[:, 0])] = Tneigh.mean(axis=0)47 del meds, isneighbor, Tneigh48 49 if T.ndim == 2:50 grads = T[neighbors[0, [2, 1, 4, 3]], neighbors[1, [2, 1, 4, 3]]]51 del neighbors52 dy = grads[0] - grads[1]53 dx = grads[2] - grads[3]54 del grads55 mu_torch = np.stack((dy.cpu().squeeze(0), dx.cpu().squeeze(0)), axis=-2)56 else:57 grads = T[tuple(neighbors[:, 1:])]58 del neighbors59 dz = grads[0] - grads[1]60 dy = grads[2] - grads[3]61 dx = grads[4] - grads[5]62 del grads63 mu_torch = np.stack(64 (dz.cpu().squeeze(0), dy.cpu().squeeze(0), dx.cpu().squeeze(0)), axis=-2)65 return mu_torch66 67def center_of_mass(mask):68 yi, xi = np.nonzero(mask)69 ymean = int(np.round(yi.sum() / len(yi)))70 xmean = int(np.round(xi.sum() / len(xi)))71 if not ((yi==ymean) * (xi==xmean)).sum():72 # center is closest point to (ymean, xmean) within mask73 imin = ((xi - xmean)**2 + (yi - ymean)**2).argmin()74 ymean = yi[imin]75 xmean = xi[imin]76 77 return ymean, xmean78 79def get_centers(masks, slices):80 centers = [center_of_mass(masks[slices[i]]==(i+1)) for i in range(len(slices))]81 centers = np.array([np.array([centers[i][0] + slices[i][0].start, centers[i][1] + slices[i][1].start]) 82 for i in range(len(slices))])83 exts = np.array([(slc[0].stop - slc[0].start) + (slc[1].stop - slc[1].start) + 2 for slc in slices])84 return centers, exts85 86 87def masks_to_flows_gpu(masks, device=torch.device("cpu"), niter=None):88 """Convert masks to flows using diffusion from center pixel.89 90 Center of masks where diffusion starts is defined by pixel closest to median within the mask.91 92 Args:93 masks (int, 2D or 3D array): Labelled masks. 0=NO masks; 1,2,...=mask labels.94 device (torch.device, optional): The device to run the computation on. Defaults to torch.device("cpu").95 niter (int, optional): Number of iterations for the diffusion process. Defaults to None.96 97 Returns:98 np.ndarray: A 4D array representing the flows for each pixel in Z, X, and Y.99 100 101 Returns:102 A tuple containing (mu, meds_p). mu is float 3D or 4D array of flows in (Z)XY. 103 meds_p are cell centers.104 """105 if device is None:106 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('mps') if torch.backends.mps.is_available() else None107 108 if masks.max() > 0:109 Ly0, Lx0 = masks.shape110 Ly, Lx = Ly0 + 2, Lx0 + 2111 112 masks_padded = torch.from_numpy(masks.astype("int64")).to(device)113 masks_padded = F.pad(masks_padded, (1, 1, 1, 1))114 shape = masks_padded.shape115 116 ### get mask pixel neighbors117 y, x = torch.nonzero(masks_padded, as_tuple=True)118 y = y.int()119 x = x.int()120 neighbors = torch.zeros((2, 9, y.shape[0]), dtype=torch.int, device=device)121 yxi = [[0, -1, 1, 0, 0, -1, -1, 1, 1], [0, 0, 0, -1, 1, -1, 1, -1, 1]]122 for i in range(9):123 neighbors[0, i] = y + yxi[0][i]124 neighbors[1, i] = x + yxi[1][i]125 isneighbor = torch.ones((9, y.shape[0]), dtype=torch.bool, device=device)126 m0 = masks_padded[neighbors[0, 0], neighbors[1, 0]]127 for i in range(1, 9):128 isneighbor[i] = masks_padded[neighbors[0, i], neighbors[1, i]] == m0129 del m0, masks_padded130 131 ### get center-of-mass within cell132 slices = find_objects(masks)133 centers, ext = get_centers(masks, slices)134 meds_p = torch.from_numpy(centers).to(device).long()135 meds_p += 1 # for padding136 137 ### run diffusion138 n_iter = 2 * ext.max() if niter is None else niter139 mu = _extend_centers_gpu(neighbors, meds_p, isneighbor, shape, n_iter=n_iter,140 device=device)141 mu = mu.astype("float64")142 143 # new normalization144 mu /= (1e-60 + (mu**2).sum(axis=0)**0.5)145 146 # put into original image147 mu0 = np.zeros((2, Ly0, Lx0))148 mu0[:, y.cpu().numpy() - 1, x.cpu().numpy() - 1] = mu149 else:150 # no masks, return empty flows151 mu0 = np.zeros((2, masks.shape[0], masks.shape[1]))152 return mu0153 154 155 156def flow_error(maski, dP_net, device=None):157 """Error in flows from predicted masks vs flows predicted by network run on image.158 159 This function serves to benchmark the quality of masks. It works as follows:160 1. The predicted masks are used to create a flow diagram.161 2. The mask-flows are compared to the flows that the network predicted.162 163 If there is a discrepancy between the flows, it suggests that the mask is incorrect.164 Masks with flow_errors greater than 0.4 are discarded by default. This setting can be165 changed in Cellpose.eval or CellposeModel.eval.166 167 Args:168 maski (np.ndarray, int): Masks produced from running dynamics on dP_net, where 0=NO masks; 1,2... are mask labels.169 dP_net (np.ndarray, float): ND flows where dP_net.shape[1:] = maski.shape.170 171 Returns:172 A tuple containing (flow_errors, dP_masks): flow_errors (np.ndarray, float): Mean squared error between predicted flows and flows from masks; 173 dP_masks (np.ndarray, float): ND flows produced from the predicted masks.174 """175 if dP_net.shape[1:] != maski.shape:176 print("ERROR: net flow is not same size as predicted masks")177 return178 179 # flows predicted from estimated masks180 dP_masks = masks_to_flows_gpu(maski, device=device)181 # difference between predicted flows vs mask flows182 flow_errors = np.zeros(maski.max())183 for i in range(dP_masks.shape[0]):184 flow_errors += mean((dP_masks[i] - dP_net[i] / 5.)**2, maski,185 index=np.arange(1,186 maski.max() + 1))187 188 return flow_errors, dP_masks189 190 191def steps_interp(dP, inds, niter, device=torch.device("cpu")):192 """ Run dynamics of pixels to recover masks in 2D/3D, with interpolation between pixel values.193 194 Euler integration of dynamics dP for niter steps.195 196 Args:197 p (numpy.ndarray): Array of shape (n_points, 2 or 3) representing the initial pixel locations.198 dP (numpy.ndarray): Array of shape (2, Ly, Lx) or (3, Lz, Ly, Lx) representing the flow field.199 niter (int): Number of iterations to perform.200 device (torch.device, optional): Device to use for computation. Defaults to None.201 202 Returns:203 numpy.ndarray: Array of shape (n_points, 2) or (n_points, 3) representing the final pixel locations.204 205 Raises:206 None207 208 """209 210 shape = dP.shape[1:]211 ndim = len(shape)212 213 pt = torch.zeros((*[1]*ndim, len(inds[0]), ndim), dtype=torch.float32, device=device)214 im = torch.zeros((1, ndim, *shape), dtype=torch.float32, device=device)215 # Y and X dimensions, flipped X-1, Y-1216 # pt is [1 1 1 3 n_points]217 for n in range(ndim):218 if ndim==3:219 pt[0, 0, 0, :, ndim - n - 1] = torch.from_numpy(inds[n]).to(device, dtype=torch.float32)220 else:221 pt[0, 0, :, ndim - n - 1] = torch.from_numpy(inds[n]).to(device, dtype=torch.float32)222 im[0, ndim - n - 1] = torch.from_numpy(dP[n]).to(device, dtype=torch.float32)223 shape = np.array(shape)[::-1].astype("float") - 1 224 225 # normalize pt between 0 and 1, normalize the flow226 for k in range(ndim):227 im[:, k] *= 2. / shape[k]228 pt[..., k] /= shape[k]229 230 # normalize to between -1 and 1231 pt *= 2 232 pt -= 1233 234 # dynamics235 for t in range(niter):236 dPt = torch.nn.functional.grid_sample(im, pt, align_corners=False)237 for k in range(ndim): #clamp the final pixel locations238 pt[..., k] = torch.clamp(pt[..., k] + dPt[:, k], -1., 1.)239 240 #undo the normalization from before, reverse order of operations241 pt += 1 242 pt *= 0.5243 for k in range(ndim):244 pt[..., k] *= shape[k]245 246 if ndim==3:247 pt = pt[..., [2, 1, 0]].squeeze()248 pt = pt.unsqueeze(0) if pt.ndim==1 else pt 249 return pt.T250 else:251 pt = pt[..., [1, 0]].squeeze()252 pt = pt.unsqueeze(0) if pt.ndim==1 else pt253 return pt.T254 255 256 257def remove_bad_flow_masks(masks, flows, threshold=0.4, device=torch.device("cpu")):258 """Remove masks which have inconsistent flows.259 260 Uses metrics.flow_error to compute flows from predicted masks 261 and compare flows to predicted flows from the network. Discards 262 masks with flow errors greater than the threshold.263 264 Args:265 masks (int, 2D or 3D array): Labelled masks, 0=NO masks; 1,2,...=mask labels,266 size [Ly x Lx] or [Lz x Ly x Lx].267 flows (float, 3D or 4D array): Flows [axis x Ly x Lx] or [axis x Lz x Ly x Lx].268 threshold (float, optional): Masks with flow error greater than threshold are discarded.269 Default is 0.4.270 271 Returns:272 masks (int, 2D or 3D array): Masks with inconsistent flow masks removed,273 0=NO masks; 1,2,...=mask labels, size [Ly x Lx] or [Lz x Ly x Lx].274 """275 device0 = device276 if masks.size > 10000 * 10000 and (device is not None and device.type == "cuda"):277 278 major_version, minor_version = torch.__version__.split(".")[:2]279 torch.cuda.empty_cache()280 if major_version == "1" and int(minor_version) < 10:281 # for PyTorch version lower than 1.10282 def mem_info():283 total_mem = torch.cuda.get_device_properties(device0.index).total_memory284 used_mem = torch.cuda.memory_allocated(device0.index)285 free_mem = total_mem - used_mem286 return total_mem, free_mem287 else:288 # for PyTorch version 1.10 and above289 def mem_info():290 free_mem, total_mem = torch.cuda.mem_get_info(device0.index)291 return total_mem, free_mem292 total_mem, free_mem = mem_info()293 if masks.size * 32 > free_mem:294 dynamics_logger.warning(295 "WARNING: image is very large, not using gpu to compute flows from masks for QC step flow_threshold"296 )297 dynamics_logger.info("turn off QC step with flow_threshold=0 if too slow")298 device0 = torch.device("cpu")299 300 merrors, _ = flow_error(masks, flows, device0)301 badi = 1 + (merrors > threshold).nonzero()[0]302 masks[np.isin(masks, badi)] = 0303 return masks304 305 306def max_pool1d(h, kernel_size=5, axis=1, out=None):307 """ memory efficient max_pool thanks to Mark Kittisopikul 308 309 for stride=1, padding=kernel_size//2, requires odd kernel_size >= 3310 311 """312 if out is None:313 out = h.clone()314 else:315 out.copy_(h)316 317 nd = h.shape[axis] 318 k0 = kernel_size // 2319 for d in range(-k0, k0+1):320 if axis==1:321 mv = out[:, max(-d,0):min(nd-d,nd)]322 hv = h[:, max(d,0):min(nd+d,nd)]323 elif axis==2:324 mv = out[:, :, max(-d,0):min(nd-d,nd)]325 hv = h[:, :, max(d,0):min(nd+d,nd)]326 elif axis==3:327 mv = out[:, :, :, max(-d,0):min(nd-d,nd)]328 hv = h[:, :, :, max(d,0):min(nd+d,nd)]329 torch.maximum(mv, hv, out=mv)330 return out331 332def max_pool_nd(h, kernel_size=5):333 """ memory efficient max_pool in 2d or 3d """334 ndim = h.ndim - 1335 hmax = max_pool1d(h, kernel_size=kernel_size, axis=1)336 hmax2 = max_pool1d(hmax, kernel_size=kernel_size, axis=2)337 if ndim==2:338 del hmax339 return hmax2340 else:341 hmax = max_pool1d(hmax2, kernel_size=kernel_size, axis=3, out=hmax)342 del hmax2 343 return hmax344 345def get_masks_torch(pt, inds, shape0, rpad=20, max_size_fraction=0.4):346 """Create masks using pixel convergence after running dynamics.347 348 Makes a histogram of final pixel locations p, initializes masks 349 at peaks of histogram and extends the masks from the peaks so that350 they include all pixels with more than 2 final pixels p. Discards 351 masks with flow errors greater than the threshold. 352 353 Parameters:354 p (float32, 3D or 4D array): Final locations of each pixel after dynamics,355 size [axis x Ly x Lx] or [axis x Lz x Ly x Lx].356 iscell (bool, 2D or 3D array): If iscell is not None, set pixels that are 357 iscell False to stay in their original location.358 rpad (int, optional): Histogram edge padding. Default is 20.359 max_size_fraction (float, optional): Masks larger than max_size_fraction of360 total image size are removed. Default is 0.4.361 362 Returns:363 M0 (int, 2D or 3D array): Masks with inconsistent flow masks removed, 364 0=NO masks; 1,2,...=mask labels, size [Ly x Lx] or [Lz x Ly x Lx].365 """366 367 ndim = len(shape0)368 device = pt.device369 370 rpad = 20371 pt += rpad372 pt = torch.clamp(pt, min=0)373 for i in range(len(pt)):374 pt[i] = torch.clamp(pt[i], max=shape0[i]+rpad-1)375 376 # # add extra padding to make divisible by 5377 # shape = tuple((np.ceil((shape0 + 2*rpad)/5) * 5).astype(int))378 shape = tuple(np.array(shape0) + 2*rpad)379 380 # sparse coo torch381 coo = torch.sparse_coo_tensor(pt, torch.ones(pt.shape[1], device=pt.device, dtype=torch.int), 382 shape)383 h1 = coo.to_dense()384 del coo385 386 hmax1 = max_pool_nd(h1.unsqueeze(0), kernel_size=5)387 hmax1 = hmax1.squeeze()388 seeds1 = torch.nonzero((h1 - hmax1 > -1e-6) * (h1 > 10))389 del hmax1390 if len(seeds1) == 0:391 dynamics_logger.warning("no seeds found in get_masks_torch - no masks found.")392 return np.zeros(shape0, dtype="uint16")393 394 npts = h1[tuple(seeds1.T)]395 isort1 = npts.argsort()396 seeds1 = seeds1[isort1]397 398 n_seeds = len(seeds1)399 h_slc = torch.zeros((n_seeds, *[11]*ndim), device=seeds1.device)400 for k in range(n_seeds):401 slc = tuple([slice(seeds1[k][j]-5, seeds1[k][j]+6) for j in range(ndim)])402 h_slc[k] = h1[slc]403 del h1404 seed_masks = torch.zeros((n_seeds, *[11]*ndim), device=seeds1.device)405 if ndim==2:406 seed_masks[:,5,5] = 1407 else:408 seed_masks[:,5,5,5] = 1409 410 for iter in range(5):411 seed_masks = max_pool_nd(seed_masks, kernel_size=3)412 seed_masks *= h_slc > 2413 del h_slc 414 seeds_new = [tuple((torch.nonzero(seed_masks[k]) + seeds1[k] - 5).T) 415 for k in range(n_seeds)]416 del seed_masks 417 418 dtype = torch.int32 if n_seeds < 2**16 else torch.int64419 M1 = torch.zeros(shape, dtype=dtype, device=device)420 for k in range(n_seeds):421 M1[seeds_new[k]] = 1 + k422 423 M1 = M1[tuple(pt)]424 M1 = M1.cpu().numpy()425 426 dtype = "uint16" if n_seeds < 2**16 else "uint32"427 M0 = np.zeros(shape0, dtype=dtype)428 M0[inds] = M1429 430 # remove big masks431 uniq, counts = fastremap.unique(M0, return_counts=True)432 big = np.prod(shape0) * max_size_fraction433 bigc = uniq[counts > big]434 if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):435 M0 = fastremap.mask(M0, bigc)436 fastremap.renumber(M0, in_place=True) #convenient to guarantee non-skipped labels437 M0 = M0.reshape(tuple(shape0))438 439 return M0440 441 442def resize_and_compute_masks(dP, cellprob, niter=200, cellprob_threshold=0.0,443 flow_threshold=0.4, do_3D=False, min_size=15,444 max_size_fraction=0.4, resize=None, device=torch.device("cpu")):445 """Compute masks using dynamics from dP and cellprob, and resizes masks if resize is not None.446 447 Args:448 dP (numpy.ndarray): The dynamics flow field array.449 cellprob (numpy.ndarray): The cell probability array.450 p (numpy.ndarray, optional): The pixels on which to run dynamics. Defaults to None451 niter (int, optional): The number of iterations for mask computation. Defaults to 200.452 cellprob_threshold (float, optional): The threshold for cell probability. Defaults to 0.0.453 flow_threshold (float, optional): The threshold for quality control metrics. Defaults to 0.4.454 interp (bool, optional): Whether to interpolate during dynamics computation. Defaults to True.455 do_3D (bool, optional): Whether to perform mask computation in 3D. Defaults to False.456 min_size (int, optional): The minimum size of the masks. Defaults to 15.457 max_size_fraction (float, optional): Masks larger than max_size_fraction of458 total image size are removed. Default is 0.4.459 resize (tuple, optional): The desired size for resizing the masks. Defaults to None.460 device (torch.device, optional): The device to use for computation. Defaults to torch.device("cpu").461 462 Returns:463 tuple: A tuple containing the computed masks and the final pixel locations.464 """465 mask = compute_masks(dP, cellprob, niter=niter,466 cellprob_threshold=cellprob_threshold,467 flow_threshold=flow_threshold, do_3D=do_3D,468 max_size_fraction=max_size_fraction, 469 device=device)470 471 if resize is not None:472 dynamics_logger.warning("Resizing is depricated in v4.0.1+")473 474 mask = utils.fill_holes_and_remove_small_masks(mask, min_size=min_size)475 476 return mask477 478 479def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,480 flow_threshold=0.4, do_3D=False, min_size=-1,481 max_size_fraction=0.4, device=torch.device("cpu")):482 """Compute masks using dynamics from dP and cellprob.483 484 Args:485 dP (numpy.ndarray): The dynamics flow field array.486 cellprob (numpy.ndarray): The cell probability array.487 p (numpy.ndarray, optional): The pixels on which to run dynamics. Defaults to None488 niter (int, optional): The number of iterations for mask computation. Defaults to 200.489 cellprob_threshold (float, optional): The threshold for cell probability. Defaults to 0.0.490 flow_threshold (float, optional): The threshold for quality control metrics. Defaults to 0.4.491 interp (bool, optional): Whether to interpolate during dynamics computation. Defaults to True.492 do_3D (bool, optional): Whether to perform mask computation in 3D. Defaults to False.493 min_size (int, optional): The minimum size of the masks. Defaults to 15.494 max_size_fraction (float, optional): Masks larger than max_size_fraction of495 total image size are removed. Default is 0.4.496 device (torch.device, optional): The device to use for computation. Defaults to torch.device("cpu").497 498 Returns:499 tuple: A tuple containing the computed masks and the final pixel locations.500 """501 502 if (cellprob > cellprob_threshold).sum(): #mask at this point is a cell cluster binary map, not labels503 inds = np.nonzero(cellprob > cellprob_threshold)504 if len(inds[0]) == 0:505 dynamics_logger.info("No cell pixels found.")506 shape = cellprob.shape507 mask = np.zeros(shape, "uint16")508 return mask509 510 p_final = steps_interp(dP * (cellprob > cellprob_threshold) / 5., 511 inds=inds, niter=niter, 512 device=device)513 if not torch.is_tensor(p_final):514 p_final = torch.from_numpy(p_final).to(device, dtype=torch.int)515 else:516 p_final = p_final.int()517 # calculate masks518 if device.type == "mps":519 p_final = p_final.to(torch.device("cpu"))520 mask = get_masks_torch(p_final, inds, dP.shape[1:], 521 max_size_fraction=max_size_fraction)522 del p_final523 # flow thresholding factored out of get_masks524 if not do_3D:525 if mask.max() > 0 and flow_threshold is not None and flow_threshold > 0:526 # make sure labels are unique at output of get_masks527 mask = remove_bad_flow_masks(mask, dP, threshold=flow_threshold,528 device=device)529 530 if mask.max() < 2**16 and mask.dtype != "uint16":531 mask = mask.astype("uint16")532 533 else: # nothing to compute, just make it compatible534 dynamics_logger.info("No cell pixels found.")535 shape = cellprob.shape536 mask = np.zeros(cellprob.shape, "uint16")537 return mask538 539 if min_size > 0:540 mask = utils.fill_holes_and_remove_small_masks(mask, min_size=min_size)541 542 if mask.dtype == np.uint32:543 dynamics_logger.warning(544 "more than 65535 masks in image, masks returned as np.uint32")545 546 return mask547