File size: 1,958 Bytes
2fe488a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """Simple image augmentations for training."""
import numpy as np
import torch
def brightness_augment(
imgs: torch.Tensor,
coords: torch.Tensor,
masks: torch.Tensor,
*,
rng: np.random.Generator,
shift_range: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Random additive brightness shift.
Parameters
----------
imgs : torch.Tensor
(W, *spatial) normalised images.
coords : torch.Tensor
(W, M, 3) node coordinates — passed through unchanged.
masks : torch.Tensor
(W, M) boolean masks — passed through unchanged.
rng : np.random.Generator
Random number generator.
shift_range : float
Maximum absolute shift sampled from [-shift_range, shift_range].
"""
shift = rng.uniform(-shift_range, shift_range)
return imgs + shift, coords, masks
def flip_augment(
imgs: torch.Tensor,
coords: torch.Tensor,
masks: torch.Tensor,
*,
rng: np.random.Generator,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Random spatial flip: samples uniformly from all 8 axis-aligned symmetries.
Each of Z, Y, X is independently flipped with probability 0.5,
giving 2^3 = 8 equally likely outcomes (including identity).
Creates a copy of coords so the underlying dataset is not mutated.
Only real (non-padded) coordinates are flipped; padding stays at zero.
"""
flip_mask = rng.random(3) < 0.5 # (Z, Y, X)
dims_to_flip = [1 + dim for dim, flip in enumerate(flip_mask) if flip]
if not dims_to_flip:
return imgs, coords, masks
imgs = imgs.flip(dims=dims_to_flip)
coords = coords.clone()
shape = imgs.shape[1:] # (Z, Y, X)
for dim in range(3):
if flip_mask[dim]:
dim_coords = coords[..., dim]
dim_coords[masks] = shape[dim] - dim_coords[masks] - 1
coords[..., dim] = dim_coords
return imgs, coords, masks
|