143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
"""
|
|
Noise augmentation transforms.
|
|
|
|
Provides noise effects for document image augmentation:
|
|
- GaussianNoise: Adds Gaussian noise to simulate sensor noise
|
|
- SaltPepper: Adds salt and pepper noise for impulse noise effects
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from shared.augmentation.base import AugmentationResult, BaseAugmentation
|
|
|
|
|
|
class GaussianNoise(BaseAugmentation):
|
|
"""
|
|
Adds Gaussian noise to the image.
|
|
|
|
Simulates sensor noise from cameras or scanners.
|
|
Document-safe with conservative default parameters.
|
|
|
|
Parameters:
|
|
mean: Mean of the Gaussian noise (default: 0).
|
|
std: Standard deviation, can be int or (min, max) tuple (default: (5, 15)).
|
|
"""
|
|
|
|
name = "gaussian_noise"
|
|
affects_geometry = False
|
|
|
|
def _validate_params(self) -> None:
|
|
std = self.params.get("std", (5, 15))
|
|
if isinstance(std, (int, float)):
|
|
if std < 0:
|
|
raise ValueError("std must be non-negative")
|
|
elif isinstance(std, tuple):
|
|
if len(std) != 2 or std[0] < 0 or std[1] < std[0]:
|
|
raise ValueError("std tuple must be (min, max) with min <= max >= 0")
|
|
|
|
def apply(
|
|
self,
|
|
image: np.ndarray,
|
|
bboxes: np.ndarray | None = None,
|
|
rng: np.random.Generator | None = None,
|
|
) -> AugmentationResult:
|
|
rng = rng or np.random.default_rng()
|
|
|
|
mean = self.params.get("mean", 0)
|
|
std = self.params.get("std", (5, 15))
|
|
|
|
if isinstance(std, tuple):
|
|
std = rng.uniform(std[0], std[1])
|
|
|
|
# Generate noise
|
|
noise = rng.normal(mean, std, image.shape).astype(np.float32)
|
|
|
|
# Apply noise
|
|
noisy = image.astype(np.float32) + noise
|
|
noisy = np.clip(noisy, 0, 255).astype(np.uint8)
|
|
|
|
return AugmentationResult(
|
|
image=noisy,
|
|
bboxes=bboxes.copy() if bboxes is not None else None,
|
|
metadata={"applied_std": std},
|
|
)
|
|
|
|
def get_preview_params(self) -> dict[str, Any]:
|
|
return {"mean": 0, "std": 15}
|
|
|
|
|
|
class SaltPepper(BaseAugmentation):
|
|
"""
|
|
Adds salt and pepper (impulse) noise to the image.
|
|
|
|
Simulates defects from damaged sensors or transmission errors.
|
|
Very sparse by default to preserve document readability.
|
|
|
|
Parameters:
|
|
amount: Proportion of pixels to affect, can be float or (min, max) tuple.
|
|
Default: (0.001, 0.005) for very sparse noise.
|
|
salt_vs_pepper: Ratio of salt to pepper (default: 0.5 for equal amounts).
|
|
"""
|
|
|
|
name = "salt_pepper"
|
|
affects_geometry = False
|
|
|
|
def _validate_params(self) -> None:
|
|
amount = self.params.get("amount", (0.001, 0.005))
|
|
if isinstance(amount, (int, float)):
|
|
if not (0 <= amount <= 1):
|
|
raise ValueError("amount must be between 0 and 1")
|
|
elif isinstance(amount, tuple):
|
|
if len(amount) != 2 or not (0 <= amount[0] <= amount[1] <= 1):
|
|
raise ValueError("amount tuple must be (min, max) in range [0, 1]")
|
|
|
|
def apply(
|
|
self,
|
|
image: np.ndarray,
|
|
bboxes: np.ndarray | None = None,
|
|
rng: np.random.Generator | None = None,
|
|
) -> AugmentationResult:
|
|
rng = rng or np.random.default_rng()
|
|
|
|
amount = self.params.get("amount", (0.001, 0.005))
|
|
salt_vs_pepper = self.params.get("salt_vs_pepper", 0.5)
|
|
|
|
if isinstance(amount, tuple):
|
|
amount = rng.uniform(amount[0], amount[1])
|
|
|
|
# Copy image
|
|
output = image.copy()
|
|
h, w = image.shape[:2]
|
|
total_pixels = h * w
|
|
|
|
# Calculate number of salt and pepper pixels
|
|
num_salt = int(total_pixels * amount * salt_vs_pepper)
|
|
num_pepper = int(total_pixels * amount * (1 - salt_vs_pepper))
|
|
|
|
# Add salt (white pixels)
|
|
if num_salt > 0:
|
|
salt_coords = (
|
|
rng.integers(0, h, num_salt),
|
|
rng.integers(0, w, num_salt),
|
|
)
|
|
output[salt_coords] = 255
|
|
|
|
# Add pepper (black pixels)
|
|
if num_pepper > 0:
|
|
pepper_coords = (
|
|
rng.integers(0, h, num_pepper),
|
|
rng.integers(0, w, num_pepper),
|
|
)
|
|
output[pepper_coords] = 0
|
|
|
|
return AugmentationResult(
|
|
image=output,
|
|
bboxes=bboxes.copy() if bboxes is not None else None,
|
|
metadata={"applied_amount": amount},
|
|
)
|
|
|
|
def get_preview_params(self) -> dict[str, Any]:
|
|
return {"amount": 0.01, "salt_vs_pepper": 0.5}
|