"""Orchestrates every configurable augmentation effect for one sample. Two passes are performed: 1. :meth:`AugmentationEngine.apply_to_signature` -- effects that only make sense applied to the ink itself (stroke geometry, ink density, localized blur/jitter), run on the RGBA signature layer *before* blending. 2. :meth:`AugmentationEngine.apply_to_document` -- effects that simulate physical document/scanner degradation, run on the fully blended page. Every effect is gated independently by its own ``enabled`` flag, a probability roll, and a strength value -- each resolved from a seed derived deterministically from ``(sample_seed, augmentation_name)`` (or an explicit per-effect ``seed`` override in config), guaranteeing reproducibility. Every effect's resolved parameters are collected into an :class:`AugmentationRecord` for full metadata traceability, even when the effect was not applied. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Any import numpy as np from src.augmentation import albumentations_effects as AE from src.augmentation import document_effects as DE from src.augmentation import signature_effects as SE from src.config.schema import AugmentationParam, AugmentationsConfig from src.utils.logging_utils import get_logger from src.utils.random_utils import SeedManager, resolve_scalar_or_range logger = get_logger(__name__) @dataclass class AugmentationRecord: """The resolved, applied-or-not state of a single augmentation effect.""" name: str enabled: bool applied: bool strength: float probability: float seed: int params: dict[str, Any] = field(default_factory=dict) def _prepare(name: str, cfg: AugmentationParam, seed_manager: SeedManager) -> tuple[bool, np.random.Generator, int]: seed = cfg.seed if cfg.seed is not None else seed_manager.seed_for(f"augmentation:{name}") rng = np.random.default_rng(seed) roll = float(rng.uniform(0.0, 1.0)) applied = bool(cfg.enabled and roll < cfg.probability) return applied, rng, seed def _record(name: str, cfg: AugmentationParam, applied: bool, seed: int, strength: float = 0.0, params: "dict | None" = None) -> AugmentationRecord: return AugmentationRecord( name=name, enabled=cfg.enabled, applied=applied, strength=round(strength, 4), probability=cfg.probability, seed=seed, params=params or {}, ) class AugmentationEngine: """Applies the full configurable augmentation suite to a signature layer and a document. Example: >>> engine = AugmentationEngine(cfg.augmentation) >>> rgba, sig_records = engine.apply_to_signature(placed_rgba, seed_manager) >>> doc, doc_records = engine.apply_to_document(blended_doc, seed_manager) """ def __init__(self, config: AugmentationsConfig) -> None: self.config = config # ------------------------------------------------------------------ # # Signature-layer effects # ------------------------------------------------------------------ # def apply_to_signature( self, rgba: np.ndarray, seed_manager: SeedManager, ) -> tuple[np.ndarray, dict[str, AugmentationRecord]]: cfg = self.config records: dict[str, AugmentationRecord] = {} out = rgba applied, rng, seed = _prepare("rotation", cfg.rotation, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.rotation.strength, rng) extra_deg = float(rng.uniform(-1, 1)) * strength * 8.0 out = SE.jitter_rotation(out, extra_deg) records["rotation"] = _record("rotation", cfg.rotation, True, seed, strength, {"extra_deg": round(extra_deg, 3)}) else: records["rotation"] = _record("rotation", cfg.rotation, False, seed) applied, rng, seed = _prepare("stroke_dilation", cfg.stroke_dilation, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.stroke_dilation.strength, rng) out = SE.stroke_morphology(out, strength, thicken=True) records["stroke_dilation"] = _record("stroke_dilation", cfg.stroke_dilation, True, seed, strength, {"thicken": True}) else: records["stroke_dilation"] = _record("stroke_dilation", cfg.stroke_dilation, False, seed) applied, rng, seed = _prepare("ink_fading", cfg.ink_fading, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.ink_fading.strength, rng) out = SE.ink_fading(out, strength, rng) records["ink_fading"] = _record("ink_fading", cfg.ink_fading, True, seed, strength) else: records["ink_fading"] = _record("ink_fading", cfg.ink_fading, False, seed) applied, rng, seed = _prepare("ink_bleed", cfg.ink_bleed, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.ink_bleed.strength, rng) out = SE.ink_bleed(out, strength) records["ink_bleed"] = _record("ink_bleed", cfg.ink_bleed, True, seed, strength) else: records["ink_bleed"] = _record("ink_bleed", cfg.ink_bleed, False, seed) applied, rng, seed = _prepare("motion_blur", cfg.motion_blur, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.motion_blur.strength, rng) angle = float(rng.uniform(0, 180)) out = SE.motion_blur_signature(out, strength, angle) records["motion_blur"] = _record("motion_blur", cfg.motion_blur, True, seed, strength, {"angle_deg": round(angle, 2)}) else: records["motion_blur"] = _record("motion_blur", cfg.motion_blur, False, seed) applied, rng, seed = _prepare("gaussian_blur", cfg.gaussian_blur, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.gaussian_blur.strength, rng) sigma = 0.2 + strength * 2.2 out = SE.gaussian_blur_signature(out, sigma) records["gaussian_blur"] = _record("gaussian_blur", cfg.gaussian_blur, True, seed, strength, {"sigma": round(sigma, 3)}) else: records["gaussian_blur"] = _record("gaussian_blur", cfg.gaussian_blur, False, seed) return out, records # ------------------------------------------------------------------ # # Document-level effects # ------------------------------------------------------------------ # def apply_to_document( self, image: np.ndarray, seed_manager: SeedManager, ) -> tuple[np.ndarray, dict[str, AugmentationRecord]]: cfg = self.config records: dict[str, AugmentationRecord] = {} out = image applied, rng, seed = _prepare("perspective", cfg.perspective, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.perspective.strength, rng) out = AE.perspective_warp(out, strength, seed) records["perspective"] = _record("perspective", cfg.perspective, True, seed, strength) else: records["perspective"] = _record("perspective", cfg.perspective, False, seed) applied, rng, seed = _prepare("contrast", cfg.contrast, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.contrast.strength, rng) out = AE.contrast_shift(out, strength, seed) records["contrast"] = _record("contrast", cfg.contrast, True, seed, strength) else: records["contrast"] = _record("contrast", cfg.contrast, False, seed) applied, rng, seed = _prepare("gamma", cfg.gamma, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.gamma.strength, rng) out, gamma_value = AE.gamma_shift(out, strength, seed) records["gamma"] = _record("gamma", cfg.gamma, True, seed, strength, {"gamma": gamma_value}) else: records["gamma"] = _record("gamma", cfg.gamma, False, seed) applied, rng, seed = _prepare("coffee_stain", cfg.coffee_stain, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.coffee_stain.strength, rng) out, params = DE.coffee_stain(out, strength, rng) records["coffee_stain"] = _record("coffee_stain", cfg.coffee_stain, True, seed, strength, params) else: records["coffee_stain"] = _record("coffee_stain", cfg.coffee_stain, False, seed) applied, rng, seed = _prepare("stain", cfg.stain, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.stain.strength, rng) out, params = DE.stain(out, strength, cfg.stain.type, rng) records["stain"] = _record("stain", cfg.stain, True, seed, strength, params) else: records["stain"] = _record("stain", cfg.stain, False, seed) applied, rng, seed = _prepare("shadow", cfg.shadow, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.shadow.strength, rng) out, params = DE.shadow(out, strength, cfg.shadow.type, rng) records["shadow"] = _record("shadow", cfg.shadow, True, seed, strength, params) else: records["shadow"] = _record("shadow", cfg.shadow, False, seed) applied, rng, seed = _prepare("crumple", cfg.crumple, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.crumple.strength, rng) density = resolve_scalar_or_range(cfg.crumple.wrinkle_density, rng) out, params = DE.crumple(out, strength, density, rng) records["crumple"] = _record("crumple", cfg.crumple, True, seed, strength, {**params, "wrinkle_density": round(density, 3)}) else: records["crumple"] = _record("crumple", cfg.crumple, False, seed) applied, rng, seed = _prepare("discoloration", cfg.discoloration, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.discoloration.strength, rng) out, params = DE.discoloration(out, strength, cfg.discoloration.tone, rng) records["discoloration"] = _record("discoloration", cfg.discoloration, True, seed, strength, params) else: records["discoloration"] = _record("discoloration", cfg.discoloration, False, seed) applied, rng, seed = _prepare("page_tilt", cfg.page_tilt, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.page_tilt.strength, rng) out, params = DE.page_tilt(out, strength, rng) records["page_tilt"] = _record("page_tilt", cfg.page_tilt, True, seed, strength, params) else: records["page_tilt"] = _record("page_tilt", cfg.page_tilt, False, seed) applied, rng, seed = _prepare("dust", cfg.dust, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.dust.strength, rng) out, params = DE.dust(out, strength, rng) records["dust"] = _record("dust", cfg.dust, True, seed, strength, params) else: records["dust"] = _record("dust", cfg.dust, False, seed) applied, rng, seed = _prepare("scanner_streaks", cfg.scanner_streaks, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.scanner_streaks.strength, rng) out, params = DE.scanner_streaks(out, strength, rng) records["scanner_streaks"] = _record("scanner_streaks", cfg.scanner_streaks, True, seed, strength, params) else: records["scanner_streaks"] = _record("scanner_streaks", cfg.scanner_streaks, False, seed) applied, rng, seed = _prepare("scanner_noise", cfg.scanner_noise, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.scanner_noise.strength, rng) out = AE.sensor_noise(out, strength, seed) records["scanner_noise"] = _record("scanner_noise", cfg.scanner_noise, True, seed, strength) else: records["scanner_noise"] = _record("scanner_noise", cfg.scanner_noise, False, seed) applied, rng, seed = _prepare("salt_pepper", cfg.salt_pepper, seed_manager) if applied: strength = resolve_scalar_or_range(cfg.salt_pepper.strength, rng) out = AE.salt_and_pepper(out, strength, seed) records["salt_pepper"] = _record("salt_pepper", cfg.salt_pepper, True, seed, strength) else: records["salt_pepper"] = _record("salt_pepper", cfg.salt_pepper, False, seed) # JPEG compression must run last: it is a lossy re-encode and every # subsequent pixel-level effect would otherwise mask its artifacts. jpeg_cfg = cfg.jpeg_compression applied, rng, seed = _prepare("jpeg_compression", jpeg_cfg, seed_manager) if applied: strength = resolve_scalar_or_range(jpeg_cfg.strength, rng) out, quality = AE.jpeg_compression(out, jpeg_cfg.min_quality, jpeg_cfg.max_quality, seed) records["jpeg_compression"] = _record("jpeg_compression", jpeg_cfg, True, seed, strength, {"quality": quality}) else: records["jpeg_compression"] = _record("jpeg_compression", jpeg_cfg, False, seed) return out, records