saim1309/Cell_Segmentation
0
1"""2Copyright © 2022 Howard Hughes Medical Institute, 3Authored by Carsen Stringer and Marius Pachitariu.4 5Redistribution and use in source and binary forms, with or without 6modification, are permitted provided that the following conditions are met:7 81. Redistributions of source code must retain the above copyright notice, 9 this list of conditions and the following disclaimer.10 112. Redistributions in binary form must reproduce the above copyright notice, 12 this list of conditions and the following disclaimer in the documentation 13 and/or other materials provided with the distribution.14 153. Neither the name of HHMI nor the names of its contributors may be used to 16 endorse or promote products derived from this software without specific 17 prior written permission.18 19THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE21IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 22ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 23LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 24CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 25SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 26INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 27CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 28ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 29POSSIBILITY OF SUCH DAMAGE.30 31--------------------------------------------------------------------------32MEDIAR Prediction uses CellPose's Gradient Flow Tracking.33 34This code is adapted from the following codes:35[1] https://github.com/MouseLand/cellpose/blob/main/cellpose/utils.py36[2] https://github.com/MouseLand/cellpose/blob/main/cellpose/dynamics.py37[3] https://github.com/MouseLand/cellpose/blob/main/cellpose/metrics.py38"""39 40import torch41from torch.nn.functional import grid_sample42import numpy as np43import fastremap44 45from skimage import morphology46from scipy.ndimage import mean, find_objects47from scipy.ndimage.filters import maximum_filter1d48 49torch_GPU = torch.device("cuda")50torch_CPU = torch.device("cpu")51 52 53def labels_to_flows(labels, use_gpu=False, device=None, redo_flows=False):54 """55 Convert labels (list of masks or flows) to flows for training model56 """57 58 # Labels b x 1 x h x w59 labels = labels.cpu().numpy().astype(np.int16)60 nimg = len(labels)61 62 if labels[0].ndim < 3:63 labels = [labels[n][np.newaxis, :, :] for n in range(nimg)]64 65 # Flows need to be recomputed66 if labels[0].shape[0] == 1 or labels[0].ndim < 3 or redo_flows:67 # compute flows; labels are fixed here to be unique, so they need to be passed back68 # make sure labels are unique!69 labels = [fastremap.renumber(label, in_place=True)[0] for label in labels]70 veci = [71 masks_to_flows(labels[n][0], use_gpu=use_gpu, device=device)72 for n in range(nimg)73 ]74 75 # concatenate labels, distance transform, vector flows, heat (boundary and mask are computed in augmentations)76 flows = [77 np.concatenate((labels[n], labels[n] > 0.5, veci[n]), axis=0).astype(78 np.float3279 )80 for n in range(nimg)81 ]82 83 return np.array(flows)84 85 86def compute_masks(87 dP,88 cellprob,89 p=None,90 niter=200,91 cellprob_threshold=0.4,92 flow_threshold=0.4,93 interp=True,94 resize=None,95 use_gpu=False,96 device=None,97):98 """compute masks using dynamics from dP, cellprob, and boundary"""99 100 cp_mask = cellprob > cellprob_threshold101 cp_mask = morphology.remove_small_holes(cp_mask, area_threshold=16)102 cp_mask = morphology.remove_small_objects(cp_mask, min_size=16)103 104 if np.any(cp_mask): # mask at this point is a cell cluster binary map, not labels105 # follow flows106 if p is None:107 p, inds = follow_flows(108 dP * cp_mask / 5.0,109 niter=niter,110 interp=interp,111 use_gpu=use_gpu,112 device=device,113 )114 if inds is None:115 shape = resize if resize is not None else cellprob.shape116 mask = np.zeros(shape, np.uint16)117 p = np.zeros((len(shape), *shape), np.uint16)118 return mask, p119 120 # calculate masks121 mask = get_masks(p, iscell=cp_mask)122 123 # flow thresholding factored out of get_masks124 shape0 = p.shape[1:]125 if mask.max() > 0 and flow_threshold is not None and flow_threshold > 0:126 # make sure labels are unique at output of get_masks127 mask = remove_bad_flow_masks(128 mask, dP, threshold=flow_threshold, use_gpu=use_gpu, device=device129 )130 else: # nothing to compute, just make it compatible131 shape = resize if resize is not None else cellprob.shape132 mask = np.zeros(shape, np.uint16)133 p = np.zeros((len(shape), *shape), np.uint16)134 135 return mask, p136 137 138def _extend_centers_gpu(139 neighbors, centers, isneighbor, Ly, Lx, n_iter=200, device=torch.device("cuda")140):141 if device is not None:142 device = device143 nimg = neighbors.shape[0] // 9144 pt = torch.from_numpy(neighbors).to(device)145 146 T = torch.zeros((nimg, Ly, Lx), dtype=torch.double, device=device)147 meds = torch.from_numpy(centers.astype(int)).to(device).long()148 isneigh = torch.from_numpy(isneighbor).to(device)149 for i in range(n_iter):150 T[:, meds[:, 0], meds[:, 1]] += 1151 Tneigh = T[:, pt[:, :, 0], pt[:, :, 1]]152 Tneigh *= isneigh153 T[:, pt[0, :, 0], pt[0, :, 1]] = Tneigh.mean(axis=1)154 del meds, isneigh, Tneigh155 T = torch.log(1.0 + T)156 # gradient positions157 grads = T[:, pt[[2, 1, 4, 3], :, 0], pt[[2, 1, 4, 3], :, 1]]158 del pt159 dy = grads[:, 0] - grads[:, 1]160 dx = grads[:, 2] - grads[:, 3]161 del grads162 mu_torch = np.stack((dy.cpu().squeeze(), dx.cpu().squeeze()), axis=-2)163 return mu_torch164 165 166def diameters(masks):167 _, counts = np.unique(np.int32(masks), return_counts=True)168 counts = counts[1:]169 md = np.median(counts ** 0.5)170 if np.isnan(md):171 md = 0172 md /= (np.pi ** 0.5) / 2173 return md, counts ** 0.5174 175 176def masks_to_flows_gpu(masks, device=None):177 if device is None:178 device = torch.device("cuda")179 180 Ly0, Lx0 = masks.shape181 Ly, Lx = Ly0 + 2, Lx0 + 2182 183 masks_padded = np.zeros((Ly, Lx), np.int64)184 masks_padded[1:-1, 1:-1] = masks185 186 # get mask pixel neighbors187 y, x = np.nonzero(masks_padded)188 neighborsY = np.stack((y, y - 1, y + 1, y, y, y - 1, y - 1, y + 1, y + 1), axis=0)189 neighborsX = np.stack((x, x, x, x - 1, x + 1, x - 1, x + 1, x - 1, x + 1), axis=0)190 neighbors = np.stack((neighborsY, neighborsX), axis=-1)191 192 # get mask centers193 slices = find_objects(masks)194 195 centers = np.zeros((masks.max(), 2), "int")196 for i, si in enumerate(slices):197 if si is not None:198 sr, sc = si199 200 ly, lx = sr.stop - sr.start + 1, sc.stop - sc.start + 1201 yi, xi = np.nonzero(masks[sr, sc] == (i + 1))202 yi = yi.astype(np.int32) + 1 # add padding203 xi = xi.astype(np.int32) + 1 # add padding204 ymed = np.median(yi)205 xmed = np.median(xi)206 imin = np.argmin((xi - xmed) ** 2 + (yi - ymed) ** 2)207 xmed = xi[imin]208 ymed = yi[imin]209 centers[i, 0] = ymed + sr.start210 centers[i, 1] = xmed + sc.start211 212 # get neighbor validator (not all neighbors are in same mask)213 neighbor_masks = masks_padded[neighbors[:, :, 0], neighbors[:, :, 1]]214 isneighbor = neighbor_masks == neighbor_masks[0]215 ext = np.array(216 [[sr.stop - sr.start + 1, sc.stop - sc.start + 1] for sr, sc in slices]217 )218 n_iter = 2 * (ext.sum(axis=1)).max()219 # run diffusion220 mu = _extend_centers_gpu(221 neighbors, centers, isneighbor, Ly, Lx, n_iter=n_iter, device=device222 )223 224 # normalize225 mu /= 1e-20 + (mu ** 2).sum(axis=0) ** 0.5226 227 # put into original image228 mu0 = np.zeros((2, Ly0, Lx0))229 mu0[:, y - 1, x - 1] = mu230 mu_c = np.zeros_like(mu0)231 return mu0, mu_c232 233 234def masks_to_flows(masks, use_gpu=False, device=None):235 if masks.max() == 0 or (masks != 0).sum() == 1:236 # dynamics_logger.warning('empty masks!')237 return np.zeros((2, *masks.shape), "float32")238 239 if use_gpu:240 if use_gpu and device is None:241 device = torch_GPU242 elif device is None:243 device = torch_CPU244 masks_to_flows_device = masks_to_flows_gpu245 246 if masks.ndim == 3:247 Lz, Ly, Lx = masks.shape248 mu = np.zeros((3, Lz, Ly, Lx), np.float32)249 for z in range(Lz):250 mu0 = masks_to_flows_device(masks[z], device=device)[0]251 mu[[1, 2], z] += mu0252 for y in range(Ly):253 mu0 = masks_to_flows_device(masks[:, y], device=device)[0]254 mu[[0, 2], :, y] += mu0255 for x in range(Lx):256 mu0 = masks_to_flows_device(masks[:, :, x], device=device)[0]257 mu[[0, 1], :, :, x] += mu0258 return mu259 elif masks.ndim == 2:260 mu, mu_c = masks_to_flows_device(masks, device=device)261 return mu262 263 else:264 raise ValueError("masks_to_flows only takes 2D or 3D arrays")265 266 267def steps2D_interp(p, dP, niter, use_gpu=False, device=None):268 shape = dP.shape[1:]269 if use_gpu:270 if device is None:271 device = torch_GPU272 shape = (273 np.array(shape)[[1, 0]].astype("float") - 1274 ) # Y and X dimensions (dP is 2.Ly.Lx), flipped X-1, Y-1275 pt = (276 torch.from_numpy(p[[1, 0]].T).float().to(device).unsqueeze(0).unsqueeze(0)277 ) # p is n_points by 2, so pt is [1 1 2 n_points]278 im = (279 torch.from_numpy(dP[[1, 0]]).float().to(device).unsqueeze(0)280 ) # covert flow numpy array to tensor on GPU, add dimension281 # normalize pt between 0 and 1, normalize the flow282 for k in range(2):283 im[:, k, :, :] *= 2.0 / shape[k]284 pt[:, :, :, k] /= shape[k]285 286 # normalize to between -1 and 1287 pt = pt * 2 - 1288 289 # here is where the stepping happens290 for t in range(niter):291 # align_corners default is False, just added to suppress warning292 dPt = grid_sample(im, pt, align_corners=False)293 294 for k in range(2): # clamp the final pixel locations295 pt[:, :, :, k] = torch.clamp(296 pt[:, :, :, k] + dPt[:, k, :, :], -1.0, 1.0297 )298 299 # undo the normalization from before, reverse order of operations300 pt = (pt + 1) * 0.5301 for k in range(2):302 pt[:, :, :, k] *= shape[k]303 304 p = pt[:, :, :, [1, 0]].cpu().numpy().squeeze().T305 return p306 307 else:308 assert print("ho")309 310 311def follow_flows(dP, mask=None, niter=200, interp=True, use_gpu=True, device=None):312 shape = np.array(dP.shape[1:]).astype(np.int32)313 niter = np.uint32(niter)314 315 p = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")316 p = np.array(p).astype(np.float32)317 318 inds = np.array(np.nonzero(np.abs(dP[0]) > 1e-3)).astype(np.int32).T319 320 if inds.ndim < 2 or inds.shape[0] < 5:321 return p, None322 323 if not interp:324 assert print("woo")325 326 else:327 p_interp = steps2D_interp(328 p[:, inds[:, 0], inds[:, 1]], dP, niter, use_gpu=use_gpu, device=device329 )330 p[:, inds[:, 0], inds[:, 1]] = p_interp331 332 return p, inds333 334 335def flow_error(maski, dP_net, use_gpu=False, device=None):336 if dP_net.shape[1:] != maski.shape:337 print("ERROR: net flow is not same size as predicted masks")338 return339 340 # flows predicted from estimated masks341 dP_masks = masks_to_flows(maski, use_gpu=use_gpu, device=device)342 # difference between predicted flows vs mask flows343 flow_errors = np.zeros(maski.max())344 for i in range(dP_masks.shape[0]):345 flow_errors += mean(346 (dP_masks[i] - dP_net[i] / 5.0) ** 2,347 maski,348 index=np.arange(1, maski.max() + 1),349 )350 351 return flow_errors, dP_masks352 353 354def remove_bad_flow_masks(masks, flows, threshold=0.4, use_gpu=False, device=None):355 merrors, _ = flow_error(masks, flows, use_gpu, device)356 badi = 1 + (merrors > threshold).nonzero()[0]357 masks[np.isin(masks, badi)] = 0358 return masks359 360 361def get_masks(p, iscell=None, rpad=20):362 pflows = []363 edges = []364 shape0 = p.shape[1:]365 dims = len(p)366 367 for i in range(dims):368 pflows.append(p[i].flatten().astype("int32"))369 edges.append(np.arange(-0.5 - rpad, shape0[i] + 0.5 + rpad, 1))370 371 h, _ = np.histogramdd(tuple(pflows), bins=edges)372 hmax = h.copy()373 for i in range(dims):374 hmax = maximum_filter1d(hmax, 5, axis=i)375 376 seeds = np.nonzero(np.logical_and(h - hmax > -1e-6, h > 10))377 Nmax = h[seeds]378 isort = np.argsort(Nmax)[::-1]379 for s in seeds:380 s = s[isort]381 382 pix = list(np.array(seeds).T)383 384 shape = h.shape385 if dims == 3:386 expand = np.nonzero(np.ones((3, 3, 3)))387 else:388 expand = np.nonzero(np.ones((3, 3)))389 for e in expand:390 e = np.expand_dims(e, 1)391 392 for iter in range(5):393 for k in range(len(pix)):394 if iter == 0:395 pix[k] = list(pix[k])396 newpix = []397 iin = []398 for i, e in enumerate(expand):399 epix = e[:, np.newaxis] + np.expand_dims(pix[k][i], 0) - 1400 epix = epix.flatten()401 iin.append(np.logical_and(epix >= 0, epix < shape[i]))402 newpix.append(epix)403 iin = np.all(tuple(iin), axis=0)404 for p in newpix:405 p = p[iin]406 newpix = tuple(newpix)407 igood = h[newpix] > 2408 for i in range(dims):409 pix[k][i] = newpix[i][igood]410 if iter == 4:411 pix[k] = tuple(pix[k])412 413 M = np.zeros(h.shape, np.uint32)414 for k in range(len(pix)):415 M[pix[k]] = 1 + k416 417 for i in range(dims):418 pflows[i] = pflows[i] + rpad419 M0 = M[tuple(pflows)]420 421 # remove big masks422 uniq, counts = fastremap.unique(M0, return_counts=True)423 big = np.prod(shape0) * 0.9424 bigc = uniq[counts > big]425 if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):426 M0 = fastremap.mask(M0, bigc)427 fastremap.renumber(M0, in_place=True) # convenient to guarantee non-skipped labels428 M0 = np.reshape(M0, shape0)429 return M0430