VisionLanguageGroup/MicroscopyMatching
0
1"""2Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer, Michael Rariden and Marius Pachitariu.3"""4 5import os, time6# from pathlib import Path7import numpy as np8from tqdm import trange9import torch10from scipy.ndimage import gaussian_filter11# import gc12import cv213 14import logging15 16models_logger = logging.getLogger(__name__)17 18from . import transforms, dynamics, utils19from .vit import Transformer20from .core import assign_device, run_net21 22# _MODEL_DIR_ENV = os.environ.get("CELLPOSE_LOCAL_MODELS_PATH")23# _MODEL_DIR_DEFAULT = Path("/media/data1/huix/seg/cellpose_models")24# MODEL_DIR = Path(_MODEL_DIR_ENV) if _MODEL_DIR_ENV else _MODEL_DIR_DEFAULT25 26# MODEL_NAMES = ["cpsam"]27 28# MODEL_LIST_PATH = os.fspath(MODEL_DIR.joinpath("gui_models.txt"))29 30normalize_default = {31 "lowhigh": None,32 "percentile": None,33 "normalize": True,34 "norm3D": True,35 "sharpen_radius": 0,36 "smooth_radius": 0,37 "tile_norm_blocksize": 0,38 "tile_norm_smooth3D": 1,39 "invert": False40}41 42 43# def get_user_models():44# model_strings = []45# if os.path.exists(MODEL_LIST_PATH):46# with open(MODEL_LIST_PATH, "r") as textfile:47# lines = [line.rstrip() for line in textfile]48# if len(lines) > 0:49# model_strings.extend(lines)50# return model_strings51 52 53class SegModel():54 """55 Class representing a Cellpose model.56 57 Attributes:58 diam_mean (float): Mean "diameter" value for the model.59 builtin (bool): Whether the model is a built-in model or not.60 device (torch device): Device used for model running / training.61 nclasses (int): Number of classes in the model.62 nbase (list): List of base values for the model.63 net (CPnet): Cellpose network.64 pretrained_model (str): Path to pretrained cellpose model.65 pretrained_model_ortho (str): Path or model_name for pretrained cellpose model for ortho views in 3D.66 backbone (str): Type of network ("default" is the standard res-unet, "transformer" for the segformer).67 68 Methods:69 __init__(self, gpu=False, pretrained_model=False, model_type=None, diam_mean=30., device=None):70 Initialize the CellposeModel.71 72 eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None, z_axis=None, normalize=True, invert=False, rescale=None, diameter=None, flow_threshold=0.4, cellprob_threshold=0.0, do_3D=False, anisotropy=None, stitch_threshold=0.0, min_size=15, niter=None, augment=False, tile_overlap=0.1, bsize=224, interp=True, compute_masks=True, progress=None):73 Segment list of images x, or 4D array - Z x C x Y x X.74 75 """76 77 def __init__(self, gpu=False, pretrained_model="", model_type=None,78 diam_mean=None, device=None, nchan=None, use_bfloat16=True, vit_checkpoint=None):79 """80 Initialize the CellposeModel.81 82 Parameters:83 gpu (bool, optional): Whether or not to save model to GPU, will check if GPU available.84 pretrained_model (str or list of strings, optional): Full path to pretrained cellpose model(s), if None or False, no model loaded.85 model_type (str, optional): Any model that is available in the GUI, use name in GUI e.g. "livecell" (can be user-trained or model zoo).86 diam_mean (float, optional): Mean "diameter", 30. is built-in value for "cyto" model; 17. is built-in value for "nuclei" model; if saved in custom model file (cellpose>=2.0) then it will be loaded automatically and overwrite this value.87 device (torch device, optional): Device used for model running / training (torch.device("cuda") or torch.device("cpu")), overrides gpu input, recommended if you want to use a specific GPU (e.g. torch.device("cuda:1")).88 use_bfloat16 (bool, optional): Use 16bit float precision instead of 32bit for model weights. Default to 16bit (True).89 """90 ### assign model device91 self.device = assign_device(gpu=gpu)[0] if device is None else device92 if torch.cuda.is_available():93 device_gpu = self.device.type == "cuda"94 elif torch.backends.mps.is_available():95 device_gpu = self.device.type == "mps"96 else:97 device_gpu = False98 self.gpu = device_gpu99 100 if pretrained_model is None:101 # raise ValueError("Must specify a pretrained model, training from scratch is not implemented")102 pretrained_model = ""103 104 self.pretrained_model = pretrained_model105 dtype = torch.bfloat16 if use_bfloat16 else torch.float32106 self.net = Transformer(dtype=dtype, checkpoint=vit_checkpoint).to(self.device)107 108 109 def eval(self, x, feat=None, batch_size=8, resample=True, channels=None, channel_axis=None,110 z_axis=None, normalize=True, invert=False, rescale=None, diameter=None,111 flow_threshold=0.4, cellprob_threshold=0.0, do_3D=False, anisotropy=None,112 flow3D_smooth=0, stitch_threshold=0.0, 113 min_size=15, max_size_fraction=0.4, niter=None, 114 augment=False, tile_overlap=0.1, bsize=256, 115 compute_masks=True, progress=None):116 117 if isinstance(x, list) or x.squeeze().ndim == 5:118 self.timing = []119 masks, styles, flows = [], [], []120 tqdm_out = utils.TqdmToLogger(models_logger, level=logging.INFO)121 nimg = len(x)122 iterator = trange(nimg, file=tqdm_out,123 mininterval=30) if nimg > 1 else range(nimg)124 for i in iterator:125 tic = time.time()126 maski, flowi, stylei = self.eval(127 x[i], 128 feat=None if feat is None else feat[i],129 batch_size=batch_size,130 channel_axis=channel_axis, 131 z_axis=z_axis,132 normalize=normalize, 133 invert=invert,134 diameter=diameter[i] if isinstance(diameter, list) or135 isinstance(diameter, np.ndarray) else diameter, 136 do_3D=do_3D,137 anisotropy=anisotropy, 138 augment=augment, 139 tile_overlap=tile_overlap, 140 bsize=bsize, 141 resample=resample,142 flow_threshold=flow_threshold,143 cellprob_threshold=cellprob_threshold, 144 compute_masks=compute_masks,145 min_size=min_size, 146 max_size_fraction=max_size_fraction, 147 stitch_threshold=stitch_threshold, 148 flow3D_smooth=flow3D_smooth,149 progress=progress, 150 niter=niter)151 masks.append(maski)152 flows.append(flowi)153 styles.append(stylei)154 self.timing.append(time.time() - tic)155 return masks, flows, styles156 157 ############# actual eval code ############158 # reshape image159 x = transforms.convert_image(x, channel_axis=channel_axis,160 z_axis=z_axis, 161 do_3D=(do_3D or stitch_threshold > 0))162 163 # Add batch dimension if not present164 if x.ndim < 4:165 x = x[np.newaxis, ...]166 if feat is not None:167 if feat.ndim < 4:168 feat = feat[np.newaxis, ...]169 nimg = x.shape[0]170 171 image_scaling = None172 Ly_0 = x.shape[1]173 Lx_0 = x.shape[2]174 Lz_0 = None175 if stitch_threshold > 0:176 Lz_0 = x.shape[0]177 if diameter is not None:178 image_scaling = 30. / diameter179 x = transforms.resize_image(x,180 Ly=int(x.shape[1] * image_scaling),181 Lx=int(x.shape[2] * image_scaling))182 if feat is not None:183 feat = transforms.resize_image(feat,184 Ly=int(feat.shape[1] * image_scaling),185 Lx=int(feat.shape[2] * image_scaling))186 187 188 # normalize image189 normalize_params = normalize_default190 if isinstance(normalize, dict):191 normalize_params = {**normalize_params, **normalize}192 elif not isinstance(normalize, bool):193 raise ValueError("normalize parameter must be a bool or a dict")194 else:195 normalize_params["normalize"] = normalize196 normalize_params["invert"] = invert197 198 # pre-normalize if 3D stack for stitching or do_3D199 do_normalization = True if normalize_params["normalize"] else False200 if nimg > 1 and do_normalization and (stitch_threshold or do_3D):201 normalize_params["norm3D"] = True if do_3D else normalize_params["norm3D"]202 x = transforms.normalize_img(x, **normalize_params)203 do_normalization = False # do not normalize again204 else:205 if normalize_params["norm3D"] and nimg > 1 and do_normalization:206 models_logger.warning(207 "normalize_params['norm3D'] is True but do_3D is False and stitch_threshold=0, so setting to False"208 )209 normalize_params["norm3D"] = False210 if do_normalization:211 x = transforms.normalize_img(x, **normalize_params)212 213 if feat is not None:214 if feat.shape[-1] > feat.shape[1]:215 # transpose feat to have channels last216 feat = np.moveaxis(feat, 1, -1)217 218 # adjust the anisotropy when diameter is specified and images are resized:219 if isinstance(anisotropy, (float, int)) and image_scaling:220 anisotropy = image_scaling * anisotropy221 222 dP, cellprob, styles = self._run_net(223 x, 224 feat=feat,225 augment=augment, 226 batch_size=batch_size, 227 tile_overlap=tile_overlap, 228 bsize=bsize,229 do_3D=do_3D, 230 anisotropy=anisotropy)231 232 233 if resample:234 # upsample flows before computing them: 235 dP = self._resize_gradients(dP, to_y_size=Ly_0, to_x_size=Lx_0, to_z_size=Lz_0)236 cellprob = self._resize_cellprob(cellprob, to_x_size=Lx_0, to_y_size=Ly_0, to_z_size=Lz_0)237 238 239 if compute_masks:240 niter0 = 200241 niter = niter0 if niter is None or niter == 0 else niter242 masks = self._compute_masks(x.shape, dP, cellprob, flow_threshold=flow_threshold,243 cellprob_threshold=cellprob_threshold, min_size=min_size,244 max_size_fraction=max_size_fraction, niter=niter,245 stitch_threshold=stitch_threshold, do_3D=do_3D)246 else:247 masks = np.zeros(0) #pass back zeros if not compute_masks248 249 masks = masks.squeeze()250 251 # undo resizing:252 if image_scaling is not None or anisotropy is not None:253 254 if compute_masks:255 masks = transforms.resize_image(masks, Ly=Ly_0, Lx=Lx_0, no_channels=True, interpolation=cv2.INTER_NEAREST)256 257 return masks258 259 260 261 def _resize_cellprob(self, prob: np.ndarray, to_y_size: int, to_x_size: int, to_z_size: int = None) -> np.ndarray:262 """263 Resize cellprob array to specified dimensions for either 2D or 3D.264 265 Parameters:266 prob (numpy.ndarray): The cellprobs to resize, either in 2D or 3D. Returns the same ndim as provided.267 to_y_size (int): The target size along the Y-axis.268 to_x_size (int): The target size along the X-axis.269 to_z_size (int, optional): The target size along the Z-axis. Required270 for 3D cellprobs.271 272 Returns:273 numpy.ndarray: The resized cellprobs array with the same number of dimensions274 as the input.275 276 Raises:277 ValueError: If the input cellprobs array does not have 3 or 4 dimensions.278 """279 prob_shape = prob.shape280 prob = prob.squeeze()281 squeeze_happened = prob.shape != prob_shape282 prob_shape = np.array(prob_shape)283 284 if prob.ndim == 2:285 # 2D case:286 prob = transforms.resize_image(prob, Ly=to_y_size, Lx=to_x_size, no_channels=True)287 if squeeze_happened:288 prob = np.expand_dims(prob, int(np.argwhere(prob_shape == 1))) # add back empty axis for compatibility289 elif prob.ndim == 3:290 # 3D case: 291 prob = transforms.resize_image(prob, Ly=to_y_size, Lx=to_x_size, no_channels=True)292 prob = prob.transpose(1, 0, 2)293 prob = transforms.resize_image(prob, Ly=to_z_size, Lx=to_x_size, no_channels=True)294 prob = prob.transpose(1, 0, 2)295 else:296 raise ValueError(f'gradients have incorrect dimension after squeezing. Should be 2 or 3, prob shape: {prob.shape}')297 298 return prob299 300 301 def _resize_gradients(self, grads: np.ndarray, to_y_size: int, to_x_size: int, to_z_size: int = None) -> np.ndarray:302 """303 Resize gradient arrays to specified dimensions for either 2D or 3D gradients.304 305 Parameters:306 grads (np.ndarray): The gradients to resize, either in 2D or 3D. Returns the same ndim as provided.307 to_y_size (int): The target size along the Y-axis.308 to_x_size (int): The target size along the X-axis.309 to_z_size (int, optional): The target size along the Z-axis. Required310 for 3D gradients.311 312 Returns:313 numpy.ndarray: The resized gradient array with the same number of dimensions314 as the input.315 316 Raises:317 ValueError: If the input gradient array does not have 3 or 4 dimensions.318 """319 grads_shape = grads.shape320 grads = grads.squeeze()321 squeeze_happened = grads.shape != grads_shape322 grads_shape = np.array(grads_shape)323 324 if grads.ndim == 3:325 # 2D case, with XY flows in 2 channels:326 grads = np.moveaxis(grads, 0, -1) # Put gradients last327 grads = transforms.resize_image(grads, Ly=to_y_size, Lx=to_x_size, no_channels=False)328 grads = np.moveaxis(grads, -1, 0) # Put gradients first329 330 if squeeze_happened:331 grads = np.expand_dims(grads, int(np.argwhere(grads_shape == 1))) # add back empty axis for compatibility332 elif grads.ndim == 4:333 # dP has gradients that can be treated as channels:334 grads = grads.transpose(1, 2, 3, 0) # move gradients last:335 grads = transforms.resize_image(grads, Ly=to_y_size, Lx=to_x_size, no_channels=False)336 grads = grads.transpose(1, 0, 2, 3) # switch axes to resize again337 grads = transforms.resize_image(grads, Ly=to_z_size, Lx=to_x_size, no_channels=False)338 grads = grads.transpose(3, 1, 0, 2) # undo transposition339 else:340 raise ValueError(f'gradients have incorrect dimension after squeezing. Should be 3 or 4, grads shape: {grads.shape}')341 342 return grads343 344 345 def _run_net(self, x, feat=None,346 augment=False, 347 batch_size=8, tile_overlap=0.1,348 bsize=224, anisotropy=1.0, do_3D=False):349 """ run network on image x """350 tic = time.time()351 shape = x.shape352 nimg = shape[0]353 354 355 356 yf, styles = run_net(self.net, x, feat=feat, bsize=bsize, augment=augment,357 batch_size=batch_size, 358 tile_overlap=tile_overlap, 359 )360 cellprob = yf[..., -1]361 dP = yf[..., -3:-1].transpose((3, 0, 1, 2))362 if yf.shape[-1] > 3:363 styles = yf[..., :-3]364 365 styles = styles.squeeze()366 367 net_time = time.time() - tic368 if nimg > 1:369 models_logger.info("network run in %2.2fs" % (net_time))370 371 return dP, cellprob, styles372 373 def _compute_masks(self, shape, dP, cellprob, flow_threshold=0.4, cellprob_threshold=0.0,374 min_size=15, max_size_fraction=0.4, niter=None,375 do_3D=False, stitch_threshold=0.0):376 """ compute masks from flows and cell probability """377 changed_device_from = None378 if self.device.type == "mps" and do_3D:379 models_logger.warning("MPS does not support 3D post-processing, switching to CPU")380 self.device = torch.device("cpu")381 changed_device_from = "mps"382 Lz, Ly, Lx = shape[:3]383 tic = time.time()384 # if do_3D:385 # masks = dynamics.resize_and_compute_masks(386 # dP, cellprob, niter=niter, cellprob_threshold=cellprob_threshold,387 # flow_threshold=flow_threshold, do_3D=do_3D,388 # min_size=min_size, max_size_fraction=max_size_fraction, 389 # resize=shape[:3] if (np.array(dP.shape[-3:])!=np.array(shape[:3])).sum() 390 # else None,391 # device=self.device)392 # else:393 nimg = shape[0]394 Ly0, Lx0 = cellprob[0].shape 395 resize = None if Ly0==Ly and Lx0==Lx else [Ly, Lx]396 tqdm_out = utils.TqdmToLogger(models_logger, level=logging.INFO)397 iterator = trange(nimg, file=tqdm_out,398 mininterval=30) if nimg > 1 else range(nimg)399 for i in iterator:400 # turn off min_size for 3D stitching401 min_size0 = min_size if stitch_threshold == 0 or nimg == 1 else -1402 outputs = dynamics.resize_and_compute_masks(403 dP[:, i], cellprob[i],404 niter=niter, cellprob_threshold=cellprob_threshold,405 flow_threshold=flow_threshold, resize=resize,406 min_size=min_size0, max_size_fraction=max_size_fraction,407 device=self.device)408 if i==0 and nimg > 1:409 masks = np.zeros((nimg, shape[1], shape[2]), outputs.dtype)410 if nimg > 1:411 masks[i] = outputs412 else:413 masks = outputs414 415 if stitch_threshold > 0 and nimg > 1:416 models_logger.info(417 f"stitching {nimg} planes using stitch_threshold={stitch_threshold:0.3f} to make 3D masks"418 )419 masks = utils.stitch3D(masks, stitch_threshold=stitch_threshold)420 masks = utils.fill_holes_and_remove_small_masks(421 masks, min_size=min_size)422 elif nimg > 1:423 models_logger.warning(424 "3D stack used, but stitch_threshold=0 and do_3D=False, so masks are made per plane only"425 )426 427 flow_time = time.time() - tic428 if shape[0] > 1:429 models_logger.info("masks created in %2.2fs" % (flow_time))430 431 if changed_device_from is not None:432 models_logger.info("switching back to device %s" % self.device)433 self.device = torch.device(changed_device_from)434 return masks435