File size: 4,595 Bytes
255b4a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
"""Rebuild the README figures from a completed run (clean vs protected).

  python scripts/make_figures.py --run runs/v021_base

Writes:
  assets/compare_clean_vs_adv.png  a few clean|protected pairs, high-res
  assets/grid_before_after_16.png  16 images, clean over protected, 4x4
"""
from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

IMAGES = Path("examples/testset60")
ASSETS = Path("assets")
CELL = 320
PAD = 14
LABEL_H = 26


def load_manifest() -> list[tuple[str, str]]:
    rows = []
    for line in Path("examples/testset60.csv").read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#"):
            p, t = line.split(",", 1)
            rows.append((Path(p).name, t.strip()))
    return rows


def font(size: int) -> ImageFont.FreeTypeFont:
    for p in ("/System/Library/Fonts/Supplemental/Arial.ttf",
              "/System/Library/Fonts/Helvetica.ttc",
              "/Library/Fonts/Arial.ttf"):
        if Path(p).exists():
            return ImageFont.truetype(p, size)
    return ImageFont.load_default()


def sq(img: Image.Image, n: int = CELL) -> Image.Image:
    return img.convert("RGB").resize((n, n), Image.LANCZOS)


def labeled(img: Image.Image, text: str, f) -> Image.Image:
    canvas = Image.new("RGB", (img.width, img.height + LABEL_H), "white")
    canvas.paste(img, (0, LABEL_H))
    d = ImageDraw.Draw(canvas)
    d.text((4, 4), text, fill="black", font=f)
    return canvas


def compare(run: Path, rows: list[tuple[str, str]], picks: list[str]) -> None:
    f = font(18)
    fh = font(20)
    items = [(fn, t) for fn, t in rows if fn in picks]
    w = 2 * CELL + 3 * PAD
    h = len(items) * (CELL + LABEL_H + PAD) + PAD + 30
    canvas = Image.new("RGB", (w, h), "white")
    d = ImageDraw.Draw(canvas)
    d.text((PAD, 6), "Clean", fill="black", font=fh)
    d.text((2 * PAD + CELL, 6), "Protected (VEIL-PGD)", fill="black", font=fh)
    y = 34
    for fn, truth in items:
        stem = Path(fn).stem
        clean = sq(Image.open(IMAGES / fn))
        adv = sq(Image.open(run / "adv" / f"{stem}.png"))
        canvas.paste(labeled(clean, truth, f), (PAD, y))
        canvas.paste(labeled(adv, truth, f), (2 * PAD + CELL, y))
        y += CELL + LABEL_H + PAD
    ASSETS.mkdir(exist_ok=True)
    canvas.save(ASSETS / "compare_clean_vs_adv.png")
    print("wrote assets/compare_clean_vs_adv.png", canvas.size)


def grid(run: Path, rows: list[tuple[str, str]], picks: list[str]) -> None:
    f = font(16)
    items = [(fn, t) for fn, t in rows if fn in picks][:16]
    cols = 4
    cellw = CELL
    cellh = 2 * CELL + LABEL_H
    n = len(items)
    grows = (n + cols - 1) // cols
    w = cols * cellw + (cols + 1) * PAD
    h = grows * (cellh + PAD) + PAD
    canvas = Image.new("RGB", (w, h), "white")
    for i, (fn, truth) in enumerate(items):
        r, c = divmod(i, cols)
        x = PAD + c * (cellw + PAD)
        y = PAD + r * (cellh + PAD)
        stem = Path(fn).stem
        clean = sq(Image.open(IMAGES / fn))
        adv = sq(Image.open(run / "adv" / f"{stem}.png"))
        block = Image.new("RGB", (cellw, cellh), "white")
        d = ImageDraw.Draw(block)
        d.text((4, 4), truth, fill="black", font=f)
        block.paste(clean, (0, LABEL_H))
        block.paste(adv, (0, LABEL_H + CELL))
        canvas.paste(block, (x, y))
    ASSETS.mkdir(exist_ok=True)
    canvas.save(ASSETS / "grid_before_after_16.png")
    print("wrote assets/grid_before_after_16.png", canvas.size)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", default="runs/v021_base")
    args = ap.parse_args()
    run = Path(args.run)
    rows = load_manifest()
    have = [fn for fn, _ in rows if (run / "adv" / f"{Path(fn).stem}.png").exists()]
    # one image per class for the small comparison
    seen, comp_picks = set(), []
    for fn, t in rows:
        if fn in have and t not in seen:
            seen.add(t)
            comp_picks.append(fn)
    compare(run, rows, comp_picks[:4])
    # 16 images spread across classes for the grid
    grid_picks = []
    by_class: dict[str, list[str]] = {}
    for fn, t in rows:
        if fn in have:
            by_class.setdefault(t, []).append(fn)
    rr = 0
    while len(grid_picks) < 16 and any(len(v) > rr for v in by_class.values()):
        for t in by_class:
            if len(by_class[t]) > rr and len(grid_picks) < 16:
                grid_picks.append(by_class[t][rr])
        rr += 1
    grid(run, rows, grid_picks)


if __name__ == "__main__":
    main()