Blame for scripts/degrade_cli.py

#!/usr/bin/env python
"""Interactive terminal wizard for generating synthetic signed documents with
hand-picked degradation effects.

A friendlier front end over main.py's --set override system: arrow-key menus
and checkboxes instead of memorizing config dot-paths. Same pipeline underneath
(GeneratorContext), so output is identical to what main.py would produce with
the equivalent --set flags.

Usage:
    python scripts/degrade_cli.py
"""

from __future__ import annotations

import json
import os
import sys
from datetime import datetime
from pathlib import Path

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

import questionary
from questionary import Choice, Separator
from rich.console import Console
from rich.panel import Panel
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from rich.table import Table

from src.config.loader import load_config
from src.generator import GeneratorContext
from src.templates.template_loader import TemplateLoader
from src.utils.image_io import save_image_rgb
from src.visualization.visualizer import save_visualization
from scripts.generate_before_after_assets import ALL_EFFECTS, DISPLAY_NAME, DOCUMENT_EFFECTS, SIGNATURE_EFFECTS
from scripts.generate_multiple_degradation_extreme_gallery import build_overrides as extreme_overrides

ROOT = Path(__file__).resolve().parents[1]
console = Console()

# ---- palette, matched to Metrics_Thresholds_Formulas.pptx -----------------
NAVY, TEAL, ORANGE, GRAY = "#0F2A43", "#1B7A7A", "#E07A2C", "#6B767C"

QSTYLE = questionary.Style([
    ("qmark", f"fg:{ORANGE} bold"),
    ("question", "bold"),
    ("answer", f"fg:{TEAL} bold"),
    ("pointer", f"fg:{ORANGE} bold"),
    ("highlighted", f"fg:{TEAL} bold"),
    ("selected", f"fg:{TEAL}"),
    ("instruction", f"fg:{GRAY}"),
    ("separator", f"fg:{GRAY}"),
])


def banner() -> None:
    console.print(Panel.fit(
        "[bold white]Synthetic Signed Document Generator[/bold white]\n"
        f"[{TEAL}]pick degradations  ->  generate  ->  inspect[/{TEAL}]",
        border_style=TEAL, padding=(1, 4),
    ))


def discover_templates() -> list[tuple[str, str]]:
    cfg = load_config()
    loader = TemplateLoader(cfg.paths.image_dir, cfg.paths.annotations_dir)
    loader.discover()
    return sorted((tid, t.annotation.document_type) for tid, t in loader._templates.items())


def valid_float01(v: str) -> "bool | str":
    try:
        f = float(v)
    except ValueError:
        return "Enter a number between 0.0 and 1.0"
    return True if 0.0 <= f <= 1.0 else "Must be between 0.0 and 1.0"


def valid_posint(v: str) -> "bool | str":
    return True if v.isdigit() and int(v) > 0 else "Enter a positive whole number"


# ------------------------------------------------------------- questions ---
def ask_mode() -> str:
    return questionary.select(
        "How should degradations be applied?",
        choices=[
            Choice(title="Pick effects - isolated  (ONLY the chosen effects fire, nothing else)", value="isolate"),
            Choice(title="Pick effects - mixed     (chosen effects forced on, others may also roll)", value="mixed"),
            Choice(title="Preset - normal severity  (realistic random mix, configs/generate.yaml)", value="normal"),
            Choice(title="Preset - extreme severity (worst case, everything cranked)", value="extreme"),
        ],
        style=QSTYLE,
    ).ask()


def ask_template() -> "str | None":
    templates = discover_templates()
    choices = [Choice(title="Random each sample", value=None)]
    choices += [Choice(title=f"{tid}  ({doc_type})", value=tid) for tid, doc_type in templates]
    return questionary.select("Which document template?", choices=choices, style=QSTYLE).ask()


def ask_effects() -> list[str]:
    choices = [Separator("-- signature-level --")]
    choices += [Choice(title=DISPLAY_NAME[e], value=e) for e in SIGNATURE_EFFECTS]
    choices += [Separator("-- document-level --")]
    choices += [Choice(title=DISPLAY_NAME[e], value=e) for e in DOCUMENT_EFFECTS]
    selected = questionary.checkbox(
        "Pick degradations (space to toggle, enter to confirm)", choices=choices, style=QSTYLE,
    ).ask()
    if selected is None:
        raise KeyboardInterrupt
    if not selected:
        console.print(f"[{ORANGE}]Pick at least one effect.[/{ORANGE}]")
        return ask_effects()
    return selected


def ask_strengths(effects: list[str]) -> dict[str, float]:
    global_raw = questionary.text(
        "Strength for all selected effects (0.0-1.0)", default="0.75",
        validate=valid_float01, style=QSTYLE,
    ).ask()
    global_strength = float(global_raw)
    strengths = {e: global_strength for e in effects}
    if questionary.confirm("Customize strength per effect?", default=False, style=QSTYLE).ask():
        for e in effects:
            raw = questionary.text(
                f"  {DISPLAY_NAME[e]} strength", default=str(global_strength),
                validate=valid_float01, style=QSTYLE,
            ).ask()
            strengths[e] = float(raw)
    return strengths


def ask_count() -> int:
    raw = questionary.text("How many samples?", default="3", validate=valid_posint, style=QSTYLE).ask()
    return int(raw)


def ask_save_viz() -> bool:
    return questionary.confirm(
        "Also save the debug-panel visualization (crop / mask / metrics / difficulty)?",
        default=True, style=QSTYLE,
    ).ask()


def ask_output_dir() -> Path:
    default = f"output/interactive_{datetime.now():%Y%m%d_%H%M%S}"
    raw = questionary.text("Output folder", default=default, style=QSTYLE).ask()
    return ROOT / raw


# ------------------------------------------------------------ overrides ----
def build_pick_overrides(mode: str, template: "str | None", effects: list[str], strengths: dict[str, float]) -> list[str]:
    overrides = ["generation.run_validation=false"]
    if template:
        overrides.append(f"document.template={template}")

    if mode == "isolate":
        for name in ALL_EFFECTS:
            on = name in effects
            overrides.append(f"augmentation.{name}.enabled={'true' if on else 'false'}")
            overrides.append(f"augmentation.{name}.probability={'1.0' if on else '0.0'}")
        for name in effects:
            overrides.append(f"augmentation.{name}.strength={strengths[name]}")
    else:  # mixed
        for name in effects:
            overrides.append(f"augmentation.{name}.enabled=true")
            overrides.append(f"augmentation.{name}.probability=1.0")
            overrides.append(f"augmentation.{name}.strength={strengths[name]}")
    return overrides


# ---------------------------------------------------------------- report ---
def show_plan(mode: str, template, effects, strengths, count: int, out_dir: Path) -> None:
    table = Table(title="Plan", title_style=f"bold {NAVY}", show_header=False, border_style=GRAY)
    table.add_row("Mode", mode)
    table.add_row("Template", template or "random per sample")
    if effects:
        table.add_row("Effects", ", ".join(f"{DISPLAY_NAME[e]} ({strengths[e]:.2f})" for e in effects))
    table.add_row("Samples", str(count))
    table.add_row("Output", str(out_dir))
    console.print(table)


def show_results(results: list[tuple[str, float, str, list[str]]], out_dir: Path) -> None:
    table = Table(title="Results", border_style=GRAY, header_style=f"bold white on {NAVY}")
    table.add_column("Sample")
    table.add_column("Difficulty")
    table.add_column("Tier")
    table.add_column("Effects applied")
    tier_color = {"easy": "green", "medium": ORANGE, "hard": "red"}
    for sid, score, tier, applied in results:
        table.add_row(
            sid, f"{score:.3f}", f"[{tier_color.get(tier, 'white')}]{tier}[/]",
            ", ".join(applied) if applied else "[dim]none[/dim]",
        )
    console.print(table)
    console.print(f"\n[{TEAL}]Wrote {len(results)} sample(s) -> {out_dir}[/{TEAL}]")


# ------------------------------------------------------------------ main ---
def main() -> int:
    banner()
    try:
        mode = ask_mode()
        if mode is None:
            raise KeyboardInterrupt

        template, effects, strengths = None, [], {}
        if mode in ("isolate", "mixed"):
            template = ask_template()
            effects = ask_effects()
            strengths = ask_strengths(effects)

        count = ask_count()
        save_viz = ask_save_viz()
        out_dir = ask_output_dir()

        show_plan(mode, template, effects, strengths, count, out_dir)
        if not questionary.confirm("Generate now?", default=True, style=QSTYLE).ask():
            console.print("Cancelled.")
            return 0

        if mode == "normal":
            config = load_config(
                config_path=ROOT / "configs" / "generate.yaml",
                overrides=["document.template=bank_doc_05", "generation.run_validation=false"],
            )
        elif mode == "extreme":
            config = load_config(overrides=extreme_overrides())
        else:
            config = load_config(overrides=build_pick_overrides(mode, template, effects, strengths))

        out_dir.mkdir(parents=True, exist_ok=True)
        ctx = GeneratorContext(config)
        prefix, digits = config.generation.filename_prefix, config.generation.filename_digits

        results = []
        with Progress(
            SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
            BarColumn(complete_style=TEAL), TextColumn("{task.completed}/{task.total}"),
            TimeElapsedColumn(), console=console,
        ) as progress:
            task = progress.add_task("Generating", total=count)
            for i in range(1, count + 1):
                metadata, final_image, _ = ctx.run_sample(i)
                sample_id = f"{prefix}_{i:0{digits}d}"

                save_image_rgb(out_dir / f"{sample_id}.png", final_image)
                ctx.metadata_exporter.export(metadata, out_dir / f"{sample_id}.json")
                if save_viz:
                    save_visualization(out_dir / f"{sample_id}_viz.png", final_image, metadata)

                applied = [n for n, r in metadata["augmentation"]["effects"].items() if r["applied"]]
                results.append((sample_id, metadata["difficulty"]["score"], metadata["difficulty"]["tier"], applied))
                progress.advance(task)

        console.print()
        show_results(results, out_dir)

        if questionary.confirm("Open the output folder now?", default=True, style=QSTYLE).ask():
            if os.name == "nt":
                os.startfile(out_dir)  # noqa: S606
            else:
                console.print(str(out_dir))

        return 0

    except KeyboardInterrupt:
        console.print("\n[dim]Cancelled.[/dim]")
        return 1


if __name__ == "__main__":
    sys.exit(main())