torch_em.transform.raw

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