torch_em.transform.raw

  1from functools import partial
  2from typing import Union, Optional, Tuple, Dict, Callable
  3
  4import numpy as np
  5
  6import torch
  7from torchvision import transforms
  8
  9
 10#
 11# normalization functions
 12#
 13
 14
 15TORCH_DTYPES = {
 16    "float16": torch.float16,
 17    "float32": torch.float32,
 18    "float64": torch.float64,
 19    "complex64": torch.complex64,
 20    "complex128": torch.complex128,
 21    "uint8": torch.uint8,
 22    "int8": torch.int8,
 23    "int16": torch.int16,
 24    "int32": torch.int32,
 25    "int64": torch.int64,
 26    "bool": torch.bool,
 27}
 28"""@private
 29"""
 30
 31
 32def cast(inpt: Union[np.ndarray, torch.tensor], typestring: torch.dtype):
 33    """@private
 34    """
 35    if torch.is_tensor(inpt):
 36        assert typestring in TORCH_DTYPES, f"{typestring} not in TORCH_DTYPES"
 37        return inpt.to(TORCH_DTYPES[typestring])
 38    return inpt.astype(typestring)
 39
 40
 41def standardize(
 42    raw: np.ndarray,
 43    mean: Optional[float] = None,
 44    std: Optional[float] = None,
 45    axis: Optional[Union[int, Tuple[int, ...]]] = None,
 46    eps: float = 1e-7,
 47) -> np.ndarray:
 48    """Standardize the input data by subtracting its mean and dividing by its standard deviation.
 49
 50    Args:
 51        raw: The input data.
 52        mean: The mean value. If None, it will be computed from the data.
 53        std: The standard deviation. If None, it will be computed from the data.
 54        axis: The axis along which to compute the mean and standard deviation.
 55        eps: The epsilon value for numerical stability.
 56
 57    Returns:
 58        The standardized input data.
 59    """
 60    raw = cast(raw, "float32")
 61    mean = raw.mean(axis=axis, keepdims=True) if mean is None else mean
 62    raw -= mean
 63
 64    std = raw.std(axis=axis, keepdims=True) if std is None else std
 65    raw /= (std + eps)
 66    return raw
 67
 68
 69def _normalize_torch(tensor, minval, maxval, axis, eps):
 70    if axis:  # torch returns torch.return_types.min or torch.return_types.max
 71        minval = torch.amin(tensor, dim=axis, keepdim=True) if minval is None else minval
 72        tensor -= minval
 73
 74        maxval = torch.amax(tensor, dim=axis, keepdim=True) if maxval is None else maxval
 75        tensor /= (maxval + eps)
 76
 77        return tensor
 78
 79    # keepdim can only be used in combination with dim
 80    minval = tensor.min() if minval is None else minval
 81    tensor -= minval
 82
 83    maxval = tensor.max() if maxval is None else maxval
 84    tensor /= (maxval + eps)
 85
 86    return tensor
 87
 88
 89def normalize(
 90    raw: Union[torch.tensor, np.ndarray],
 91    minval: Optional[float] = None,
 92    maxval: Optional[float] = None,
 93    axis: Optional[Union[int, Tuple[int, ...]]] = None,
 94    eps: float = 1e-7,
 95) -> np.ndarray:
 96    """Normalize the input data so that it is in range [0, 1].
 97
 98    Args:
 99        raw: The input data.
100        minval: The minimum data value. If None, it will be computed from the data.
101        maxval: The maximum data value. If None, it will be computed from the data.
102        axis: The axis along which to compute the min and max value.
103        eps: The epsilon value for numerical stability.
104
105    Returns:
106        The normalized input data.
107    """
108    raw = cast(raw, "float32")
109    if torch.is_tensor(raw):
110        return _normalize_torch(raw, minval=minval, maxval=maxval, axis=axis, eps=eps)
111
112    minval = raw.min(axis=axis, keepdims=True) if minval is None else minval
113    raw -= minval
114
115    maxval = raw.max(axis=axis, keepdims=True) if maxval is None else maxval
116    raw /= (maxval + eps)
117    return raw
118
119
120def normalize_percentile(
121    raw: np.ndarray,
122    lower: float = 1.0,
123    upper: float = 99.0,
124    axis: Optional[Union[int, Tuple[int, ...]]] = None,
125    eps: float = 1e-7,
126) -> np.ndarray:
127    """Normalize the input data based on percentile values.
128
129    Args:
130        raw: The input data.
131        lower: The lower percentile.
132        upper: The upper percentile.
133        axis: The axis along which to compute the percentiles.
134        eps: The epsilon value for numerical stability.
135
136    Returns:
137        The normalized input data.
138    """
139    v_lower = np.percentile(raw, lower, axis=axis, keepdims=True)
140    v_upper = np.percentile(raw, upper, axis=axis, keepdims=True) - v_lower
141    return normalize(raw, v_lower, v_upper, eps=eps)
142
143
144def _normalize_percentile_to_range(
145    raw: Union[np.ndarray, torch.Tensor], clip_max: int
146) -> Union[np.ndarray, torch.Tensor]:
147    raw = normalize_percentile(raw)
148    if torch.is_tensor(raw):
149        return torch.clamp(raw, min=0.0, max=1.0) * clip_max
150    return np.clip(raw, 0.0, 1.0) * clip_max
151
152
153class RandomPercentileNormalization:
154    """Normalize inputs with randomly sampled percentile bounds.
155
156    By default, the lower and upper percentiles are sampled uniformly from
157    ``lower_percentile_bounds`` and ``upper_percentile_bounds``. If no upper bounds are given,
158    they are inferred by mirroring the lower bounds around 50. Normal (Gaussian) sampling can be
159    enabled explicitly with ``distribution="normal"`` and
160    ``distribution_kwargs={"mean": ..., "std": ...}``. The sampled percentile intensities are
161    mapped to 0 and 1, and values outside them are clipped, so the output is always in ``[0, 1]``.
162
163    Examples:
164        Uniform sampling with the default percentile bounds and reproducible random draws:
165
166        ```python
167        normalization = RandomPercentileNormalization(seed=42)
168        ```
169
170        Normal sampling with explicit distribution parameters:
171
172        ```python
173        normalization = RandomPercentileNormalization(
174            distribution="normal",
175            distribution_kwargs={"mean": 2.0, "std": 1.0},
176            seed=42,
177        )
178        ```
179
180    Args:
181        lower_percentile_bounds: Inclusive clipping bounds for the lower percentile.
182        upper_percentile_bounds: Inclusive clipping bounds for the upper percentile. If None, the
183            bounds are inferred by mirroring ``lower_percentile_bounds`` around 50.
184        distribution: Sampling distribution for the percentiles. Supported values are
185            ``"uniform"`` and ``"normal"``.
186        distribution_kwargs: Parameters for normal sampling, which must contain exactly
187            ``{"mean": ..., "std": ...}``. The upper percentile uses the mirrored normal
188            distribution. Uniform sampling does not take additional parameters.
189        rounding_decimals: Number of decimal places used to round sampled percentiles. Set to None
190            to disable rounding.
191        axis: Axes over which to compute the intensity percentiles.
192        seed: Optional non-negative integer seed for reproducible sampling. NumPy integer types
193            are also supported. Each DataLoader worker derives a distinct stream from this seed.
194            By default, the global NumPy random state is used, which respects DataLoader worker seeding.
195        eps: Epsilon used for numerical stability during normalization.
196    """
197
198    def __init__(
199        self,
200        lower_percentile_bounds: Tuple[float, float] = (0.0, 5.0),
201        upper_percentile_bounds: Optional[Tuple[float, float]] = None,
202        distribution: str = "uniform",
203        distribution_kwargs: Optional[Dict[str, float]] = None,
204        rounding_decimals: Optional[int] = 1,
205        axis: Optional[Union[int, Tuple[int, ...]]] = None,
206        seed: Optional[int] = None,
207        eps: float = 1e-7,
208    ):
209        lower_percentile_bounds = self._validate_bounds(lower_percentile_bounds, upper=False)
210        if upper_percentile_bounds is None:
211            upper_percentile_bounds = tuple(100.0 - bound for bound in reversed(lower_percentile_bounds))
212        upper_percentile_bounds = self._validate_bounds(upper_percentile_bounds, upper=True)
213        if distribution not in ("uniform", "normal"):
214            raise ValueError("distribution must be 'uniform' or 'normal'.")
215
216        if distribution == "uniform":
217            if distribution_kwargs is not None:
218                raise ValueError("Uniform sampling does not accept distribution_kwargs.")
219        else:
220            if not isinstance(distribution_kwargs, dict) or set(distribution_kwargs) != {"mean", "std"}:
221                raise ValueError("Normal sampling requires exactly the distribution_kwargs 'mean' and 'std'.")
222            mean, std = float(distribution_kwargs["mean"]), float(distribution_kwargs["std"])
223            if not np.isfinite(mean) or not lower_percentile_bounds[0] <= mean <= lower_percentile_bounds[1]:
224                raise ValueError("The normal distribution mean must be finite and within lower_percentile_bounds.")
225            if not np.isfinite(std) or std < 0.0:
226                raise ValueError("The normal distribution std must be finite and non-negative.")
227            distribution_kwargs = {"mean": mean, "std": std}
228
229        if rounding_decimals is not None and (
230            not isinstance(rounding_decimals, int) or isinstance(rounding_decimals, bool) or rounding_decimals < 0
231        ):
232            raise ValueError("rounding_decimals must be a non-negative integer or None.")
233        if not np.isfinite(eps) or eps <= 0.0:
234            raise ValueError("eps must be finite and greater than zero.")
235        if seed is not None:
236            if not isinstance(seed, (int, np.integer)) or isinstance(seed, bool):
237                raise TypeError("seed must be an integer or None.")
238            if seed < 0:
239                raise ValueError("seed must be non-negative.")
240            seed = int(seed)
241
242        self.lower_percentile_bounds = lower_percentile_bounds
243        self.upper_percentile_bounds = upper_percentile_bounds
244        self.distribution = distribution
245        self.distribution_kwargs = distribution_kwargs
246        self.rounding_decimals = rounding_decimals
247        self.axis = axis
248        self.seed = seed
249        self.eps = float(eps)
250        self._random_generator = None
251        self._random_generator_worker_id = None
252
253    @staticmethod
254    def _validate_bounds(values, upper):
255        name = "upper_percentile_bounds" if upper else "lower_percentile_bounds"
256        if not isinstance(values, (tuple, list)) or len(values) != 2:
257            raise ValueError(f"{name} must contain exactly two values.")
258        lower_bound, upper_bound = (float(value) for value in values)
259        finite = np.isfinite(lower_bound) and np.isfinite(upper_bound)
260        if upper:
261            valid = 50.0 < lower_bound <= upper_bound <= 100.0
262            interval = "(50, 100]"
263        else:
264            valid = 0.0 <= lower_bound <= upper_bound < 50.0
265            interval = "[0, 50)"
266        if not finite or not valid:
267            raise ValueError(f"{name} must be a finite interval within {interval}.")
268        return lower_bound, upper_bound
269
270    def _round(self, value):
271        return float(value) if self.rounding_decimals is None else round(float(value), self.rounding_decimals)
272
273    def _get_random_generator(self):
274        if self.seed is None:
275            return np.random
276
277        worker_info = torch.utils.data.get_worker_info()
278        worker_id = None if worker_info is None else worker_info.id
279        if self._random_generator is None or self._random_generator_worker_id != worker_id:
280            seed_sequence = np.random.SeedSequence([self.seed, 0 if worker_id is None else worker_id])
281            self._random_generator = np.random.default_rng(seed_sequence)
282            self._random_generator_worker_id = worker_id
283        return self._random_generator
284
285    def sample_percentiles(self) -> Tuple[float, float]:
286        """Sample and return a valid ``(lower, upper)`` percentile pair."""
287        random_generator = self._get_random_generator()
288        if self.distribution == "uniform":
289            lower = random_generator.uniform(*self.lower_percentile_bounds)
290            upper = random_generator.uniform(*self.upper_percentile_bounds)
291        else:
292            mean = self.distribution_kwargs["mean"]
293            std = self.distribution_kwargs["std"]
294            lower = mean if std == 0.0 else random_generator.normal(mean, std)
295            upper = 100.0 - (mean if std == 0.0 else random_generator.normal(mean, std))
296
297        # Normal distribution tails may leave the configured percentile interval.
298        lower = float(np.clip(self._round(lower), *self.lower_percentile_bounds))
299        upper = float(np.clip(self._round(upper), *self.upper_percentile_bounds))
300        return lower, upper
301
302    def __call__(self, raw: Union[np.ndarray, torch.tensor]) -> Union[np.ndarray, torch.tensor]:
303        lower, upper = self.sample_percentiles()
304        normalized = normalize_percentile(raw, lower=lower, upper=upper, axis=self.axis, eps=self.eps)
305        if torch.is_tensor(normalized):
306            return torch.clamp(normalized, min=0.0, max=1.0)
307        return np.clip(normalized, 0.0, 1.0)
308
309
310#
311# Intensity Augmentations / Noise Augmentations.
312#
313
314# modified from https://github.com/kreshuklab/spoco/blob/main/spoco/transforms.py
315class RandomContrast:
316    """Transformation to adjust contrast by scaling image to `mean + alpha * (image - mean)`.
317
318    Args:
319        alpha: Minimal and maximal alpha value for adjusting the contrast.
320            The value for the transformation will be drawn uniformly from the corresponding interval.
321        mean: Mean value for the image data.
322        clip_kwargs: Keyword arguments for clipping the data after the contrast augmentation.
323    """
324    def __init__(
325        self, alpha: Tuple[float, float] = (0.5, 2), mean: float = 0.5, clip_kwargs: Dict = {"a_min": 0, "a_max": 1}
326    ):
327        self.alpha = alpha
328        self.mean = mean
329        self.clip_kwargs = clip_kwargs
330
331    def __call__(self, img: np.ndarray) -> np.ndarray:
332        """Apply the augmentation to data.
333
334        Args:
335            img: The input image.
336
337        Returns:
338            The transformed image.
339        """
340        alpha = np.random.uniform(self.alpha[0], self.alpha[1])
341        result = self.mean + alpha * (img - self.mean)
342        if self.clip_kwargs:
343            return np.clip(result, **self.clip_kwargs)
344        return result
345
346
347class AdditiveGaussianNoise:
348    """Transformation to add random Gaussian noise to image.
349
350    Args:
351        scale: Scale for the noise.
352        clip_kwargs: Keyword arguments for clipping the data after the transformation.
353    """
354    def __init__(self, scale: Tuple[float, float] = (0.0, 0.3), clip_kwargs: Dict = {"a_min": 0, "a_max": 1}):
355        self.scale = scale
356        self.clip_kwargs = clip_kwargs
357
358    def __call__(self, img: np.ndarray) -> np.ndarray:
359        """Apply the augmentation to data.
360
361        Args:
362            img: The input image.
363
364        Returns:
365            The transformed image.
366        """
367        std = np.random.uniform(self.scale[0], self.scale[1])
368        gaussian_noise = np.random.normal(0, std, size=img.shape)
369
370        if self.clip_kwargs:
371            return np.clip(img + gaussian_noise, 0, 1)
372
373        return img + gaussian_noise
374
375
376class AdditivePoissonNoise:
377    """Transformation to add random additive Poisson noise to image.
378
379    Args:
380        lam: Lambda value for the Poisson transformation.
381        clip_kwargs: Keyword arguments for clipping the data after the transformation.
382    """
383    # Not sure if Poisson noise like this does make sense for data that is already normalized
384    def __init__(self, lam: Tuple[float, float] = (0.0, 0.1), clip_kwargs: Dict = {"a_min": 0, "a_max": 1}):
385        self.lam = lam
386        self.clip_kwargs = clip_kwargs
387
388    def __call__(self, img: np.ndarray) -> np.ndarray:
389        """Apply the augmentation to data.
390
391        Args:
392            img: The input image.
393
394        Returns:
395            The transformed image.
396        """
397        lam = np.random.uniform(self.lam[0], self.lam[1])
398        poisson_noise = np.random.poisson(lam, size=img.shape) / lam
399        if self.clip_kwargs:
400            return np.clip(img + poisson_noise, 0, 1)
401        return img + poisson_noise
402
403
404class PoissonNoise:
405    """Transformation to add random data-dependant Poisson noise to image.
406
407    Args:
408        multiplier: Multiplicative factors for deriving the lambda factor from the data.
409            The factor used for the transformation will be uniformly sampled form the range of this parameter.
410        clip_kwargs: Keyword arguments for clipping the data after the transformation.
411    """
412    def __init__(self, multiplier: Tuple[float, float] = (5.0, 10.0), clip_kwargs: Dict = {"a_min": 0, "a_max": 1}):
413        self.multiplier = multiplier
414        self.clip_kwargs = clip_kwargs
415
416    def __call__(self, img: np.ndarray) -> np.ndarray:
417        """Apply the augmentation to data.
418
419        Args:
420            img: The input image.
421
422        Returns:
423            The transformed image.
424        """
425        multiplier = np.random.uniform(self.multiplier[0], self.multiplier[1])
426        offset = img.min()
427        poisson_noise = np.random.poisson((img - offset) * multiplier)
428
429        if isinstance(img, torch.Tensor):
430            poisson_noise = torch.Tensor(poisson_noise)
431        poisson_noise = poisson_noise / multiplier + offset
432
433        if self.clip_kwargs:
434            return np.clip(poisson_noise, **self.clip_kwargs)
435        return poisson_noise
436
437
438class GaussianBlur:
439    """Transformation to blur the image with a randomly drawn sigma value.
440
441    Args:
442        sigma: The sigma value for the transformation.
443            The value used in the transformation will be uniformly drawn from the range specified here.
444    """
445    def __init__(self, sigma: Tuple[float, float] = (0.0, 3.0)):
446        self.sigma = sigma
447
448    def __call__(self, img: np.ndarray) -> np.ndarray:
449        """Apply the augmentation to data.
450
451        Args:
452            img: The input image.
453
454        Returns:
455            The transformed image.
456        """
457        # Sample the sigma value. Note that we switch the bounds to ensure zero is excluded from sampling.
458        sigma = np.random.uniform(self.sigma[1], self.sigma[0])
459        # Determine the kernel size based on the sigma value.
460        kernel_size = int(2 * np.ceil(3 * sigma) + 1)
461        if isinstance(img, np.ndarray):
462            img = torch.from_numpy(img)
463
464        return transforms.GaussianBlur(kernel_size, sigma=sigma)(img)
465
466
467#
468# Default Transformation: Apply intensity augmentations and normalize.
469#
470
471class RawTransform:
472    """The transformation for raw data during training.
473
474    Args:
475        normalizer: The normalization function.
476        augmentation1: Intensity augmentation applied before the normalization.
477        augmentation2: Intensity augmentation applied after the normalization.
478    """
479    def __init__(
480        self, normalizer: Callable, augmentation1: Optional[Callable] = None, augmentation2: Optional[Callable] = None
481    ):
482        self.normalizer = normalizer
483        self.augmentation1 = augmentation1
484        self.augmentation2 = augmentation2
485
486    def __call__(self, raw: np.ndarray) -> np.ndarray:
487        """Apply the raw transformation.
488
489        Args:
490            raw: The raw data.
491
492        Returns:
493            The transformed raw data.
494        """
495        if self.augmentation1 is not None:
496            raw = self.augmentation1(raw)
497
498        raw = self.normalizer(raw)
499
500        if self.augmentation2 is not None:
501            raw = self.augmentation2(raw)
502        return raw
503
504
505def get_raw_transform(
506    normalizer: Callable = standardize,
507    augmentation1: Optional[Callable] = None,
508    augmentation2: Optional[Callable] = None
509) -> Callable:
510    """Get the raw transformation.
511
512    Args:
513        normalizer: The normalization function.
514        augmentation1: Intensity augmentation applied before the normalization.
515        augmentation2: Intensity augmentation applied after the normalization.
516
517    Returns:
518        The raw transformation.
519    """
520    return RawTransform(normalizer, augmentation1=augmentation1, augmentation2=augmentation2)
521
522
523def get_raw_transform_for_pretrained_vit(backbone: Optional[str]) -> Tuple[Optional[Callable], Optional[int]]:
524    """Get the raw transformation and clipping limit for a pretrained ViT backbone.
525
526    Args:
527        backbone: The pretrained ViT backbone. Options: "sam", "sam2", "sam3", "cellpose_sam", "mae", "scalemae",
528            "dinov2", "dinov3", "torchvision", or None.
529
530    Returns:
531        The raw transformation. It normalizes percentiles to [0, 255] for "sam" and [0, 1] for the other backbones.
532        The maximum value `clip_max` for intensity augmentations. Use it to clip their output to [0, clip_max].
533            Both values are None if `backbone` is None, so data loaders use the default standardization.
534    """
535    if backbone is None:
536        return None, None
537
538    from ..model.util import VIT_CONFIGS
539
540    if backbone not in VIT_CONFIGS:
541        raise ValueError(f"Unsupported backbone: {backbone}.")
542
543    clip_max = VIT_CONFIGS[backbone].clip_max
544    normalizer = partial(_normalize_percentile_to_range, clip_max=clip_max)
545    return get_raw_transform(normalizer=normalizer), clip_max
546
547
548def get_default_mean_teacher_augmentations(
549    p: float = 0.3,
550    norm: Optional[Callable] = None,
551    blur_kwargs: Optional[Dict] = None,
552    poisson_kwargs: Optional[Dict] = None,
553    gaussian_kwargs: Optional[Dict] = None,
554) -> Callable:
555    """Get the default augmentations for mean teacher training.
556
557    The default values for the augmentations are designed for an image with pixel values in range [0, 1].
558    By default, a normalization transformation is applied for this reason.
559
560    Args:
561        p: The probability for applying the individual intensity transformations.
562        norm: The noromaization function.
563        blur_kwargs: The keyword arguments for `GaussianBlur`.
564        poisson_kwargs: The keyword arguments for `PoissonNoise`.
565        gaussian_kwargs: The keyword arguments for `AdditiveGaussianNoise`.
566
567    Returns:
568        The raw transformation with augmentations.
569    """
570    if norm is None:
571        norm = normalize
572
573    aug1 = transforms.Compose([
574        norm,
575        transforms.RandomApply([GaussianBlur(**({} if blur_kwargs is None else blur_kwargs))], p=p),
576        transforms.RandomApply([PoissonNoise(**({} if poisson_kwargs is None else poisson_kwargs))], p=p/2),
577        transforms.RandomApply([AdditiveGaussianNoise(**({} if gaussian_kwargs is None else gaussian_kwargs))], p=p/2),
578    ])
579
580    aug2 = transforms.RandomApply([RandomContrast(clip_kwargs={"a_min": 0, "a_max": 1})], p=p)
581    return get_raw_transform(normalizer=norm, augmentation1=aug1, augmentation2=aug2)