File size: 3,934 Bytes
6279422 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """
Data augmentation for infrared thermal images.
Uses Albumentations for efficient, GPU-friendly augmentation
with thermal-image-appropriate transforms.
"""
import numpy as np
import albumentations as A
from typing import Optional
class ThermalAugmentor:
"""
Augmentation pipeline tailored for thermal images.
Default transforms: rotation, horizontal flip, brightness/contrast
shift, Gaussian noise, and random crop-resize.
"""
def __init__(
self,
rotation_limit: int = 15,
horizontal_flip: bool = True,
vertical_flip: bool = False,
brightness_limit: float = 0.1,
contrast_limit: float = 0.1,
image_size: tuple = (224, 224),
enabled: bool = True,
):
self.enabled = enabled
self.image_size = image_size
if not self.enabled:
self.transform = A.Compose([A.NoOp()])
return
self.transform = A.Compose([
A.Rotate(limit=rotation_limit, border_mode=0, p=0.5),
A.HorizontalFlip(p=0.5 if horizontal_flip else 0.0),
A.VerticalFlip(p=0.5 if vertical_flip else 0.0),
A.RandomBrightnessContrast(
brightness_limit=brightness_limit,
contrast_limit=contrast_limit,
p=0.5,
),
A.GaussNoise(p=0.3),
A.RandomResizedCrop(
size=image_size,
scale=(0.85, 1.0),
ratio=(0.9, 1.1),
p=0.3,
),
])
@classmethod
def from_config(cls, config) -> "ThermalAugmentor":
"""Create augmentor from a Config object."""
aug = config.augmentation
return cls(
rotation_limit=aug.rotation_limit,
horizontal_flip=aug.horizontal_flip,
vertical_flip=aug.vertical_flip,
brightness_limit=aug.brightness_limit,
contrast_limit=aug.contrast_limit,
image_size=tuple(config.data.image_size),
enabled=aug.enabled,
)
def __call__(self, image: np.ndarray) -> np.ndarray:
"""
Apply augmentation pipeline to a single image.
Args:
image: Grayscale float32 image in [0, 1], shape (H, W).
Returns:
Augmented image, same shape and range.
"""
if not self.enabled:
return image
# Albumentations expects uint8 or float32 with shape (H, W) or (H, W, C)
if image.ndim == 2:
image = image[:, :, np.newaxis] # (H, W, 1)
result = self.transform(image=image)
augmented = result["image"]
if augmented.ndim == 3 and augmented.shape[2] == 1:
augmented = augmented[:, :, 0]
# Clamp to [0, 1]
augmented = np.clip(augmented, 0.0, 1.0)
return augmented.astype(np.float32)
def augment_sequence(self, images: list) -> list:
"""
Apply the *same* augmentation to every image in a sequence.
Uses a shared random seed so every frame receives identical
spatial transforms (important for temporal consistency).
Args:
images: List of grayscale float32 images.
Returns:
List of augmented images.
"""
if not self.enabled or len(images) == 0:
return images
# Pick a random seed for this sequence
seed = np.random.randint(0, 2**31)
augmented_images = []
for img in images:
# Reset random state so each frame gets identical transforms
import random
random.seed(seed)
np.random.seed(seed)
augmented_images.append(self(img))
# Restore randomness
np.random.seed(None)
return augmented_images
|