Blame for scripts/generate_multiple_degradation_extreme_gallery.py

#!/usr/bin/env python
"""Generate WORST-CASE samples: many degradations stacked at high strength on one template.

Like scripts/generate_multiple_degradation_gallery.py, but every augmentation
family's probability and strength range are pushed high (instead of using
configs/generate.yaml's mild "realistic batch" defaults), and placement/
rendering are also pushed toward their hardest settings (heavy misplacement,
small/faint signature). Most effects fire on most samples, so each sample
ends up with a large, clearly-visible stack of simultaneous degradations.

The caption band at the top of each image lists not just which effects
fired but their resolved *extent*: strength, plus any effect-specific
parameter (blob count, shadow type, JPEG quality, wrinkle count/amplitude,
etc.) alongside the resolved placement/rendering severity for that sample.

Usage:
    python scripts/generate_multiple_degradation_extreme_gallery.py
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from PIL import Image, ImageDraw, ImageFont

from src.config.loader import load_config
from src.generator import GeneratorContext
from src.visualization.visualizer import save_visualization
from scripts.generate_degradation_gallery import ALL_EFFECTS

ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "multiple_degradation_extreme"
N_SAMPLES = 10
TEMPLATE_ID = "bank_doc_05"

EFFECT_PROBABILITY = 0.85   # near-certain per effect, but still leaves sample-to-sample variety
STRENGTH_MIN, STRENGTH_MAX = 0.7, 1.0


def _load_font(size: int):
    for candidate in (r"C:\Windows\Fonts\consola.ttf", r"C:\Windows\Fonts\arial.ttf"):
        if Path(candidate).exists():
            return ImageFont.truetype(candidate, size)
    return ImageFont.load_default()


def build_overrides() -> list[str]:
    overrides = [
        f"document.template={TEMPLATE_ID}",
        "generation.run_validation=false",
        # Worst-case placement: large rotation, small scale, frequent/strong misplacement.
        "placement.rotation_deg.min=-15", "placement.rotation_deg.max=15",
        "placement.scale_factor.min=0.3", "placement.scale_factor.max=0.6",
        "placement.margin=0.03",
        "placement.misplacement_probability=0.5",
        "placement.misplacement_strength.min=0.5", "placement.misplacement_strength.max=1.0",
        "placement.allow_partial_outside=true",
        # Worst-case rendering: faint, blurred, heavily feathered ink.
        "rendering.ink_opacity.min=0.3", "rendering.ink_opacity.max=0.55",
        "rendering.blur_sigma.min=0.6", "rendering.blur_sigma.max=1.4",
        "rendering.feather_px.min=1.5", "rendering.feather_px.max=3.0",
        "rendering.scanner_noise=0.05",
    ]
    for name in ALL_EFFECTS:
        overrides.append(f"augmentation.{name}.enabled=true")
        overrides.append(f"augmentation.{name}.probability={EFFECT_PROBABILITY}")
        if name == "jpeg_compression":
            overrides += [
                "augmentation.jpeg_compression.strength=0.9",
                "augmentation.jpeg_compression.min_quality=10",
                "augmentation.jpeg_compression.max_quality=30",
            ]
        else:
            overrides.append(f"augmentation.{name}.strength.min={STRENGTH_MIN}")
            overrides.append(f"augmentation.{name}.strength.max={STRENGTH_MAX}")
    return overrides


def _fmt_effect(name: str, rec: dict) -> str:
    params = rec.get("params") or {}
    if name == "jpeg_compression" and "quality" in params:
        return f"{name}=q{params['quality']}"
    tag = f"{name}={rec['strength']:.2f}"
    if "resolved_type" in params:
        tag += f"({params['resolved_type']})"
    elif "resolved_tone" in params:
        tag += f"({params['resolved_tone']})"
    elif name == "crumple" and "wrinkle_lines" in params:
        tag += f"({params['wrinkle_lines']}L,{params['displacement_amplitude_px']:.1f}px)"
    elif name == "coffee_stain" and "blob_count" in params:
        tag += f"({params['blob_count']}blobs)"
    elif name == "page_tilt" and "tilt_deg" in params:
        tag += f"({params['tilt_deg']:.1f}deg)"
    elif name == "dust" and "speck_count" in params:
        tag += f"({params['speck_count']}px)"
    elif name == "scanner_streaks" and "streak_count" in params:
        tag += f"({params['streak_count']}streaks)"
    return tag


def _annotate(path: Path, metadata: dict) -> list[str]:
    effects = metadata["augmentation"]["effects"]
    applied_names = [n for n, r in effects.items() if r["applied"]]
    applied_tags = [_fmt_effect(n, effects[n]) for n in applied_names]

    placement, rendering, difficulty = metadata["placement"], metadata["rendering"], metadata["difficulty"]
    header = (
        f"{metadata['document_id']}  |  difficulty={difficulty['score']:.3f} ({difficulty['tier']})  |  "
        f"placement: rot={placement['rotation_deg']:.1f}deg scale={placement['scale_factor']:.2f} "
        f"misplaced={placement['misplaced']}  |  "
        f"render: ink_opacity={rendering['ink_opacity']:.2f} blur_sigma={rendering['blur_sigma']:.2f}"
    )
    body = f"Applied ({len(applied_tags)}/{len(effects)}): " + (", ".join(applied_tags) if applied_tags else "none")

    img = Image.open(path).convert("RGB")
    font = _load_font(15)
    draw_probe = ImageDraw.Draw(img)
    max_w = img.width - 28

    # Pixel-accurate wrap of the (potentially long) applied-effects line.
    words = body.replace(", ", " ,SPLIT_ME ").split(" ,SPLIT_ME ")
    lines, current = [], ""
    for tok in words:
        candidate = (current + ", " + tok) if current else tok
        if draw_probe.textlength(candidate, font=font) <= max_w:
            current = candidate
        else:
            if current:
                lines.append(current)
            current = tok
    if current:
        lines.append(current)

    all_lines = [header] + lines
    line_h = 22
    band_h = 14 + line_h * len(all_lines)

    canvas = Image.new("RGB", (img.width, img.height + band_h), "white")
    canvas.paste(img, (0, band_h))
    draw = ImageDraw.Draw(canvas)
    draw.rectangle([0, 0, canvas.width, band_h], fill=(60, 18, 18))
    y = 8
    for i, line in enumerate(all_lines):
        draw.text((14, y), line, fill="white" if i == 0 else (255, 214, 194), font=font)
        y += line_h
    canvas.save(path)
    return applied_names


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    config = load_config(overrides=build_overrides())
    ctx = GeneratorContext(config)

    for i in range(1, N_SAMPLES + 1):
        metadata, final_image, _ = ctx.run_sample(i)
        out_path = OUT_DIR / f"sample_{i}.png"
        save_visualization(out_path, final_image, metadata)
        applied = _annotate(out_path, metadata)
        print(f"[sample_{i}] {len(applied)}/{len(metadata['augmentation']['effects'])} effects -> {', '.join(applied)}")

    print(f"\nWrote {N_SAMPLES} extreme samples -> {OUT_DIR}")


if __name__ == "__main__":
    main()