Blame for src/augmentation/document_effects.py

"""Custom augmentation effects operating on the full, already-blended document image.

These simulate physical document degradation and scanning artifacts that
Albumentations has no domain-specific transform for: coffee/water/paper
stains, directional shadows, page creasing, discoloration, dust, and
scanner streaks. Each function returns ``(image, actual_params)`` so the
exact resolved parameters (e.g. which corner a shadow fell in) can be
recorded in metadata.
"""

from __future__ import annotations

import cv2
import numpy as np

from src.utils.transform_utils import smooth_noise_field


def _blob_mask(shape: tuple[int, int], center: tuple[float, float], radii: tuple[float, float],
                angle_deg: float, blur_sigma: float) -> np.ndarray:
    """Build a soft-edged elliptical blob mask in ``[0, 1]``."""
    h, w = shape
    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.ellipse(mask, (int(center[0]), int(center[1])), (int(radii[0]), int(radii[1])), angle_deg, 0, 360, 255, -1)
    mask_f = cv2.GaussianBlur(mask.astype(np.float32), (0, 0), sigmaX=blur_sigma) / 255.0
    return np.clip(mask_f, 0, 1)


def coffee_stain(image: np.ndarray, strength: float, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Overlay 1-3 translucent brown ring-edged blotches, like a coffee cup mark."""
    h, w = image.shape[:2]
    out = image.astype(np.float32)
    n_blobs = 1 + int(rng.integers(0, 1 + int(round(strength * 3))))
    centers = []
    for _ in range(max(n_blobs, 1)):
        cx, cy = rng.uniform(0.1, 0.9) * w, rng.uniform(0.1, 0.9) * h
        base_r = (0.03 + strength * 0.10) * min(w, h)
        rx, ry = base_r * rng.uniform(0.8, 1.3), base_r * rng.uniform(0.7, 1.2)
        angle = float(rng.uniform(0, 180))
        fill = _blob_mask((h, w), (cx, cy), (rx, ry), angle, blur_sigma=base_r * 0.35)
        rim = fill - _blob_mask((h, w), (cx, cy), (rx * 0.82, ry * 0.82), angle, blur_sigma=base_r * 0.2)
        color = np.array([139, 98, 60], dtype=np.float32)
        opacity = 0.10 + strength * 0.30
        for c in range(3):
            out[:, :, c] = out[:, :, c] * (1 - fill * opacity * 0.5) + color[c] * fill * opacity * 0.5
            out[:, :, c] = out[:, :, c] * (1 - rim * opacity * 0.9) + color[c] * 0.7 * rim * opacity * 0.9
        centers.append((round(float(cx), 1), round(float(cy), 1)))
    return np.clip(out, 0, 255).astype(np.uint8), {"blob_count": len(centers), "centers": centers}


def stain(image: np.ndarray, strength: float, stain_type: str, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Overlay a paper or water stain blotch.

    ``paper`` stains are mottled brownish-grey discolorations; ``water``
    stains are pale, soft-edged, faintly blue-grey tide marks.
    """
    h, w = image.shape[:2]
    out = image.astype(np.float32)
    resolved_type = stain_type if stain_type != "random" else rng.choice(["paper", "water"])
    cx, cy = rng.uniform(0.15, 0.85) * w, rng.uniform(0.15, 0.85) * h
    base_r = (0.06 + strength * 0.16) * min(w, h)

    if resolved_type == "water":
        color = np.array([170, 178, 182], dtype=np.float32)
        opacity = 0.06 + strength * 0.18
        rx, ry = base_r * rng.uniform(1.0, 1.6), base_r * rng.uniform(0.8, 1.3)
        fill = _blob_mask((h, w), (cx, cy), (rx, ry), float(rng.uniform(0, 180)), blur_sigma=base_r * 0.6)
        edge = fill - _blob_mask((h, w), (cx, cy), (rx * 0.9, ry * 0.9), 0, blur_sigma=base_r * 0.3)
        for c in range(3):
            out[:, :, c] = out[:, :, c] * (1 - fill * opacity * 0.4) + color[c] * fill * opacity * 0.4
            out[:, :, c] = out[:, :, c] * (1 - edge * opacity) + color[c] * edge * opacity
    else:
        color = np.array([120, 106, 84], dtype=np.float32)
        opacity = 0.08 + strength * 0.22
        for _ in range(3):
            rr = base_r * rng.uniform(0.4, 1.0)
            ox, oy = cx + rng.uniform(-base_r, base_r) * 0.5, cy + rng.uniform(-base_r, base_r) * 0.5
            fill = _blob_mask((h, w), (ox, oy), (rr, rr * rng.uniform(0.7, 1.1)), float(rng.uniform(0, 180)),
                               blur_sigma=rr * 0.5)
            for c in range(3):
                out[:, :, c] = out[:, :, c] * (1 - fill * opacity * 0.5) + color[c] * fill * opacity * 0.5

    return np.clip(out, 0, 255).astype(np.uint8), {
        "resolved_type": resolved_type, "center": (round(float(cx), 1), round(float(cy), 1)),
    }


def shadow(image: np.ndarray, strength: float, shadow_type: str, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Apply an edge, corner, or fold shadow gradient."""
    h, w = image.shape[:2]
    resolved_type = shadow_type if shadow_type != "random" else rng.choice(["edge", "corner", "fold"])
    out = image.astype(np.float32)
    darkness = 0.15 + strength * 0.55

    if resolved_type == "edge":
        side = rng.choice(["top", "bottom", "left", "right"])
        band_frac = 0.12 + strength * 0.18
        yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
        if side == "top":
            grad = np.clip(1 - yy / (h * band_frac), 0, 1)
        elif side == "bottom":
            grad = np.clip(1 - (h - 1 - yy) / (h * band_frac), 0, 1)
        elif side == "left":
            grad = np.clip(1 - xx / (w * band_frac), 0, 1)
        else:
            grad = np.clip(1 - (w - 1 - xx) / (w * band_frac), 0, 1)
        gain = 1 - grad * darkness
        out *= gain[:, :, None]
        extra = {"side": side}
    elif resolved_type == "corner":
        corner = rng.choice(["top_left", "top_right", "bottom_left", "bottom_right"])
        cx = 0 if "left" in corner else w
        cy = 0 if "top" in corner else h
        yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
        dist = np.hypot(xx - cx, yy - cy)
        radius = 0.55 * min(w, h) * (0.7 + strength * 0.6)
        grad = np.clip(1 - dist / radius, 0, 1)
        gain = 1 - grad * darkness
        out *= gain[:, :, None]
        extra = {"corner": corner}
    else:  # fold
        n_folds = 1 + int(strength * 2)
        extra_lines = []
        for _ in range(n_folds):
            vertical = bool(rng.integers(0, 2))
            pos = rng.uniform(0.25, 0.75)
            width = (0.015 + strength * 0.03) * (w if vertical else h)
            yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
            axis = xx if vertical else yy
            center = pos * (w if vertical else h)
            dist = np.abs(axis - center)
            grad = np.clip(1 - dist / max(width, 1.0), 0, 1)
            gain = 1 - grad * darkness * 0.7
            out *= gain[:, :, None]
            extra_lines.append({"vertical": vertical, "position_frac": round(float(pos), 3)})
        extra = {"folds": extra_lines}

    return np.clip(out, 0, 255).astype(np.uint8), {"resolved_type": resolved_type, **extra}


def crumple(image: np.ndarray, strength: float, wrinkle_density: float, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Simulate paper crumpling: subtle geometric displacement plus faint wrinkle-line shading."""
    h, w = image.shape[:2]
    field_x = smooth_noise_field((h, w), rng, blur_frac=0.06) - 0.5
    field_y = smooth_noise_field((h, w), rng, blur_frac=0.06) - 0.5
    amplitude = strength * min(w, h) * 0.015

    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
    map_x = xx + field_x * amplitude
    map_y = yy + field_y * amplitude
    warped = cv2.remap(image, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)

    shade_field = smooth_noise_field((h, w), rng, blur_frac=0.03)
    n_lines = int(3 + wrinkle_density * 12)
    line_mask = np.zeros((h, w), dtype=np.float32)
    for _ in range(n_lines):
        x1, y1 = rng.uniform(0, w), rng.uniform(0, h)
        angle = rng.uniform(0, 2 * np.pi)
        length = rng.uniform(0.15, 0.5) * min(w, h)
        x2, y2 = x1 + np.cos(angle) * length, y1 + np.sin(angle) * length
        cv2.line(line_mask, (int(x1), int(y1)), (int(x2), int(y2)), 1.0, thickness=1)
    line_mask = cv2.GaussianBlur(line_mask, (0, 0), sigmaX=1.0)

    shading = (shade_field - 0.5) * strength * 0.25 - line_mask * strength * 0.18
    out = warped.astype(np.float32) * (1 + shading[:, :, None])
    return np.clip(out, 0, 255).astype(np.uint8), {"wrinkle_lines": n_lines, "displacement_amplitude_px": round(float(amplitude), 3)}


_SEPIA_MATRIX = np.array([
    [0.393, 0.769, 0.189],
    [0.349, 0.686, 0.168],
    [0.272, 0.534, 0.131],
], dtype=np.float32)


def discoloration(image: np.ndarray, strength: float, tone: str, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Apply an overall color-cast discoloration: yellowing, greying, or sepia toning."""
    resolved_tone = tone if tone != "random" else rng.choice(["yellow", "grey", "sepia"])
    img_f = image.astype(np.float32)

    if resolved_tone == "yellow":
        tint = np.array([1.0 + strength * 0.12, 1.0 + strength * 0.04, 1.0 - strength * 0.28], dtype=np.float32)
        out = img_f * tint[None, None, :]
    elif resolved_tone == "grey":
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY).astype(np.float32)
        gray3 = np.repeat(gray[:, :, None], 3, axis=2)
        out = img_f * (1 - strength * 0.6) + gray3 * (strength * 0.6)
    else:  # sepia
        sepia = img_f @ _SEPIA_MATRIX.T
        out = img_f * (1 - strength * 0.7) + sepia * (strength * 0.7)

    unevenness = smooth_noise_field(image.shape[:2], rng, blur_frac=0.05)
    gain = 1 - (1 - unevenness) * strength * 0.08
    out *= gain[:, :, None]

    return np.clip(out, 0, 255).astype(np.uint8), {"resolved_tone": resolved_tone}


def page_tilt(image: np.ndarray, strength: float, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Rotate the whole page by a small angle, simulating a crooked scanner feed."""
    h, w = image.shape[:2]
    angle = float(rng.uniform(-1.0, 1.0) * (1.5 + strength * 6.5))
    matrix = cv2.getRotationMatrix2D((w / 2.0, h / 2.0), angle, 1.0)
    rotated = cv2.warpAffine(image, matrix, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    return rotated, {"tilt_deg": round(angle, 3)}


def dust(image: np.ndarray, strength: float, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Scatter small light/dark specks simulating dust and sensor debris."""
    h, w = image.shape[:2]
    out = image.astype(np.float32)
    n_specks = int((h * w) * 0.00003 * (1 + strength * 6))
    xs = rng.integers(0, w, size=n_specks)
    ys = rng.integers(0, h, size=n_specks)
    for x, y in zip(xs, ys):
        r = int(rng.integers(1, 3))
        val = float(rng.choice([255.0, 40.0]))
        cv2.circle(out, (int(x), int(y)), r, (val, val, val), -1)
    return np.clip(out, 0, 255).astype(np.uint8), {"speck_count": int(n_specks)}


def scanner_streaks(image: np.ndarray, strength: float, rng: np.random.Generator) -> tuple[np.ndarray, dict]:
    """Add faint vertical scan-line streaks, simulating a dirty scanner sensor/glass."""
    h, w = image.shape[:2]
    out = image.astype(np.float32)
    n_streaks = 1 + int(strength * 6)
    xs = sorted(int(v) for v in rng.integers(0, w, size=n_streaks))
    for x in xs:
        width = int(rng.integers(1, 3))
        delta = float(rng.uniform(-1, 1)) * (10 + strength * 35)
        x1, x2 = max(x - width, 0), min(x + width, w)
        out[:, x1:x2, :] += delta
    return np.clip(out, 0, 255).astype(np.uint8), {"streak_count": n_streaks, "positions_x": xs}