| """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: |
| 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 = 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} |
|
|
|
|
|
|
