VisionLanguageGroup/MicroscopyMatching
0
1"""2Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer , Michael Rariden and Marius Pachitariu.3"""4import logging5 6import cv27import numpy as np8import torch9from scipy.ndimage import gaussian_filter1d10from torch.fft import fft2, fftshift, ifft211 12transforms_logger = logging.getLogger(__name__)13 14 15def _taper_mask(ly=224, lx=224, sig=7.5):16 """17 Generate a taper mask.18 19 Args:20 ly (int): The height of the mask. Default is 224.21 lx (int): The width of the mask. Default is 224.22 sig (float): The sigma value for the tapering function. Default is 7.5.23 24 Returns:25 numpy.ndarray: The taper mask.26 27 """28 bsize = max(224, max(ly, lx))29 xm = np.arange(bsize)30 xm = np.abs(xm - xm.mean())31 mask = 1 / (1 + np.exp((xm - (bsize / 2 - 20)) / sig))32 mask = mask * mask[:, np.newaxis]33 mask = mask[bsize // 2 - ly // 2:bsize // 2 + ly // 2 + ly % 2,34 bsize // 2 - lx // 2:bsize // 2 + lx // 2 + lx % 2]35 return mask36 37 38def unaugment_tiles(y):39 """Reverse test-time augmentations for averaging (includes flipping of flowsY and flowsX).40 41 Args:42 y (float32): Array of shape (ntiles_y, ntiles_x, chan, Ly, Lx) where chan = (flowsY, flowsX, cell prob).43 44 Returns:45 float32: Array of shape (ntiles_y, ntiles_x, chan, Ly, Lx).46 47 """48 for j in range(y.shape[0]):49 for i in range(y.shape[1]):50 if j % 2 == 0 and i % 2 == 1:51 y[j, i] = y[j, i, :, ::-1, :]52 y[j, i, 0] *= -153 elif j % 2 == 1 and i % 2 == 0:54 y[j, i] = y[j, i, :, :, ::-1]55 y[j, i, 1] *= -156 elif j % 2 == 1 and i % 2 == 1:57 y[j, i] = y[j, i, :, ::-1, ::-1]58 y[j, i, 0] *= -159 y[j, i, 1] *= -160 return y61 62 63def average_tiles(y, ysub, xsub, Ly, Lx):64 """65 Average the results of the network over tiles.66 67 Args:68 y (float): Output of cellpose network for each tile. Shape: [ntiles x nclasses x bsize x bsize]69 ysub (list): List of arrays with start and end of tiles in Y of length ntiles70 xsub (list): List of arrays with start and end of tiles in X of length ntiles71 Ly (int): Size of pre-tiled image in Y (may be larger than original image if image size is less than bsize)72 Lx (int): Size of pre-tiled image in X (may be larger than original image if image size is less than bsize)73 74 Returns:75 yf (float32): Network output averaged over tiles. Shape: [nclasses x Ly x Lx]76 """77 Navg = np.zeros((Ly, Lx))78 yf = np.zeros((y.shape[1], Ly, Lx), np.float32)79 # taper edges of tiles80 mask = _taper_mask(ly=y.shape[-2], lx=y.shape[-1])81 for j in range(len(ysub)):82 yf[:, ysub[j][0]:ysub[j][1], xsub[j][0]:xsub[j][1]] += y[j] * mask83 Navg[ysub[j][0]:ysub[j][1], xsub[j][0]:xsub[j][1]] += mask84 yf /= Navg85 return yf86 87 88def make_tiles(imgi, bsize=224, augment=False, tile_overlap=0.1):89 """Make tiles of image to run at test-time.90 91 Args:92 imgi (np.ndarray): Array of shape (nchan, Ly, Lx) representing the input image.93 bsize (int, optional): Size of tiles. Defaults to 224.94 augment (bool, optional): Whether to flip tiles and set tile_overlap=2. Defaults to False.95 tile_overlap (float, optional): Fraction of overlap of tiles. Defaults to 0.1.96 97 Returns:98 A tuple containing (IMG, ysub, xsub, Ly, Lx):99 IMG (np.ndarray): Array of shape (ntiles, nchan, bsize, bsize) representing the tiles.100 ysub (list): List of arrays with start and end of tiles in Y of length ntiles.101 xsub (list): List of arrays with start and end of tiles in X of length ntiles.102 Ly (int): Height of the input image.103 Lx (int): Width of the input image.104 """105 nchan, Ly, Lx = imgi.shape106 if augment:107 bsize = np.int32(bsize)108 # pad if image smaller than bsize109 if Ly < bsize:110 imgi = np.concatenate((imgi, np.zeros((nchan, bsize - Ly, Lx))), axis=1)111 Ly = bsize112 if Lx < bsize:113 imgi = np.concatenate((imgi, np.zeros((nchan, Ly, bsize - Lx))), axis=2)114 Ly, Lx = imgi.shape[-2:]115 116 # tiles overlap by half of tile size117 ny = max(2, int(np.ceil(2. * Ly / bsize)))118 nx = max(2, int(np.ceil(2. * Lx / bsize)))119 ystart = np.linspace(0, Ly - bsize, ny).astype(int)120 xstart = np.linspace(0, Lx - bsize, nx).astype(int)121 122 ysub = []123 xsub = []124 125 # flip tiles so that overlapping segments are processed in rotation126 IMG = np.zeros((len(ystart), len(xstart), nchan, bsize, bsize), np.float32)127 for j in range(len(ystart)):128 for i in range(len(xstart)):129 ysub.append([ystart[j], ystart[j] + bsize])130 xsub.append([xstart[i], xstart[i] + bsize])131 IMG[j, i] = imgi[:, ysub[-1][0]:ysub[-1][1], xsub[-1][0]:xsub[-1][1]]132 # flip tiles to allow for augmentation of overlapping segments133 if j % 2 == 0 and i % 2 == 1:134 IMG[j, i] = IMG[j, i, :, ::-1, :]135 elif j % 2 == 1 and i % 2 == 0:136 IMG[j, i] = IMG[j, i, :, :, ::-1]137 elif j % 2 == 1 and i % 2 == 1:138 IMG[j, i] = IMG[j, i, :, ::-1, ::-1]139 else:140 tile_overlap = min(0.5, max(0.05, tile_overlap))141 bsizeY, bsizeX = min(bsize, Ly), min(bsize, Lx)142 bsizeY = np.int32(bsizeY)143 bsizeX = np.int32(bsizeX)144 # tiles overlap by 10% tile size145 ny = 1 if Ly <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Ly / bsize))146 nx = 1 if Lx <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Lx / bsize))147 ystart = np.linspace(0, Ly - bsizeY, ny).astype(int)148 xstart = np.linspace(0, Lx - bsizeX, nx).astype(int)149 150 ysub = []151 xsub = []152 IMG = np.zeros((len(ystart), len(xstart), nchan, bsizeY, bsizeX), np.float32)153 for j in range(len(ystart)):154 for i in range(len(xstart)):155 ysub.append([ystart[j], ystart[j] + bsizeY])156 xsub.append([xstart[i], xstart[i] + bsizeX])157 IMG[j, i] = imgi[:, ysub[-1][0]:ysub[-1][1], xsub[-1][0]:xsub[-1][1]]158 159 return IMG, ysub, xsub, Ly, Lx160 161 162def normalize99(Y, lower=1, upper=99, copy=True, downsample=False):163 """164 Normalize the image so that 0.0 corresponds to the 1st percentile and 1.0 corresponds to the 99th percentile.165 166 Args:167 Y (ndarray): The input image (for downsample, use [Ly x Lx] or [Lz x Ly x Lx]).168 lower (int, optional): The lower percentile. Defaults to 1.169 upper (int, optional): The upper percentile. Defaults to 99.170 copy (bool, optional): Whether to create a copy of the input image. Defaults to True.171 downsample (bool, optional): Whether to downsample image to compute percentiles. Defaults to False.172 173 Returns:174 ndarray: The normalized image.175 """176 X = Y.copy() if copy else Y177 X = X.astype("float32") if X.dtype!="float64" and X.dtype!="float32" else X178 if downsample and X.size > 224**3:179 nskip = [max(1, X.shape[i] // 224) for i in range(X.ndim)]180 nskip[0] = max(1, X.shape[0] // 50) if X.ndim == 3 else nskip[0]181 slc = tuple([slice(0, X.shape[i], nskip[i]) for i in range(X.ndim)])182 x01 = np.percentile(X[slc], lower)183 x99 = np.percentile(X[slc], upper)184 else:185 x01 = np.percentile(X, lower)186 x99 = np.percentile(X, upper)187 if x99 - x01 > 1e-3:188 X -= x01 189 X /= (x99 - x01)190 else:191 X[:] = 0192 return X193 194 195def normalize99_tile(img, blocksize=100, lower=1., upper=99., tile_overlap=0.1,196 norm3D=False, smooth3D=1, is3D=False):197 """Compute normalization like normalize99 function but in tiles.198 199 Args:200 img (numpy.ndarray): Array of shape (Lz x) Ly x Lx (x nchan) containing the image.201 blocksize (float, optional): Size of tiles. Defaults to 100.202 lower (float, optional): Lower percentile for normalization. Defaults to 1.0.203 upper (float, optional): Upper percentile for normalization. Defaults to 99.0.204 tile_overlap (float, optional): Fraction of overlap of tiles. Defaults to 0.1.205 norm3D (bool, optional): Use same tiled normalization for each z-plane. Defaults to False.206 smooth3D (int, optional): Smoothing factor for 3D normalization. Defaults to 1.207 is3D (bool, optional): Set to True if image is a 3D stack. Defaults to False.208 209 Returns:210 numpy.ndarray: Normalized image array of shape (Lz x) Ly x Lx (x nchan).211 """212 is1c = True if img.ndim == 2 or (is3D and img.ndim == 3) else False213 is3D = True if img.ndim > 3 or (is3D and img.ndim == 3) else False214 img = img[..., np.newaxis] if is1c else img215 img = img[np.newaxis, ...] if img.ndim == 3 else img216 Lz, Ly, Lx, nchan = img.shape217 218 tile_overlap = min(0.5, max(0.05, tile_overlap))219 blocksizeY, blocksizeX = min(blocksize, Ly), min(blocksize, Lx)220 blocksizeY = np.int32(blocksizeY)221 blocksizeX = np.int32(blocksizeX)222 # tiles overlap by 10% tile size223 ny = 1 if Ly <= blocksize else int(np.ceil(224 (1. + 2 * tile_overlap) * Ly / blocksize))225 nx = 1 if Lx <= blocksize else int(np.ceil(226 (1. + 2 * tile_overlap) * Lx / blocksize))227 ystart = np.linspace(0, Ly - blocksizeY, ny).astype(int)228 xstart = np.linspace(0, Lx - blocksizeX, nx).astype(int)229 ysub = []230 xsub = []231 for j in range(len(ystart)):232 for i in range(len(xstart)):233 ysub.append([ystart[j], ystart[j] + blocksizeY])234 xsub.append([xstart[i], xstart[i] + blocksizeX])235 236 x01_tiles_z = []237 x99_tiles_z = []238 for z in range(Lz):239 IMG = np.zeros((len(ystart), len(xstart), blocksizeY, blocksizeX, nchan),240 "float32")241 k = 0242 for j in range(len(ystart)):243 for i in range(len(xstart)):244 IMG[j, i] = img[z, ysub[k][0]:ysub[k][1], xsub[k][0]:xsub[k][1], :]245 k += 1246 x01_tiles = np.percentile(IMG, lower, axis=(-3, -2))247 x99_tiles = np.percentile(IMG, upper, axis=(-3, -2))248 249 # fill areas with small differences with neighboring squares250 to_fill = np.zeros(x01_tiles.shape[:2], "bool")251 for c in range(nchan):252 to_fill = x99_tiles[:, :, c] - x01_tiles[:, :, c] < +1e-3253 if to_fill.sum() > 0 and to_fill.sum() < x99_tiles[:, :, c].size:254 fill_vals = np.nonzero(to_fill)255 fill_neigh = np.nonzero(~to_fill)256 nearest_neigh = (257 (fill_vals[0] - fill_neigh[0][:, np.newaxis])**2 +258 (fill_vals[1] - fill_neigh[1][:, np.newaxis])**2).argmin(axis=0)259 x01_tiles[fill_vals[0], fill_vals[1],260 c] = x01_tiles[fill_neigh[0][nearest_neigh],261 fill_neigh[1][nearest_neigh], c]262 x99_tiles[fill_vals[0], fill_vals[1],263 c] = x99_tiles[fill_neigh[0][nearest_neigh],264 fill_neigh[1][nearest_neigh], c]265 elif to_fill.sum() > 0 and to_fill.sum() == x99_tiles[:, :, c].size:266 x01_tiles[:, :, c] = 0267 x99_tiles[:, :, c] = 1268 x01_tiles_z.append(x01_tiles)269 x99_tiles_z.append(x99_tiles)270 271 x01_tiles_z = np.array(x01_tiles_z)272 x99_tiles_z = np.array(x99_tiles_z)273 # do not smooth over z-axis if not normalizing separately per plane274 for a in range(2):275 x01_tiles_z = gaussian_filter1d(x01_tiles_z, 1, axis=a)276 x99_tiles_z = gaussian_filter1d(x99_tiles_z, 1, axis=a)277 if norm3D:278 smooth3D = 1 if smooth3D == 0 else smooth3D279 x01_tiles_z = gaussian_filter1d(x01_tiles_z, smooth3D, axis=a)280 x99_tiles_z = gaussian_filter1d(x99_tiles_z, smooth3D, axis=a)281 282 if not norm3D and Lz > 1:283 x01 = np.zeros((len(x01_tiles_z), Ly, Lx, nchan), "float32")284 x99 = np.zeros((len(x01_tiles_z), Ly, Lx, nchan), "float32")285 for z in range(Lz):286 x01_rsz = cv2.resize(x01_tiles_z[z], (Lx, Ly),287 interpolation=cv2.INTER_LINEAR)288 x01[z] = x01_rsz[..., np.newaxis] if nchan == 1 else x01_rsz289 x99_rsz = cv2.resize(x99_tiles_z[z], (Lx, Ly),290 interpolation=cv2.INTER_LINEAR)291 x99[z] = x99_rsz[..., np.newaxis] if nchan == 1 else x01_rsz292 if (x99 - x01).min() < 1e-3:293 raise ZeroDivisionError(294 "cannot use norm3D=False with tile_norm, sample is too sparse; set norm3D=True or tile_norm=0"295 )296 else:297 x01 = cv2.resize(x01_tiles_z.mean(axis=0), (Lx, Ly),298 interpolation=cv2.INTER_LINEAR)299 x99 = cv2.resize(x99_tiles_z.mean(axis=0), (Lx, Ly),300 interpolation=cv2.INTER_LINEAR)301 if x01.ndim < 3:302 x01 = x01[..., np.newaxis]303 x99 = x99[..., np.newaxis]304 305 if is1c:306 img, x01, x99 = img.squeeze(), x01.squeeze(), x99.squeeze()307 elif not is3D:308 img, x01, x99 = img[0], x01[0], x99[0]309 310 # normalize311 img -= x01 312 img /= (x99 - x01)313 314 return img315 316 317def gaussian_kernel(sigma, Ly, Lx, device=torch.device("cpu")):318 """319 Generates a 2D Gaussian kernel.320 321 Args:322 sigma (float): Standard deviation of the Gaussian distribution.323 Ly (int): Number of pixels in the y-axis.324 Lx (int): Number of pixels in the x-axis.325 device (torch.device, optional): Device to store the kernel tensor. Defaults to torch.device("cpu").326 327 Returns:328 torch.Tensor: 2D Gaussian kernel tensor.329 330 """331 y = torch.linspace(-Ly / 2, Ly / 2 + 1, Ly, device=device)332 x = torch.linspace(-Ly / 2, Ly / 2 + 1, Lx, device=device)333 y, x = torch.meshgrid(y, x, indexing="ij")334 kernel = torch.exp(-(y**2 + x**2) / (2 * sigma**2))335 kernel /= kernel.sum()336 return kernel337 338 339def smooth_sharpen_img(img, smooth_radius=6, sharpen_radius=12,340 device=torch.device("cpu"), is3D=False):341 """Sharpen blurry images with surround subtraction and/or smooth noisy images.342 343 Args:344 img (float32): Array that's (Lz x) Ly x Lx (x nchan).345 smooth_radius (float, optional): Size of gaussian smoothing filter, recommended to be 1/10-1/4 of cell diameter346 (if also sharpening, should be 2-3x smaller than sharpen_radius). Defaults to 6.347 sharpen_radius (float, optional): Size of gaussian surround filter, recommended to be 1/8-1/2 of cell diameter348 (if also smoothing, should be 2-3x larger than smooth_radius). Defaults to 12.349 device (torch.device, optional): Device on which to perform sharpening.350 Will be faster on GPU but need to ensure GPU has RAM for image. Defaults to torch.device("cpu").351 is3D (bool, optional): If image is 3D stack (only necessary to set if img.ndim==3). Defaults to False.352 353 Returns:354 img_sharpen (float32): Array that's (Lz x) Ly x Lx (x nchan).355 """356 img_sharpen = torch.from_numpy(img.astype("float32")).to(device)357 shape = img_sharpen.shape358 359 is1c = True if img_sharpen.ndim == 2 or (is3D and img_sharpen.ndim == 3) else False360 is3D = True if img_sharpen.ndim > 3 or (is3D and img_sharpen.ndim == 3) else False361 img_sharpen = img_sharpen.unsqueeze(-1) if is1c else img_sharpen362 img_sharpen = img_sharpen.unsqueeze(0) if img_sharpen.ndim == 3 else img_sharpen363 Lz, Ly, Lx, nchan = img_sharpen.shape364 365 if smooth_radius > 0:366 kernel = gaussian_kernel(smooth_radius, Ly, Lx, device=device)367 if sharpen_radius > 0:368 kernel += -1 * gaussian_kernel(sharpen_radius, Ly, Lx, device=device)369 elif sharpen_radius > 0:370 kernel = -1 * gaussian_kernel(sharpen_radius, Ly, Lx, device=device)371 kernel[Ly // 2, Lx // 2] = 1372 373 fhp = fft2(kernel)374 for z in range(Lz):375 for c in range(nchan):376 img_filt = torch.real(ifft2(377 fft2(img_sharpen[z, :, :, c]) * torch.conj(fhp)))378 img_filt = fftshift(img_filt)379 img_sharpen[z, :, :, c] = img_filt380 381 img_sharpen = img_sharpen.reshape(shape)382 return img_sharpen.cpu().numpy()383 384 385def move_axis(img, m_axis=-1, first=True):386 """ move axis m_axis to first or last position """387 if m_axis == -1:388 m_axis = img.ndim - 1389 m_axis = min(img.ndim - 1, m_axis)390 axes = np.arange(0, img.ndim)391 if first:392 axes[1:m_axis + 1] = axes[:m_axis]393 axes[0] = m_axis394 else:395 axes[m_axis:-1] = axes[m_axis + 1:]396 axes[-1] = m_axis397 img = img.transpose(tuple(axes))398 return img399 400 401 402def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):403 """Converts the image to have the z-axis first, channels last. Image will be converted to 3 channels if it is not already.404 If more than 3 channels are provided, only the first 3 channels will be used. 405 406 Accepts: 407 - 2D images with no channel dimension: `z_axis` and `channel_axis` must be `None`408 - 2D images with channel dimension: `channel_axis` will be guessed between first or last axis, can also specify `channel_axis`. `z_axis` must be `None`409 - 3D images with or without channels: 410 411 Args:412 x (numpy.ndarray or torch.Tensor): The input image.413 channel_axis (int or None): The axis of the channels in the input image. If None, the axis is determined automatically.414 z_axis (int or None): The axis of the z-dimension in the input image. If None, the axis is determined automatically.415 do_3D (bool): Whether to process the image in 3D mode. Defaults to False.416 417 Returns:418 numpy.ndarray: The converted image.419 420 Raises:421 ValueError: If the input image is 2D and do_3D is True.422 ValueError: If the input image is 4D and do_3D is False.423 """424 425 # check if image is a torch array instead of numpy array, convert to numpy426 ndim = x.ndim427 if torch.is_tensor(x):428 transforms_logger.warning("torch array used as input, converting to numpy")429 x = x.cpu().numpy()430 431 # should be 2D432 if z_axis is not None and not do_3D:433 raise ValueError("2D image provided, but z_axis is not None. Set z_axis=None to process 2D images of ndim=2 or 3.")434 435 if ndim == 4 and not do_3D:436 raise ValueError("3D input image provided, but do_3D is False. Set do_3D=True to process 3D images. ndims=4")437 438 439 ######################## 2D reshaping ########################440 # if user specifies channel axis, return early441 if channel_axis is not None:442 if ndim == 2:443 raise ValueError("2D image provided, but channel_axis is not None. Set channel_axis=None to process 2D images of ndim=2.")444 445 # Put channel axis last:446 # Find the indices of the dims that need to be put in dim 0 and 1447 n_channels = x.shape[channel_axis]448 x_shape_dims = list(x.shape)449 del x_shape_dims[channel_axis]450 dimension_indicies = [i for i in range(x.ndim)]451 del dimension_indicies[channel_axis]452 453 x = x.transpose((dimension_indicies[0], dimension_indicies[1], channel_axis))454 455 if n_channels != 3:456 x_chans_to_copy = min(3, n_channels)457 458 if n_channels > 3: 459 transforms_logger.warning("more than 3 channels provided, only segmenting on first 3 channels")460 x = x[..., :x_chans_to_copy]461 else: 462 x_out = np.zeros((x_shape_dims[0], x_shape_dims[1], 3), dtype=x.dtype)463 x_out[..., :x_chans_to_copy] = x[...]464 x = x_out465 del x_out466 467 return x468 469 # do image padding and channel conversion470 if ndim == 2:471 # grayscale image, make 3 channels472 x_out = np.zeros((x.shape[0], x.shape[1], 3), dtype=x.dtype)473 x_out[..., 0] = x474 x = x_out475 del x_out476 elif ndim == 3:477 # assume 2d with channels478 # find dim with smaller size between first and last dims479 move_channel_axis = x.shape[0] < x.shape[2]480 if move_channel_axis:481 x = x.transpose((1, 2, 0))482 483 # zero padding up to 3 channels: 484 num_channels = x.shape[-1]485 if num_channels > 3: 486 transforms_logger.warning("Found more than 3 channels, only using first 3")487 num_channels = 3488 x_out = np.zeros((x.shape[0], x.shape[1], 3), dtype=x.dtype)489 x_out[..., :num_channels] = x[..., :num_channels]490 x = x_out491 del x_out492 else:493 # something is wrong: yell494 expected_shapes = "2D (H, W), 3D (H, W, C), or 4D (Z, H, W, C)"495 transforms_logger.critical(f"ERROR: Unexpected image shape: {str(x.shape)}. Expected shapes: {expected_shapes}")496 raise ValueError(f"ERROR: Unexpected image shape: {str(x.shape)}. Expected shapes: {expected_shapes}")497 498 return x499 500 501def normalize_img(img, normalize=True, norm3D=True, invert=False, lowhigh=None,502 percentile=(1., 99.), sharpen_radius=0, smooth_radius=0,503 tile_norm_blocksize=0, tile_norm_smooth3D=1, axis=-1):504 """Normalize each channel of the image with optional inversion, smoothing, and sharpening.505 506 Args:507 img (ndarray): The input image. It should have at least 3 dimensions.508 If it is 4-dimensional, it assumes the first non-channel axis is the Z dimension.509 normalize (bool, optional): Whether to perform normalization. Defaults to True.510 norm3D (bool, optional): Whether to normalize in 3D. If True, the entire 3D stack will511 be normalized per channel. If False, normalization is applied per Z-slice. Defaults to False.512 invert (bool, optional): Whether to invert the image. Useful if cells are dark instead of bright.513 Defaults to False.514 lowhigh (tuple or ndarray, optional): The lower and upper bounds for normalization.515 Can be a tuple of two values (applied to all channels) or an array of shape (nchan, 2)516 for per-channel normalization. Incompatible with smoothing and sharpening.517 Defaults to None.518 percentile (tuple, optional): The lower and upper percentiles for normalization. If provided, it should be519 a tuple of two values. Each value should be between 0 and 100. Defaults to (1.0, 99.0).520 sharpen_radius (int, optional): The radius for sharpening the image. Defaults to 0.521 smooth_radius (int, optional): The radius for smoothing the image. Defaults to 0.522 tile_norm_blocksize (int, optional): The block size for tile-based normalization. Defaults to 0.523 tile_norm_smooth3D (int, optional): The smoothness factor for tile-based normalization in 3D. Defaults to 1.524 axis (int, optional): The channel axis to loop over for normalization. Defaults to -1.525 526 Returns:527 ndarray: The normalized image of the same size.528 529 Raises:530 ValueError: If the image has less than 3 dimensions.531 ValueError: If the provided lowhigh or percentile values are invalid.532 ValueError: If the image is inverted without normalization.533 534 """535 if img.ndim < 3:536 error_message = "Image needs to have at least 3 dimensions"537 transforms_logger.critical(error_message)538 raise ValueError(error_message)539 540 img_norm = img if img.dtype=="float32" else img.astype(np.float32)541 if axis != -1 and axis != img_norm.ndim - 1:542 img_norm = np.moveaxis(img_norm, axis, -1) # Move channel axis to last543 544 nchan = img_norm.shape[-1]545 546 # Validate and handle lowhigh bounds547 if lowhigh is not None:548 lowhigh = np.array(lowhigh)549 if lowhigh.shape == (2,):550 lowhigh = np.tile(lowhigh, (nchan, 1)) # Expand to per-channel bounds551 elif lowhigh.shape != (nchan, 2):552 error_message = "`lowhigh` must have shape (2,) or (nchan, 2)"553 transforms_logger.critical(error_message)554 raise ValueError(error_message)555 556 # Validate percentile557 if percentile is None:558 percentile = (1.0, 99.0)559 elif not (0 <= percentile[0] < percentile[1] <= 100):560 error_message = "Invalid percentile range, should be between 0 and 100"561 transforms_logger.critical(error_message)562 raise ValueError(error_message)563 564 # Apply normalization based on lowhigh or percentile565 cgood = np.zeros(nchan, "bool")566 if lowhigh is not None:567 for c in range(nchan):568 lower = lowhigh[c, 0]569 upper = lowhigh[c, 1]570 img_norm[..., c] -= lower 571 img_norm[..., c] /= (upper - lower)572 cgood[c] = True573 else:574 # Apply sharpening and smoothing if specified575 if sharpen_radius > 0 or smooth_radius > 0:576 img_norm = smooth_sharpen_img(577 img_norm, sharpen_radius=sharpen_radius, smooth_radius=smooth_radius578 )579 580 # Apply tile-based normalization or standard normalization581 if tile_norm_blocksize > 0:582 img_norm = normalize99_tile(583 img_norm,584 blocksize=tile_norm_blocksize,585 lower=percentile[0],586 upper=percentile[1],587 smooth3D=tile_norm_smooth3D,588 norm3D=norm3D,589 )590 cgood[:] = True591 elif normalize:592 if img_norm.ndim == 3 or norm3D: # i.e. if YXC, or ZYXC with norm3D=True593 for c in range(nchan):594 if np.ptp(img_norm[..., c]) > 0.:595 img_norm[..., c] = normalize99(596 img_norm[..., c],597 lower=percentile[0],598 upper=percentile[1],599 copy=False, downsample=True,600 )601 cgood[c] = True602 else: # i.e. if ZYXC with norm3D=False then per Z-slice603 for z in range(img_norm.shape[0]):604 for c in range(nchan):605 if np.ptp(img_norm[z, ..., c]) > 0.:606 img_norm[z, ..., c] = normalize99(607 img_norm[z, ..., c],608 lower=percentile[0],609 upper=percentile[1],610 copy=False, downsample=True,611 )612 cgood[c] = True613 614 615 if invert:616 if lowhigh is not None or tile_norm_blocksize > 0 or normalize:617 for c in range(nchan):618 if cgood[c]:619 img_norm[..., c] = 1 - img_norm[..., c]620 else:621 error_message = "Cannot invert image without normalization"622 transforms_logger.critical(error_message)623 raise ValueError(error_message)624 625 # Move channel axis back to the original position626 if axis != -1 and axis != img_norm.ndim - 1:627 img_norm = np.moveaxis(img_norm, -1, axis)628 629 # The transformer can get confused if a channel is all 1's instead of all 0's:630 for i, chan_did_normalize in enumerate(cgood):631 if not chan_did_normalize:632 if img_norm.ndim == 3:633 img_norm[:, :, i] = 0634 if img_norm.ndim == 4:635 img_norm[:, :, :, i] = 0636 637 return img_norm638 639def resize_safe(img, Ly, Lx, interpolation=cv2.INTER_LINEAR):640 """OpenCV resize function does not support uint32.641 642 This function converts the image to float32 before resizing and then converts it back to uint32. Not safe!643 References issue: https://github.com/MouseLand/cellpose/issues/937644 645 Implications:646 * Runtime: Runtime increases by 5x-50x due to type casting. However, with resizing being very efficient, this is not647 a big issue. A 10,000x10,000 image takes 0.47s instead of 0.016s to cast and resize on 32 cores on GPU.648 * Memory: However, memory usage increases. Not tested by how much.649 650 Args:651 img (ndarray): Image of size [Ly x Lx].652 Ly (int): Desired height of the resized image.653 Lx (int): Desired width of the resized image.654 interpolation (int, optional): OpenCV interpolation method. Defaults to cv2.INTER_LINEAR.655 656 Returns:657 ndarray: Resized image of size [Ly x Lx].658 659 """660 661 # cast image662 cast = img.dtype == np.uint32663 if cast:664 img = img.astype(np.float32)665 666 # resize667 img = cv2.resize(img, (Lx, Ly), interpolation=interpolation)668 669 # cast back670 if cast:671 img = img.round().astype(np.uint32)672 673 return img674 675 676def resize_image(img0, Ly=None, Lx=None, rsz=None, interpolation=cv2.INTER_LINEAR,677 no_channels=False):678 """Resize image for computing flows / unresize for computing dynamics.679 680 Args:681 img0 (ndarray): Image of size [Y x X x nchan] or [Lz x Y x X x nchan] or [Lz x Y x X].682 Ly (int, optional): Desired height of the resized image. Defaults to None.683 Lx (int, optional): Desired width of the resized image. Defaults to None.684 rsz (float, optional): Resize coefficient(s) for the image. If Ly is None, rsz is used. Defaults to None.685 interpolation (int, optional): OpenCV interpolation method. Defaults to cv2.INTER_LINEAR.686 no_channels (bool, optional): Flag indicating whether to treat the third dimension as a channel.687 Defaults to False.688 689 Returns:690 ndarray: Resized image of size [Ly x Lx x nchan] or [Lz x Ly x Lx x nchan].691 692 Raises:693 ValueError: If Ly is None and rsz is None.694 695 """696 if Ly is None and rsz is None:697 error_message = "must give size to resize to or factor to use for resizing"698 transforms_logger.critical(error_message)699 raise ValueError(error_message)700 701 if Ly is None:702 # determine Ly and Lx using rsz703 if not isinstance(rsz, list) and not isinstance(rsz, np.ndarray):704 rsz = [rsz, rsz]705 if no_channels:706 Ly = int(img0.shape[-2] * rsz[-2])707 Lx = int(img0.shape[-1] * rsz[-1])708 else:709 Ly = int(img0.shape[-3] * rsz[-2])710 Lx = int(img0.shape[-2] * rsz[-1])711 712 # no_channels useful for z-stacks, so the third dimension is not treated as a channel713 # but if this is called for grayscale images, they first become [Ly,Lx,2] so ndim=3 but714 if (img0.ndim > 2 and no_channels) or (img0.ndim == 4 and not no_channels):715 if Ly == 0 or Lx == 0:716 raise ValueError(717 "anisotropy too high / low -- not enough pixels to resize to ratio")718 for i, img in enumerate(img0):719 imgi = resize_safe(img, Ly, Lx, interpolation=interpolation)720 if i==0:721 if no_channels:722 imgs = np.zeros((img0.shape[0], Ly, Lx), imgi.dtype)723 else:724 imgs = np.zeros((img0.shape[0], Ly, Lx, img0.shape[-1]), imgi.dtype)725 imgs[i] = imgi if imgi.ndim > 2 or no_channels else imgi[..., np.newaxis]726 else:727 imgs = resize_safe(img0, Ly, Lx, interpolation=interpolation)728 return imgs729 730def get_pad_yx(Ly, Lx, div=16, extra=1, min_size=None):731 if min_size is None or Ly >= min_size[-2]:732 Lpad = int(div * np.ceil(Ly / div) - Ly)733 else:734 Lpad = min_size[-2] - Ly735 ypad1 = extra * div // 2 + Lpad // 2736 ypad2 = extra * div // 2 + Lpad - Lpad // 2737 if min_size is None or Lx >= min_size[-1]:738 Lpad = int(div * np.ceil(Lx / div) - Lx)739 else:740 Lpad = min_size[-1] - Lx741 xpad1 = extra * div // 2 + Lpad // 2742 xpad2 = extra * div // 2 + Lpad - Lpad // 2743 744 return ypad1, ypad2, xpad1, xpad2745 746 747def pad_image_ND(img0, div=16, extra=1, min_size=None, zpad=False):748 """Pad image for test-time so that its dimensions are a multiple of 16 (2D or 3D).749 750 Args:751 img0 (ndarray): Image of size [nchan (x Lz) x Ly x Lx].752 div (int, optional): Divisor for padding. Defaults to 16.753 extra (int, optional): Extra padding. Defaults to 1.754 min_size (tuple, optional): Minimum size of the image. Defaults to None.755 756 Returns:757 A tuple containing (I, ysub, xsub) or (I, ysub, xsub, zsub), I is padded image, -sub are ranges of pixels in the padded image corresponding to img0.758 759 """760 Ly, Lx = img0.shape[-2:]761 ypad1, ypad2, xpad1, xpad2 = get_pad_yx(Ly, Lx, div=div, extra=extra, min_size=min_size)762 763 if img0.ndim > 3:764 if zpad:765 Lpad = int(div * np.ceil(img0.shape[-3] / div) - img0.shape[-3])766 zpad1 = extra * div // 2 + Lpad // 2767 zpad2 = extra * div // 2 + Lpad - Lpad // 2768 else:769 zpad1, zpad2 = 0, 0770 pads = np.array([[0, 0], [zpad1, zpad2], [ypad1, ypad2], [xpad1, xpad2]])771 else:772 pads = np.array([[0, 0], [ypad1, ypad2], [xpad1, xpad2]])773 774 I = np.pad(img0, pads, mode="constant")775 776 ysub = np.arange(ypad1, ypad1 + Ly)777 xsub = np.arange(xpad1, xpad1 + Lx)778 if zpad:779 zsub = np.arange(zpad1, zpad1 + img0.shape[-3])780 return I, ysub, xsub, zsub781 else:782 return I, ysub, xsub783 784 785def random_rotate_and_resize(X, Y=None, scale_range=1., xy=(224, 224), do_3D=False,786 zcrop=48, do_flip=True, rotate=True, rescale=None, unet=False,787 random_per_image=True):788 """Augmentation by random rotation and resizing.789 790 Args:791 X (list of ND-arrays, float): List of image arrays of size [nchan x Ly x Lx] or [Ly x Lx].792 Y (list of ND-arrays, float, optional): List of image labels of size [nlabels x Ly x Lx] or [Ly x Lx].793 The 1st channel of Y is always nearest-neighbor interpolated (assumed to be masks or 0-1 representation).794 If Y.shape[0]==3 and not unet, then the labels are assumed to be [cell probability, Y flow, X flow].795 If unet, second channel is dist_to_bound. Defaults to None.796 scale_range (float, optional): Range of resizing of images for augmentation.797 Images are resized by (1-scale_range/2) + scale_range * np.random.rand(). Defaults to 1.0.798 xy (tuple, int, optional): Size of transformed images to return. Defaults to (224,224).799 do_flip (bool, optional): Whether or not to flip images horizontally. Defaults to True.800 rotate (bool, optional): Whether or not to rotate images. Defaults to True.801 rescale (array, float, optional): How much to resize images by before performing augmentations. Defaults to None.802 unet (bool, optional): Whether or not to use unet. Defaults to False.803 random_per_image (bool, optional): Different random rotate and resize per image. Defaults to True.804 805 Returns:806 A tuple containing (imgi, lbl, scale): imgi (ND-array, float): Transformed images in array [nimg x nchan x xy[0] x xy[1]]; 807 lbl (ND-array, float): Transformed labels in array [nimg x nchan x xy[0] x xy[1]]; 808 scale (array, float): Amount each image was resized by.809 """810 scale_range = max(0, min(2, float(scale_range))) if scale_range is not None else scale_range811 nimg = len(X)812 if X[0].ndim > 2:813 nchan = X[0].shape[0]814 else:815 nchan = 1816 # if do_3D and X[0].ndim > 3:817 # shape = (zcrop, xy[0], xy[1])818 # else:819 shape = (xy[0], xy[1])820 imgi = np.zeros((nimg, nchan, *shape), "float32")821 822 lbl = []823 if Y is not None:824 if Y[0].ndim > 2:825 nt = Y[0].shape[0]826 else:827 nt = 1828 lbl = np.zeros((nimg, nt, *shape), np.float32)829 830 scale = np.ones(nimg, np.float32)831 832 for n in range(nimg):833 834 if random_per_image or n == 0:835 Ly, Lx = X[n].shape[-2:]836 # generate random augmentation parameters837 flip = np.random.rand() > .5838 theta = np.random.rand() * np.pi * 2 if rotate else 0.839 if scale_range is None:840 scale[n] = 2 ** (4 * np.random.rand() - 2)841 else:842 scale[n] = (1 - scale_range / 2) + scale_range * np.random.rand()843 if rescale is not None:844 scale[n] *= 1. / rescale[n]845 dxy = np.maximum(0, np.array([Lx * scale[n] - xy[1],846 Ly * scale[n] - xy[0]]))847 dxy = (np.random.rand(2,) - .5) * dxy848 849 # create affine transform850 cc = np.array([Lx / 2, Ly / 2])851 cc1 = cc - np.array([Lx - xy[1], Ly - xy[0]]) / 2 + dxy852 pts1 = np.float32([cc, cc + np.array([1, 0]), cc + np.array([0, 1])])853 pts2 = np.float32([854 cc1,855 cc1 + scale[n] * np.array([np.cos(theta), np.sin(theta)]),856 cc1 + scale[n] *857 np.array([np.cos(np.pi / 2 + theta),858 np.sin(np.pi / 2 + theta)])859 ])860 M = cv2.getAffineTransform(pts1, pts2)861 862 img = X[n].copy()863 if Y is not None:864 labels = Y[n].copy()865 if labels.ndim < 3:866 labels = labels[np.newaxis, :, :]867 868 if do_flip:869 if flip:870 img = img[..., ::-1]871 if Y is not None:872 labels = labels[..., ::-1]873 if nt > 1 and not unet:874 labels[-1] = -labels[-1]875 876 for k in range(nchan):877 I = cv2.warpAffine(img[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)878 imgi[n, k] = I879 880 if Y is not None:881 for k in range(nt):882 flag = cv2.INTER_NEAREST if k < nt-2 else cv2.INTER_LINEAR883 lbl[n, k] = cv2.warpAffine(labels[k], M, (xy[1], xy[0]), flags=flag)884 885 if nt > 1 and not unet:886 v1 = lbl[n, -1].copy()887 v2 = lbl[n, -2].copy()888 lbl[n, -2] = (-v1 * np.sin(-theta) + v2 * np.cos(-theta))889 lbl[n, -1] = (v1 * np.cos(-theta) + v2 * np.sin(-theta))890 891 return imgi, lbl, scale892 893 894def random_rotate_and_resize_with_feat(X, Y=None, feat=None, scale_range=1., xy=(224, 224), do_3D=False,895 zcrop=48, do_flip=True, rotate=True, rescale=None, unet=False,896 random_per_image=True):897 """Augmentation by random rotation and resizing.898 899 Args:900 X (list of ND-arrays, float): List of image arrays of size [nchan x Ly x Lx] or [Ly x Lx].901 Y (list of ND-arrays, float, optional): List of image labels of size [nlabels x Ly x Lx] or [Ly x Lx].902 The 1st channel of Y is always nearest-neighbor interpolated (assumed to be masks or 0-1 representation).903 If Y.shape[0]==3 and not unet, then the labels are assumed to be [cell probability, Y flow, X flow].904 If unet, second channel is dist_to_bound. Defaults to None.905 scale_range (float, optional): Range of resizing of images for augmentation.906 Images are resized by (1-scale_range/2) + scale_range * np.random.rand(). Defaults to 1.0.907 xy (tuple, int, optional): Size of transformed images to return. Defaults to (224,224).908 do_flip (bool, optional): Whether or not to flip images horizontally. Defaults to True.909 rotate (bool, optional): Whether or not to rotate images. Defaults to True.910 rescale (array, float, optional): How much to resize images by before performing augmentations. Defaults to None.911 unet (bool, optional): Whether or not to use unet. Defaults to False.912 random_per_image (bool, optional): Different random rotate and resize per image. Defaults to True.913 914 Returns:915 A tuple containing (imgi, lbl, scale): imgi (ND-array, float): Transformed images in array [nimg x nchan x xy[0] x xy[1]]; 916 lbl (ND-array, float): Transformed labels in array [nimg x nchan x xy[0] x xy[1]]; 917 scale (array, float): Amount each image was resized by.918 """919 scale_range = max(0, min(2, float(scale_range))) if scale_range is not None else scale_range920 nimg = len(X)921 if X[0].ndim > 2:922 nchan = X[0].shape[0]923 else:924 nchan = 1925 shape = (xy[0], xy[1])926 imgi = np.zeros((nimg, nchan, *shape), "float32")927 928 lbl = []929 if Y is not None:930 if Y[0].ndim > 2:931 nt = Y[0].shape[0]932 else:933 nt = 1934 lbl = np.zeros((nimg, nt, *shape), np.float32)935 936 if feat is not None:937 if feat[0].ndim > 2:938 nf = feat[0].shape[0]939 else:940 nf = 1941 feat_out = np.zeros((nimg, nf, *shape), "float32")942 943 scale = np.ones(nimg, np.float32)944 945 for n in range(nimg):946 947 if random_per_image or n == 0:948 Ly, Lx = X[n].shape[-2:]949 # generate random augmentation parameters950 flip = np.random.rand() > .5951 theta = np.random.rand() * np.pi * 2 if rotate else 0.952 if scale_range is None:953 scale[n] = 2 ** (4 * np.random.rand() - 2)954 else:955 scale[n] = (1 - scale_range / 2) + scale_range * np.random.rand()956 if rescale is not None:957 scale[n] *= 1. / rescale[n]958 dxy = np.maximum(0, np.array([Lx * scale[n] - xy[1],959 Ly * scale[n] - xy[0]]))960 dxy = (np.random.rand(2,) - .5) * dxy961 962 # create affine transform963 cc = np.array([Lx / 2, Ly / 2])964 cc1 = cc - np.array([Lx - xy[1], Ly - xy[0]]) / 2 + dxy965 pts1 = np.float32([cc, cc + np.array([1, 0]), cc + np.array([0, 1])])966 pts2 = np.float32([967 cc1,968 cc1 + scale[n] * np.array([np.cos(theta), np.sin(theta)]),969 cc1 + scale[n] *970 np.array([np.cos(np.pi / 2 + theta),971 np.sin(np.pi / 2 + theta)])972 ])973 M = cv2.getAffineTransform(pts1, pts2)974 975 img = X[n].copy()976 if Y is not None:977 labels = Y[n].copy()978 if labels.ndim < 3:979 labels = labels[np.newaxis, :, :]980 if feat is not None:981 feats = feat[n].copy()982 if feats.ndim < 3:983 feats = feats[np.newaxis, :, :]984 985 if do_flip:986 if flip:987 img = img[..., ::-1]988 if Y is not None:989 labels = labels[..., ::-1]990 if nt > 1 and not unet:991 labels[-1] = -labels[-1]992 if feat is not None:993 feats = feats[..., ::-1]994 995 996 for k in range(nchan):997 I = cv2.warpAffine(img[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)998 imgi[n, k] = I999 1000 if Y is not None:1001 for k in range(nt):1002 flag = cv2.INTER_NEAREST if k < nt-2 else cv2.INTER_LINEAR1003 lbl[n, k] = cv2.warpAffine(labels[k], M, (xy[1], xy[0]), flags=flag)1004 1005 if nt > 1 and not unet:1006 v1 = lbl[n, -1].copy()1007 v2 = lbl[n, -2].copy()1008 lbl[n, -2] = (-v1 * np.sin(-theta) + v2 * np.cos(-theta))1009 lbl[n, -1] = (v1 * np.cos(-theta) + v2 * np.sin(-theta))1010 1011 if feat is not None:1012 for k in range(nf):1013 feat_out[n, k] = cv2.warpAffine(feats[k], M, (xy[1], xy[0]), flags=cv2.INTER_LINEAR)1014 1015 1016 1017 return imgi, lbl, feat_out, scale1018