build_demo_showcase.py

#!/usr/bin/env python3
"""Render stroke-stratified, lossless eight-method demo contact sheets.

The full benchmark data is intentionally excluded from Git. This utility
selects five source images shared by every method—one from each Unicode
stroke-count bin—then creates one annotated PNG per generator for the
repository README.

Run from the repository root:

    python scripts/build_demo_showcase.py \
        --data-root ../posion_attack
"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Final

from PIL import Image, ImageDraw, ImageFont


REPO_ROOT: Final = Path(__file__).resolve().parents[1]
METHODS: Final = (
    ("ASPL", "old", "aspl_eps0.05_steps200"),
    ("Glaze", "old", "mi_eps16_steps300"),
    ("AMP", "old", "attackvlm_eps8_steps300"),
    ("XTransfer", "old", "xtransfer_eps12_steps300"),
    ("AnyAttack", "new", None),
    ("Nightshade", "old", "nightshade_eps0.05_steps500"),
    ("MMCoA", "old", "mmcoa_eps1_steps100"),
    ("CoA", "new", None),
)
METHOD_STYLES: Final = {
    "ASPL": ("IMAGE-ONLY", "ε=0.05 · 200 steps", "#eaf3ff", "#2563eb"),
    "Glaze": ("IMAGE-ONLY", "ε=16/255 · 300 steps", "#eaf3ff", "#2563eb"),
    "AMP": ("IMAGE-ONLY", "ε=8/255 · 300 steps", "#eaf3ff", "#2563eb"),
    "XTransfer": ("IMAGE-ONLY", "ε=12/255 · 300 steps", "#eaf3ff", "#2563eb"),
    "AnyAttack": ("IMAGE-ONLY", "ε=16/255 · 1 step", "#eaf3ff", "#2563eb"),
    "Nightshade": ("TEXT-ONLY", "ε=0.05 · 500 steps", "#fff4e5", "#c2410c"),
    "MMCoA": ("IMAGE + TEXT", "ε=1/255 · 100 steps", "#f3e8ff", "#7e22ce"),
    "CoA": ("IMAGE + TEXT", "ε=8/255 · 100 steps", "#f3e8ff", "#7e22ce"),
}
STROKE_BINS: Final = (
    ("1–5", 1, 5, "白"),
    ("6–10", 6, 10, "伞"),
    ("11–15", 11, 15, "歪"),
    ("16–20", 16, 20, "壁"),
    ("21+", 21, 10_000, "囊"),
)
GENERATOR_CONFIG: Final = {
    "ID": {
        "label": "Illusion Diffusion ControlNet (ID)",
        "clean": "outputs/ID_2k/batch_test_20260306_142512",
        "old": "outputs/run_ID_full_20260310_223019/images",
        "new": "outputs/run_new_attacks_20260602_102643/images",
        "folders": {
            "AnyAttack": "anyattack_ID_eps16",
            "CoA": "coa_ID_eps8_steps100",
        },
        "output": "id_eight_methods.png",
    },
    "SDXL": {
        "label": "Stable Diffusion ControlNet (SD)",
        "clean": "outputs/sdxl_2k/batch_test_20260306_162734",
        "old": "outputs/run_SDXL_full_20260311_101341/images",
        "new": "outputs/run_new_attacks_20260602_152412/images",
        "folders": {
            "AnyAttack": "anyattack_SDXL_eps16",
            "CoA": "coa_SDXL_eps8_steps100",
        },
        "output": "sd_eight_methods.png",
    },
}
CELL_SIZE: Final = 280
LABEL_WIDTH: Final = 210
TITLE_HEIGHT: Final = 62
HEADER_HEIGHT: Final = 108
FOOTER_HEIGHT: Final = 44
ROW_GAP: Final = 5
BACKGROUND: Final = "#ffffff"
GRID: Final = "#d0d7de"
TEXT: Final = "#24292f"
MUTED_TEXT: Final = "#57606a"
TITLE_BACKGROUND: Final = "#0f172a"


@dataclass(frozen=True)
class Sample:
    """One source image selected for a stroke-complexity bin."""

    stem: str
    label: str
    strokes: int
    stroke_bin: str


def image_font(size: int) -> ImageFont.ImageFont:
    """Return a readable system font, with Pillow's default as fallback."""
    for candidate in (
        "/System/Library/Fonts/Supplemental/Arial Unicode.ttf",
        "/System/Library/Fonts/PingFang.ttc",
        "/System/Library/Fonts/Supplemental/Arial.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ):
        if Path(candidate).is_file():
            try:
                return ImageFont.truetype(candidate, size)
            except OSError:
                continue
    return ImageFont.load_default()


def attack_images(folder: Path) -> dict[str, Path]:
    images: dict[str, Path] = {}
    for image_path in folder.glob("*_adv.png"):
        stem = image_path.stem.removesuffix("_adv")
        if stem in images:
            raise ValueError(f"Duplicate attacked image stem: {stem}")
        images[stem] = image_path
    if not images:
        raise FileNotFoundError(f"No adversarial PNGs found in {folder}")
    return images


def load_strokes(path: Path) -> dict[str, int]:
    """Read Chinese-character stroke counts from the official Unihan data."""
    strokes: dict[str, int] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        fields = line.split("\t")
        if len(fields) == 3 and fields[1] == "kTotalStrokes":
            try:
                strokes[chr(int(fields[0][2:], 16))] = int(fields[2].split()[0])
            except (IndexError, ValueError):
                continue
    if not strokes:
        raise ValueError(f"No kTotalStrokes records found in {path}")
    return strokes


def source_label(path: Path) -> str:
    """Return the single-character label stored beside a clean source PNG."""
    payload = json.loads(path.with_suffix(".json").read_text(encoding="utf-8"))
    annotations = payload.get("annotations", [])
    if len(annotations) != 1 or not isinstance(annotations[0].get("label"), str):
        raise ValueError(f"Expected one string annotation label in {path.with_suffix('.json')}")
    return annotations[0]["label"]


def stroke_bin(count: int) -> str | None:
    for label, lower, upper, _ in STROKE_BINS:
        if lower <= count <= upper:
            return label
    return None


def select_samples(
    clean_images: dict[str, Path],
    method_images: dict[str, dict[str, Path]],
    strokes: dict[str, int],
) -> list[Sample]:
    """Select one shared source from every stroke-count bin."""
    shared = set(clean_images)
    for images in method_images.values():
        shared.intersection_update(images)
    if len(shared) < len(STROKE_BINS):
        raise ValueError(
            f"Expected at least {len(STROKE_BINS)} shared samples; found {len(shared)}"
        )
    candidates: list[Sample] = []
    for stem in sorted(shared):
        label = source_label(clean_images[stem])
        count = strokes.get(label)
        if count is None:
            continue
        sample_bin = stroke_bin(count)
        if sample_bin:
            candidates.append(Sample(stem, label, count, sample_bin))

    selected: list[Sample] = []
    for bin_label, _, _, preferred_label in STROKE_BINS:
        in_bin = [sample for sample in candidates if sample.stroke_bin == bin_label]
        preferred = [sample for sample in in_bin if sample.label == preferred_label]
        if not in_bin:
            raise ValueError(f"No shared source image found for stroke bin {bin_label}")
        selected.append((preferred or in_bin)[0])
    return selected


def center_text(
    draw: ImageDraw.ImageDraw,
    box: tuple[int, int, int, int],
    text: str,
    font: ImageFont.ImageFont,
    fill: str = TEXT,
) -> None:
    bbox = draw.multiline_textbbox((0, 0), text, font=font, align="center", spacing=3)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    x0, y0, x1, y1 = box
    draw.multiline_text(
        (x0 + (x1 - x0 - text_width) / 2, y0 + (y1 - y0 - text_height) / 2),
        text,
        fill=fill,
        font=font,
        align="center",
        spacing=3,
    )


def fit_square(source: Image.Image, size: int) -> Image.Image:
    """Force every tile to the same square size.

    Attack outputs are square but vary in native resolution (224 / 384 /
    512 / 1024). ``Image.thumbnail`` never enlarges, so small images left
    white borders. Resize both up and down so every cell fills ``size``.
    """
    source = source.convert("RGB")
    if source.size != (size, size):
        source = source.resize((size, size), Image.Resampling.LANCZOS)
    return source


def paste_image(canvas: Image.Image, path: Path, x: int, y: int) -> None:
    with Image.open(path) as source:
        tile = fit_square(source, CELL_SIZE)
    canvas.paste(tile, (x, y))


def render_generator(
    data_root: Path, generator: str, strokes: dict[str, int]
) -> dict[str, object]:
    config = GENERATOR_CONFIG[generator]
    clean_root = data_root / config["clean"]
    clean_images = {path.stem: path for path in clean_root.glob("*.png")}
    if not clean_images:
        raise FileNotFoundError(f"No clean PNGs found in {clean_root}")

    method_images: dict[str, dict[str, Path]] = {}
    for method, run_kind, default_folder in METHODS:
        folder_name = default_folder or config["folders"][method]
        attack_root = data_root / config[run_kind] / folder_name
        method_images[method] = attack_images(attack_root)

    selected = select_samples(clean_images, method_images, strokes)
    row_labels = ["Source (clean)"] + [method[0] for method in METHODS]
    width = LABEL_WIDTH + len(selected) * CELL_SIZE
    height = (
        TITLE_HEIGHT
        + HEADER_HEIGHT
        + len(row_labels) * (CELL_SIZE + ROW_GAP)
        + FOOTER_HEIGHT
    )
    canvas = Image.new("RGB", (width, height), BACKGROUND)
    draw = ImageDraw.Draw(canvas)
    title_font = image_font(26)
    header_font = image_font(22)
    character_font = image_font(34)
    label_font = image_font(25)
    detail_font = image_font(16)

    draw.rectangle((0, 0, width, TITLE_HEIGHT), fill=TITLE_BACKGROUND)
    center_text(
        draw,
        (0, 0, width, TITLE_HEIGHT),
        f"{config['label']}  •  Eight methods across five stroke-complexity bins",
        title_font,
        fill="#ffffff",
    )
    center_text(
        draw,
        (0, TITLE_HEIGHT, LABEL_WIDTH, TITLE_HEIGHT + HEADER_HEIGHT),
        "METHOD\nMODALITY / PARAMETERS",
        detail_font,
        fill=MUTED_TEXT,
    )
    for column, sample in enumerate(selected):
        x0 = LABEL_WIDTH + column * CELL_SIZE
        center_text(
            draw,
            (x0, TITLE_HEIGHT + 8, x0 + CELL_SIZE, TITLE_HEIGHT + 38),
            sample.label,
            character_font,
        )
        center_text(
            draw,
            (x0, TITLE_HEIGHT + 42, x0 + CELL_SIZE, TITLE_HEIGHT + HEADER_HEIGHT),
            f"{sample.strokes} strokes\n{sample.stroke_bin} bin",
            header_font,
            fill=MUTED_TEXT,
        )

    for row, label in enumerate(row_labels):
        y = TITLE_HEIGHT + HEADER_HEIGHT + row * (CELL_SIZE + ROW_GAP)
        draw.rectangle((0, y, width, y + CELL_SIZE), outline=GRID, width=1)
        if label == "Source (clean)":
            modality, parameters, panel, accent = (
                "CLEAN SOURCE",
                "shared across every row",
                "#f6f8fa",
                "#64748b",
            )
        else:
            modality, parameters, panel, accent = METHOD_STYLES[label]
        draw.rectangle((0, y, LABEL_WIDTH, y + CELL_SIZE), fill=panel)
        draw.rectangle((0, y, 8, y + CELL_SIZE), fill=accent)
        center_text(
            draw,
            (12, y + 56, LABEL_WIDTH, y + 128),
            label,
            label_font,
        )
        center_text(
            draw,
            (12, y + 136, LABEL_WIDTH, y + 200),
            f"{modality}\n{parameters}",
            detail_font,
            fill=MUTED_TEXT,
        )
        for column, sample in enumerate(selected):
            x = LABEL_WIDTH + column * CELL_SIZE
            image_path = (
                clean_images[sample.stem]
                if label == "Source (clean)"
                else method_images[label][sample.stem]
            )
            paste_image(canvas, image_path, x, y)
            draw.rectangle((x, y, x + CELL_SIZE, y + CELL_SIZE), outline=GRID, width=1)

    footer_y = height - FOOTER_HEIGHT
    draw.rectangle((0, footer_y, width, height), fill="#f6f8fa")
    center_text(
        draw,
        (0, footer_y, width, height),
        "Each column keeps the same clean source across all rows for direct visual comparison.",
        detail_font,
        fill=MUTED_TEXT,
    )
    output_dir = REPO_ROOT / "demo" / "showcase"
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / config["output"]
    canvas.save(output_path, format="PNG", optimize=True)
    return {
        "generator": generator,
        "generator_label": config["label"],
        "samples": [asdict(sample) for sample in selected],
        "output": output_path.relative_to(REPO_ROOT).as_posix(),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-root",
        type=Path,
        default=REPO_ROOT.parents[1] / "posion_attack",
        help="Path to the local posion_attack workspace.",
    )
    parser.add_argument(
        "--unihan",
        type=Path,
        default=(
            REPO_ROOT.parents[1]
            / "Posion_paper/ICANN26/main/scripts/Unihan_IRGSources.txt"
        ),
        help="Path to the official Unihan_IRGSources.txt file.",
    )
    args = parser.parse_args()
    data_root = args.data_root.resolve()
    strokes = load_strokes(args.unihan.resolve())
    records = [
        render_generator(data_root, generator, strokes) for generator in GENERATOR_CONFIG
    ]
    manifest_path = REPO_ROOT / "demo" / "showcase" / "manifest.json"
    manifest_path.write_text(
        json.dumps({"generated_from": "CaptchaBench 16K", "showcases": records}, indent=2)
        + "\n",
        encoding="utf-8",
    )
    for record in records:
        labels = ", ".join(
            f"{sample['label']} ({sample['strokes']})" for sample in record["samples"]
        )
        print(f"Created {record['output']} from {labels}")


if __name__ == "__main__":
    main()