| """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 |
|
|
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| 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_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 |
|
|
