refactor(transforms): several updates

This commit is contained in:
Steven Palma
2026-07-29 15:43:13 +02:00
parent 2422d43fb1
commit 117216b29b
2 changed files with 331 additions and 105 deletions
+231 -95
View File
@@ -20,6 +20,7 @@ from dataclasses import dataclass, field
from typing import Any from typing import Any
import torch import torch
from torchvision.io import decode_image, encode_jpeg
from torchvision.transforms import v2 from torchvision.transforms import v2
from torchvision.transforms.v2 import ( from torchvision.transforms.v2 import (
Transform, Transform,
@@ -159,15 +160,24 @@ class GaussianNoise(Transform):
super().__init__() super().__init__()
if isinstance(std, (int, float)): if isinstance(std, (int, float)):
self.std = (0.0, float(std)) self.std = (0.0, float(std))
else: elif isinstance(std, Sequence) and len(std) == 2:
self.std = (float(std[0]), float(std[1])) self.std = (float(std[0]), float(std[1]))
else:
raise TypeError("std must be a number or a sequence with length 2.")
if not 0.0 <= self.std[0] <= self.std[1]:
raise ValueError(f"std must satisfy 0 <= min <= max, but got {self.std}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {"std": torch.empty(1).uniform_(self.std[0], self.std[1]).item()} return {
"std": torch.empty(1).uniform_(self.std[0], self.std[1]).item(),
"seed": torch.randint(0, torch.iinfo(torch.int64).max, ()).item(),
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any: def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if isinstance(inpt, torch.Tensor) and inpt.is_floating_point(): if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
return (inpt + torch.randn_like(inpt) * (params["std"] / 255.0)).clamp(0.0, 1.0) generator = torch.Generator(device=inpt.device).manual_seed(params["seed"])
noise = torch.randn(inpt.shape, device=inpt.device, dtype=inpt.dtype, generator=generator)
return (inpt + noise * (params["std"] / 255.0)).clamp(0.0, 1.0)
return inpt return inpt
@@ -177,43 +187,52 @@ class MotionBlur(Transform):
Generates a 1D averaging kernel along a random direction, applied via depthwise convolution. Generates a 1D averaging kernel along a random direction, applied via depthwise convolution.
Args: Args:
kernel_size: Range (min, max) for blur kernel size. Will be forced odd. kernel_size: An odd kernel size or a range containing at least one odd kernel size.
""" """
def __init__(self, kernel_size: int | Sequence[int] = (3, 11)) -> None: def __init__(self, kernel_size: int | Sequence[int] = (3, 11)) -> None:
super().__init__() super().__init__()
if isinstance(kernel_size, int): if isinstance(kernel_size, int):
self.kernel_size = (kernel_size, kernel_size) self.kernel_size = (kernel_size, kernel_size)
else: elif isinstance(kernel_size, Sequence) and len(kernel_size) == 2:
self.kernel_size = (int(kernel_size[0]), int(kernel_size[1])) self.kernel_size = (int(kernel_size[0]), int(kernel_size[1]))
else:
raise TypeError("kernel_size must be an int or a sequence with length 2.")
if not 1 <= self.kernel_size[0] <= self.kernel_size[1]:
raise ValueError(f"kernel_size must satisfy 1 <= min <= max, but got {self.kernel_size}.")
self._first_odd_kernel_size = self.kernel_size[0] + (self.kernel_size[0] + 1) % 2
if self._first_odd_kernel_size > self.kernel_size[1]:
raise ValueError(f"kernel_size range must contain an odd value, but got {self.kernel_size}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
ks = int(torch.randint(self.kernel_size[0], self.kernel_size[1] + 1, (1,)).item()) num_odd_sizes = (self.kernel_size[1] - self._first_odd_kernel_size) // 2 + 1
if ks % 2 == 0: size_index = int(torch.randint(0, num_odd_sizes, ()).item())
ks += 1 ks = self._first_odd_kernel_size + 2 * size_index
angle = torch.empty(1).uniform_(0, 360).item() angle = torch.empty(1).uniform_(0, 360).item()
return {"kernel_size": ks, "angle": angle} return {"kernel_size": ks, "angle": angle}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any: def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point(): if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt return inpt
ks = params["kernel_size"] if inpt.ndim < 3:
rad = params["angle"] * math.pi / 180 raise ValueError(f"MotionBlur expects [..., C, H, W] input, but got shape {inpt.shape}.")
cos_a, sin_a = abs(math.cos(rad)), abs(math.sin(rad))
x = inpt.unsqueeze(0) if inpt.dim() == 3 else inpt kernel_size = params["kernel_size"]
if cos_a > sin_a: radius = kernel_size // 2
out = torch.nn.functional.avg_pool2d( angle = math.radians(params["angle"])
torch.nn.functional.pad(x, (ks // 2, ks // 2, 0, 0), mode="replicate"), positions = torch.linspace(-radius, radius, kernel_size, device=inpt.device)
(1, ks), x_coords = (positions * math.cos(angle)).round().to(torch.long) + radius
stride=1, y_coords = (positions * math.sin(angle)).round().to(torch.long) + radius
) kernel = torch.zeros((kernel_size, kernel_size), device=inpt.device, dtype=inpt.dtype)
else: kernel[y_coords, x_coords] = 1
out = torch.nn.functional.avg_pool2d( kernel /= kernel.sum()
torch.nn.functional.pad(x, (0, 0, ks // 2, ks // 2), mode="replicate"),
(ks, 1), channels, height, width = inpt.shape[-3:]
stride=1, flat_input = inpt.reshape(-1, channels, height, width)
) depthwise_kernel = kernel.expand(channels, 1, kernel_size, kernel_size)
return (out.squeeze(0) if inpt.dim() == 3 else out).clamp(0.0, 1.0) padded = torch.nn.functional.pad(flat_input, (radius,) * 4, mode="replicate")
output = torch.nn.functional.conv2d(padded, depthwise_kernel, groups=channels)
return output.reshape(inpt.shape).clamp(0.0, 1.0)
class JPEGCompression(Transform): class JPEGCompression(Transform):
@@ -229,8 +248,12 @@ class JPEGCompression(Transform):
super().__init__() super().__init__()
if isinstance(quality, int): if isinstance(quality, int):
self.quality = (quality, quality) self.quality = (quality, quality)
else: elif isinstance(quality, Sequence) and len(quality) == 2:
self.quality = (int(quality[0]), int(quality[1])) self.quality = (int(quality[0]), int(quality[1]))
else:
raise TypeError("quality must be an int or a sequence with length 2.")
if not 1 <= self.quality[0] <= self.quality[1] <= 100:
raise ValueError(f"quality must satisfy 1 <= min <= max <= 100, but got {self.quality}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {"quality": int(torch.randint(self.quality[0], self.quality[1] + 1, (1,)).item())} return {"quality": int(torch.randint(self.quality[0], self.quality[1] + 1, (1,)).item())}
@@ -238,16 +261,20 @@ class JPEGCompression(Transform):
def transform(self, inpt: Any, params: dict[str, Any]) -> Any: def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point(): if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt return inpt
from torchvision.io import decode_image, encode_jpeg if inpt.ndim < 3:
raise ValueError(f"JPEGCompression expects [..., C, H, W] input, but got shape {inpt.shape}.")
img_uint8 = (inpt * 255).byte() channels, height, width = inpt.shape[-3:]
if img_uint8.dim() == 3: if channels not in (1, 3):
try: raise ValueError(f"JPEGCompression expects 1 or 3 channels, but got {channels}.")
buf = encode_jpeg(img_uint8.cpu(), quality=params["quality"])
return decode_image(buf).to(device=inpt.device, dtype=inpt.dtype) / 255.0 flat_input = inpt.reshape(-1, channels, height, width)
except Exception: flat_uint8 = (flat_input.clamp(0.0, 1.0) * 255).round().to(torch.uint8).cpu()
return inpt decoded_frames = [
return inpt decode_image(encode_jpeg(frame, quality=params["quality"])) for frame in flat_uint8.unbind()
]
output = torch.stack(decoded_frames).to(device=inpt.device, dtype=inpt.dtype) / 255.0
return output.reshape(inpt.shape)
class GaussianPatchBrightness(Transform): class GaussianPatchBrightness(Transform):
@@ -271,10 +298,22 @@ class GaussianPatchBrightness(Transform):
super().__init__() super().__init__()
if isinstance(num_patches, int): if isinstance(num_patches, int):
self.num_patches = (num_patches, num_patches) self.num_patches = (num_patches, num_patches)
else: elif isinstance(num_patches, Sequence) and len(num_patches) == 2:
self.num_patches = (int(num_patches[0]), int(num_patches[1])) self.num_patches = (int(num_patches[0]), int(num_patches[1]))
self.sigma_range = sigma_range else:
self.factor_range = factor_range raise TypeError("num_patches must be an int or a sequence with length 2.")
if not 1 <= self.num_patches[0] <= self.num_patches[1]:
raise ValueError(f"num_patches must satisfy 1 <= min <= max, but got {self.num_patches}.")
if not isinstance(sigma_range, Sequence) or len(sigma_range) != 2:
raise TypeError("sigma_range must be a sequence with length 2.")
self.sigma_range = (float(sigma_range[0]), float(sigma_range[1]))
if not 0.0 < self.sigma_range[0] <= self.sigma_range[1]:
raise ValueError(f"sigma_range must satisfy 0 < min <= max, but got {self.sigma_range}.")
if not isinstance(factor_range, Sequence) or len(factor_range) != 2:
raise TypeError("factor_range must be a sequence with length 2.")
self.factor_range = (float(factor_range[0]), float(factor_range[1]))
if not 0.0 <= self.factor_range[0] <= self.factor_range[1]:
raise ValueError(f"factor_range must satisfy 0 <= min <= max, but got {self.factor_range}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
n = int(torch.randint(self.num_patches[0], self.num_patches[1] + 1, (1,)).item()) n = int(torch.randint(self.num_patches[0], self.num_patches[1] + 1, (1,)).item())
@@ -297,9 +336,8 @@ class GaussianPatchBrightness(Transform):
): ):
gauss = torch.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sigma**2)) gauss = torch.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sigma**2))
mask = mask * (1.0 + (factor - 1.0) * gauss) mask = mask * (1.0 + (factor - 1.0) * gauss)
if inpt.dim() == 3: broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
return (inpt * mask.unsqueeze(0)).clamp(0.0, 1.0) return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
return (inpt * mask.unsqueeze(0).unsqueeze(0)).clamp(0.0, 1.0)
class RandomShadow(Transform): class RandomShadow(Transform):
@@ -316,33 +354,44 @@ class RandomShadow(Transform):
super().__init__() super().__init__()
if isinstance(opacity, (int, float)): if isinstance(opacity, (int, float)):
self.opacity = (float(opacity), float(opacity)) self.opacity = (float(opacity), float(opacity))
else: elif isinstance(opacity, Sequence) and len(opacity) == 2:
self.opacity = (float(opacity[0]), float(opacity[1])) self.opacity = (float(opacity[0]), float(opacity[1]))
else:
raise TypeError("opacity must be a number or a sequence with length 2.")
if not 0.0 <= self.opacity[0] <= self.opacity[1] <= 1.0:
raise ValueError(f"opacity must satisfy 0 <= min <= max <= 1, but got {self.opacity}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {"opacity": torch.empty(1).uniform_(self.opacity[0], self.opacity[1]).item()} return {
"opacity": torch.empty(1).uniform_(self.opacity[0], self.opacity[1]).item(),
"start": torch.rand(1).item(),
"width": torch.empty(1).uniform_(1 / 3, 2 / 3).item(),
"direction": -1.0 if torch.rand(1).item() < 0.5 else 1.0,
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any: def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point(): if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt return inpt
if inpt.ndim < 3:
raise ValueError(f"RandomShadow expects [..., C, H, W] input, but got shape {inpt.shape}.")
h, w = inpt.shape[-2:] h, w = inpt.shape[-2:]
x_start = int(torch.randint(0, w // 2, (1,)).item()) band_width = max(1, min(w, round(params["width"] * w)))
x_end = int(torch.randint(w // 3, w, (1,)).item()) x_start = round(params["start"] * (w - band_width))
x_end = x_start + band_width
mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype) mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
if torch.rand(1).item() < 0.5: mask[:, x_start:x_end] = 1.0 + params["direction"] * params["opacity"]
mask[:, x_start:x_end] = 1.0 - params["opacity"]
else: smoothing_size = min(8, h, w)
mask[:, x_start:x_end] = 1.0 + params["opacity"] if smoothing_size > 1:
mask = mask.unsqueeze(0).unsqueeze(0) batched_mask = mask[None, None]
small = torch.nn.functional.avg_pool2d(mask, 8, stride=8) small = torch.nn.functional.avg_pool2d(batched_mask, smoothing_size, stride=smoothing_size)
mask = ( mask = torch.nn.functional.interpolate(small, size=(h, w), mode="bilinear", align_corners=False)[
torch.nn.functional.interpolate(small, size=(h, w), mode="bilinear", align_corners=False) 0, 0
.squeeze(0) ]
.squeeze(0)
) broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
if inpt.dim() == 3: return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
return (inpt * mask.unsqueeze(0)).clamp(0.0, 1.0)
return (inpt * mask.unsqueeze(0).unsqueeze(0)).clamp(0.0, 1.0)
class CoarseDropout(Transform): class CoarseDropout(Transform):
@@ -366,6 +415,16 @@ class CoarseDropout(Transform):
fill_value: float = 0.0, fill_value: float = 0.0,
) -> None: ) -> None:
super().__init__() super().__init__()
if not isinstance(max_holes, int):
raise TypeError("max_holes must be an int.")
if max_holes < 1:
raise ValueError(f"max_holes must be at least 1, but got {max_holes}.")
if not 0.0 < max_height_frac <= 1.0:
raise ValueError(f"max_height_frac must be in (0, 1], but got {max_height_frac}.")
if not 0.0 < max_width_frac <= 1.0:
raise ValueError(f"max_width_frac must be in (0, 1], but got {max_width_frac}.")
if not 0.0 <= fill_value <= 1.0:
raise ValueError(f"fill_value must be in [0, 1], but got {fill_value}.")
self.max_holes = max_holes self.max_holes = max_holes
self.max_height_frac = max_height_frac self.max_height_frac = max_height_frac
self.max_width_frac = max_width_frac self.max_width_frac = max_width_frac
@@ -373,22 +432,27 @@ class CoarseDropout(Transform):
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
n = int(torch.randint(1, self.max_holes + 1, (1,)).item()) n = int(torch.randint(1, self.max_holes + 1, (1,)).item())
return {"n_holes": n} sizes = torch.rand(n, 2)
sizes[:, 0] *= self.max_height_frac
sizes[:, 1] *= self.max_width_frac
return {"sizes": sizes.tolist(), "positions": torch.rand(n, 2).tolist()}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any: def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point(): if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt return inpt
if inpt.ndim < 3:
raise ValueError(f"CoarseDropout expects [..., C, H, W] input, but got shape {inpt.shape}.")
h, w = inpt.shape[-2:] h, w = inpt.shape[-2:]
result = inpt.clone() result = inpt.clone()
for _ in range(params["n_holes"]): for (height_frac, width_frac), (y_frac, x_frac) in zip(
hole_h = int(torch.randint(1, max(2, int(h * self.max_height_frac)), (1,)).item()) params["sizes"], params["positions"], strict=True
hole_w = int(torch.randint(1, max(2, int(w * self.max_width_frac)), (1,)).item()) ):
y = int(torch.randint(0, h - hole_h + 1, (1,)).item()) hole_h = max(1, min(h, round(height_frac * h)))
x = int(torch.randint(0, w - hole_w + 1, (1,)).item()) hole_w = max(1, min(w, round(width_frac * w)))
if result.dim() == 3: y = round(y_frac * (h - hole_h))
result[:, y : y + hole_h, x : x + hole_w] = self.fill_value x = round(x_frac * (w - hole_w))
else: result[..., y : y + hole_h, x : x + hole_w] = self.fill_value
result[:, :, y : y + hole_h, x : x + hole_w] = self.fill_value
return result return result
@@ -406,9 +470,16 @@ class GammaCorrection(Transform):
def __init__(self, gamma: float | Sequence[float] = (0.5, 2.0)) -> None: def __init__(self, gamma: float | Sequence[float] = (0.5, 2.0)) -> None:
super().__init__() super().__init__()
if isinstance(gamma, (int, float)): if isinstance(gamma, (int, float)):
self.gamma = (1.0 / float(gamma), float(gamma)) gamma = float(gamma)
else: if gamma <= 0:
raise ValueError(f"gamma must be positive, but got {gamma}.")
self.gamma = (min(gamma, 1.0 / gamma), max(gamma, 1.0 / gamma))
elif isinstance(gamma, Sequence) and len(gamma) == 2:
self.gamma = (float(gamma[0]), float(gamma[1])) self.gamma = (float(gamma[0]), float(gamma[1]))
else:
raise TypeError("gamma must be a number or a sequence with length 2.")
if not 0.0 < self.gamma[0] <= self.gamma[1]:
raise ValueError(f"gamma must satisfy 0 < min <= max, but got {self.gamma}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
log_lo = math.log(self.gamma[0]) log_lo = math.log(self.gamma[0])
@@ -422,43 +493,112 @@ class GammaCorrection(Transform):
return inpt return inpt
# From the paper authors' MIT-licensed reference implementation:
# https://github.com/TheZino/PlanckianJitter
_PLANCKIAN_BLACKBODY_COEFFICIENTS = (
(0.6743, 0.4029, 0.0013),
(0.6281, 0.4241, 0.1665),
(0.5919, 0.4372, 0.2513),
(0.5623, 0.4457, 0.3154),
(0.5376, 0.4515, 0.3672),
(0.5163, 0.4555, 0.4103),
(0.4979, 0.4584, 0.4468),
(0.4816, 0.4604, 0.4782),
(0.4672, 0.4619, 0.5053),
(0.4542, 0.4630, 0.5289),
(0.4426, 0.4638, 0.5497),
(0.4320, 0.4644, 0.5681),
(0.4223, 0.4648, 0.5844),
(0.4135, 0.4651, 0.5990),
(0.4054, 0.4653, 0.6121),
(0.3980, 0.4654, 0.6239),
(0.3911, 0.4655, 0.6346),
(0.3847, 0.4656, 0.6444),
(0.3787, 0.4656, 0.6532),
(0.3732, 0.4656, 0.6613),
(0.3680, 0.4655, 0.6688),
(0.3632, 0.4655, 0.6756),
(0.3586, 0.4655, 0.6820),
(0.3544, 0.4654, 0.6878),
(0.3503, 0.4653, 0.6933),
)
_PLANCKIAN_MIN_TEMPERATURE = 3_000
_PLANCKIAN_MAX_TEMPERATURE = 15_000
_PLANCKIAN_TEMPERATURE_STEP = 500
class PlanckianJitter(Transform): class PlanckianJitter(Transform):
"""Simulate color temperature shift along the Planckian locus. """Simulate color temperature shift along the Planckian locus.
Models the visual effect of different light sources (LED vs fluorescent vs Samples one black-body temperature and applies the corresponding correlated red
daylight) by applying physically-motivated per-channel scaling. More accurate and blue channel scaling while preserving the green channel. Coefficients between
than arbitrary hue shift for lighting variation. the tabulated 500 K intervals are linearly interpolated.
Reference: Zini et al., "Planckian Jitter", CVPR 2022 Workshop. Reference: Zini et al., "Planckian Jitter", CVPR 2022 Workshop.
Args: Args:
strength: Range (min, max) for per-channel scale factor. temperature: A fixed color temperature or range in Kelvin. Supported values
are between 3000 K and 15000 K.
""" """
def __init__(self, strength: Sequence[float] = (0.85, 1.15)) -> None: def __init__(self, temperature: int | Sequence[int] = (3_000, 15_000)) -> None:
super().__init__() super().__init__()
self.strength = strength if isinstance(temperature, int):
self.temperature = (temperature, temperature)
elif isinstance(temperature, Sequence) and len(temperature) == 2:
self.temperature = (int(temperature[0]), int(temperature[1]))
else:
raise TypeError("temperature must be an int or a sequence with length 2.")
if not (
_PLANCKIAN_MIN_TEMPERATURE
<= self.temperature[0]
<= self.temperature[1]
<= _PLANCKIAN_MAX_TEMPERATURE
):
raise ValueError(
"temperature must satisfy "
f"{_PLANCKIAN_MIN_TEMPERATURE} <= min <= max <= {_PLANCKIAN_MAX_TEMPERATURE}, "
f"but got {self.temperature}."
)
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {"scale": torch.empty(3).uniform_(self.strength[0], self.strength[1]).tolist()} temperature = int(torch.randint(self.temperature[0], self.temperature[1] + 1, ()).item())
return {"temperature": temperature}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any: def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point(): if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt return inpt
scale = torch.tensor(params["scale"], device=inpt.device, dtype=inpt.dtype) if inpt.ndim < 3 or inpt.shape[-3] != 3:
if inpt.dim() == 3: raise ValueError(f"PlanckianJitter expects [..., 3, H, W] input, but got shape {inpt.shape}.")
return (inpt * scale.view(3, 1, 1)).clamp(0.0, 1.0)
return (inpt * scale.view(1, 3, 1, 1)).clamp(0.0, 1.0) table_position = (params["temperature"] - _PLANCKIAN_MIN_TEMPERATURE) / _PLANCKIAN_TEMPERATURE_STEP
left_index = math.floor(table_position)
right_index = min(left_index + 1, len(_PLANCKIAN_BLACKBODY_COEFFICIENTS) - 1)
interpolation_weight = table_position - left_index
left = torch.tensor(
_PLANCKIAN_BLACKBODY_COEFFICIENTS[left_index],
device=inpt.device,
dtype=inpt.dtype,
)
right = torch.tensor(
_PLANCKIAN_BLACKBODY_COEFFICIENTS[right_index],
device=inpt.device,
dtype=inpt.dtype,
)
coefficients = torch.lerp(left, right, interpolation_weight)
scale = torch.stack(
(
coefficients[0] / coefficients[1],
coefficients.new_tensor(1.0),
coefficients[2] / coefficients[1],
)
)
broadcast_shape = (1,) * (inpt.ndim - 3) + (3, 1, 1)
return (inpt * scale.reshape(broadcast_shape)).clamp(0.0, 1.0)
# Custom transform registry for make_transform_from_config _CUSTOM_TRANSFORMS: dict[str, type[Transform]] = {
_CUSTOM_TRANSFORMS: dict[str, type] = {}
def _register_custom_transforms() -> None:
"""Register all custom transforms defined in this module."""
_CUSTOM_TRANSFORMS.update(
{
"SharpnessJitter": SharpnessJitter, "SharpnessJitter": SharpnessJitter,
"GaussianNoise": GaussianNoise, "GaussianNoise": GaussianNoise,
"MotionBlur": MotionBlur, "MotionBlur": MotionBlur,
@@ -468,11 +608,7 @@ def _register_custom_transforms() -> None:
"CoarseDropout": CoarseDropout, "CoarseDropout": CoarseDropout,
"GammaCorrection": GammaCorrection, "GammaCorrection": GammaCorrection,
"PlanckianJitter": PlanckianJitter, "PlanckianJitter": PlanckianJitter,
} }
)
_register_custom_transforms()
@dataclass @dataclass
+91 -1
View File
@@ -475,7 +475,7 @@ ROBOTICS_TRANSFORMS = [
("RandomShadow", RandomShadow, {"opacity": (0.3, 0.6)}), ("RandomShadow", RandomShadow, {"opacity": (0.3, 0.6)}),
("CoarseDropout", CoarseDropout, {"max_holes": 8}), ("CoarseDropout", CoarseDropout, {"max_holes": 8}),
("GammaCorrection", GammaCorrection, {"gamma": (0.5, 2.0)}), ("GammaCorrection", GammaCorrection, {"gamma": (0.5, 2.0)}),
("PlanckianJitter", PlanckianJitter, {"strength": (0.85, 1.15)}), ("PlanckianJitter", PlanckianJitter, {"temperature": (3_000, 15_000)}),
] ]
@@ -523,3 +523,93 @@ def test_make_transform_error_message_includes_custom():
"""Error message should list all registered custom transforms.""" """Error message should list all registered custom transforms."""
with pytest.raises(ValueError, match="GaussianNoise"): with pytest.raises(ValueError, match="GaussianNoise"):
make_transform_from_config(ImageTransformConfig(type="NonExistent")) make_transform_from_config(ImageTransformConfig(type="NonExistent"))
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
@pytest.mark.parametrize("shape", [(4, 3, 32, 32), (2, 4, 3, 16, 16)])
def test_robotics_transform_supports_temporal_batches(name, cls, kwargs, shape):
img = torch.rand(shape)
out = cls(**kwargs)(img)
assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
assert out.min() >= 0
assert out.max() <= 1
@pytest.mark.parametrize(
"cls,kwargs",
[
(GaussianNoise, {"std": (25.0, 25.0)}),
(MotionBlur, {"kernel_size": 5}),
(JPEGCompression, {"quality": 10}),
(
GaussianPatchBrightness,
{"num_patches": 1, "sigma_range": (0.2, 0.2), "factor_range": (0.5, 0.5)},
),
(RandomShadow, {"opacity": 0.5}),
(CoarseDropout, {"max_holes": 1, "fill_value": 0.0}),
(GammaCorrection, {"gamma": (2.0, 2.0)}),
(PlanckianJitter, {"temperature": 3_000}),
],
)
def test_robotics_transform_is_not_silent_noop(cls, kwargs):
img = torch.rand(3, 32, 32)
out = cls(**kwargs)(img)
assert not torch.equal(out, img)
@pytest.mark.parametrize(
"transform",
[
GaussianNoise(std=25),
RandomShadow(opacity=0.5),
CoarseDropout(max_holes=4),
],
)
def test_robotics_transform_random_params_are_reused(transform):
img = torch.rand(3, 32, 32)
params = transform.make_params([img])
torch.testing.assert_close(transform.transform(img, params), transform.transform(img, params))
def test_motion_blur_kernel_size_stays_in_configured_range():
transform = MotionBlur(kernel_size=(4, 10))
sampled_sizes = {transform.make_params([])["kernel_size"] for _ in range(100)}
assert sampled_sizes <= {5, 7, 9}
assert sampled_sizes
def test_gamma_correction_scalar_below_one_defines_symmetric_range():
transform = GammaCorrection(gamma=0.5)
assert transform.gamma == (0.5, 2.0)
assert transform(torch.rand(3, 8, 8)).shape == (3, 8, 8)
def test_planckian_jitter_uses_correlated_temperature_coefficients():
img = torch.full((2, 3, 8, 8), 0.25)
out = PlanckianJitter(temperature=3_000)(img)
torch.testing.assert_close(out[:, 1], img[:, 1])
assert torch.all(out[:, 0] > out[:, 1])
assert torch.all(out[:, 2] < out[:, 1])
def test_random_shadow_supports_small_images():
img = torch.rand(3, 7, 7)
assert RandomShadow()(img).shape == img.shape
@pytest.mark.parametrize(
"cls,kwargs",
[
(GaussianNoise, {"std": (-1.0, 1.0)}),
(MotionBlur, {"kernel_size": 4}),
(JPEGCompression, {"quality": (0, 75)}),
(GaussianPatchBrightness, {"sigma_range": (0.0, 0.25)}),
(RandomShadow, {"opacity": (0.3, 1.1)}),
(CoarseDropout, {"max_holes": 0}),
(GammaCorrection, {"gamma": 0.0}),
(PlanckianJitter, {"temperature": (2_000, 6_500)}),
],
)
def test_robotics_transform_rejects_invalid_config(cls, kwargs):
with pytest.raises(ValueError):
cls(**kwargs)