| """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.""" |
| 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 while preserving valid node coordinates.""" |
|
|
| flip_mask = rng.random(3) < 0.5 |
| 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:] |
|
|
| if coords.ndim == 2: |
| |
| |
| |
| valid = masks.any(dim=1) |
|
|
| for dim in range(3): |
| if flip_mask[dim]: |
| coords[valid, dim] = ( |
| shape[dim] - coords[valid, dim] - 1 |
| ) |
|
|
| elif coords.ndim == 3: |
| |
| |
| |
| valid = masks.bool() |
|
|
| for dim in range(3): |
| if flip_mask[dim]: |
| coords[..., dim] = torch.where( |
| valid, |
| shape[dim] - coords[..., dim] - 1, |
| coords[..., dim], |
| ) |
|
|
| else: |
| raise ValueError( |
| f"Unexpected coords shape: {coords.shape}" |
| ) |
|
|
| return imgs, coords, masks |
|
|
|
|
| |
|
|
| import torch.nn.functional as _Fg |
|
|
|
|
| def noise_augment(imgs, coords, masks, *, rng, sigma_frac=0.5): |
| s = rng.uniform(0.0, sigma_frac) * float(imgs.std()) |
| return imgs + torch.randn_like(imgs) * s, coords, masks |
|
|
|
|
| def contrast_augment(imgs, coords, masks, *, rng, lo=0.6, hi=1.6): |
| c = rng.uniform(lo, hi) |
| mu = float(imgs.mean()) |
| return (imgs - mu) * c + mu, coords, masks |
|
|
|
|
| def gamma_augment(imgs, coords, masks, *, rng, lo=0.6, hi=1.6): |
| g = rng.uniform(lo, hi) |
| return imgs.clamp(min=0).pow(g), coords, masks |
|
|
|
|
| def blur_augment(imgs, coords, masks, *, rng, max_k=3): |
| k = int(rng.integers(1, max_k + 1)) | 1 |
|
|
| if k == 1: |
| return imgs, coords, masks |
|
|
| x = imgs.unsqueeze(1).float() |
| p = k // 2 |
|
|
| x = _Fg.avg_pool3d( |
| _Fg.pad( |
| x, |
| (p, p, p, p, p, p), |
| mode="replicate" |
| ), |
| k, |
| stride=1, |
| ) |
|
|
| return x[:, 0].to(imgs.dtype), coords, masks |
|
|
|
|
| generalization_augments = [ |
| brightness_augment, |
| flip_augment, |
| noise_augment, |
| contrast_augment, |
| gamma_augment, |
| blur_augment, |
| ] |
|
|
