code/make_hero_figure.py

"""
Hero / Figure 1 — the one-directional cross-source shortcut, told in one frame.

Designed composition (not a data plot): two worlds (studio, wild) joined by an
ASYMMETRIC pair of arcs with the headline number anchored in the centre, and the
mechanism SHOWN — real HiResCAM on the SAME wild coriander image: the wild-trained
model glows on the spice, the studio-trained model glows on the background.
Mono numerals for scientific precision; warm paper canvas.

GPU-light. Reuses the corner-fixed CAM from plot_gradcam_contrast.py.
"""
import sys, os, json
from pathlib import Path

_base = "/mnt/d/SpiceNet" if os.path.exists("/mnt/d/SpiceNet") else "D:/SpiceNet"
sys.path.insert(0, _base)

import numpy as np
import torch
from PIL import Image
import figstyle
from plot_gradcam_contrast import HiResCAM, _view, _postprocess, _overlay, _load, _find_ckpts

ROOT = Path(_base)
OUT = ROOT / "outputs" / "hero_shortcut"
IN_MANIFEST = ROOT / "outputs" / "manifest_overlap_indian.json"
SS_MANIFEST = ROOT / "outputs" / "manifest_overlap_ss.json"
W, H = 13.0, 10.0

PAPER = "#F6F3EE"
INK = "#1c1a17"
MUTED = "#8a8178"
RULE = "#ddd6cc"
STUDIO = figstyle.PALETTE["studio"]
WILD = figstyle.PALETTE["wild"]
COLLAPSE = figstyle.PALETTE["cross_broken"]
FREE = figstyle.PALETTE["within"]


def _wsl(p):
    if os.name != "nt" and len(p) > 2 and p[1] == ":":
        return f"/mnt/{p[0].lower()}/" + p[2:].replace("\\", "/").lstrip("/")
    return p


def _img(manifest, cls):
    m = json.load(open(manifest))
    classes = [c["name"] for c in sorted(m["classes"], key=lambda c: c["index"])]
    idx = classes.index(cls)
    for split in ("test", "val", "train"):
        for p, y in m["samples"][split]:
            if int(y) == idx:
                q = _wsl(p)
                if os.path.exists(q):
                    return q
    return None


def _square(path, size=440):
    im = Image.open(path).convert("RGB")
    w, h = im.size; s = min(w, h)
    im = im.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
    return np.asarray(im.resize((size, size)))


def _dax(fig, x, y, w, h):
    return fig.add_axes([x / W, y / H, w / W, h / H])


def _framed(fig, x, y, s, img, color, lw=3.2):
    a = _dax(fig, x, y, s, s)
    a.imshow(img); a.set_xticks([]); a.set_yticks([])
    for sp in a.spines.values():
        sp.set_edgecolor(color); sp.set_linewidth(lw)
    return a


def _cams():
    """HiResCAM overlays on the same wild coriander image + the plain input."""
    ss_ckpt, in_ckpt = _find_ckpts()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    wild_model, studio_model = _load(ss_ckpt, device), _load(in_ckpt, device)
    tensor, disp = _view(_img(SS_MANIFEST, "coriander"))
    tensor = tensor.to(device)

    def _t(m):
        try: return m.backbone.blocks[-1]
        except Exception: return m.backbone.conv_head
    gw, gs = HiResCAM(wild_model, _t(wild_model)), HiResCAM(studio_model, _t(studio_model))
    cam_w, _ = gw(tensor); cam_s, _ = gs(tensor)
    gw.remove(); gs.remove()
    return (_overlay(disp, _postprocess(cam_w, disp.shape[0])),
            _overlay(disp, _postprocess(cam_s, disp.shape[0])), disp)


def main():
    figstyle.apply()
    import matplotlib.pyplot as plt
    from matplotlib.patches import FancyArrowPatch
    plt.rcParams["font.family"] = "sans-serif"

    ov_wild, ov_studio, disp = _cams()

    fig = plt.figure(figsize=(W, H)); fig.patch.set_facecolor(PAPER)
    ax = fig.add_axes([0, 0, 1, 1]); ax.set_xlim(0, W); ax.set_ylim(0, H)
    ax.axis("off"); ax.set_facecolor(PAPER)

    def mono(x, y, s, size, color, ha="center"):
        ax.text(x, y, s, ha=ha, va="center", fontsize=size, color=color,
                fontweight="bold", family="monospace", zorder=6)

    # ── header ──────────────────────────────────────────────────────────────
    ax.text(0.7, 9.62, "C R O S S - S O U R C E   S H O R T C U T", ha="left", va="center",
            fontsize=11, color=COLLAPSE, fontweight="bold")
    ax.text(0.7, 9.06, "Studio spice benchmarks lie", ha="left", va="center",
            fontsize=25, color=INK, fontweight="bold")
    ax.text(0.7, 8.55, "A model at 100% on studio images scores 62% in the wild — the collapse is one-directional.",
            ha="left", va="center", fontsize=12.5, color=MUTED)
    ax.plot([0.7, 12.3], [8.18, 8.18], color=RULE, lw=1.2, zorder=1)

    # ── two worlds ──────────────────────────────────────────────────────────
    s = 2.6
    _framed(fig, 0.7, 5.2, s, _square(_img(IN_MANIFEST, "coriander")), STUDIO)
    ax.text(0.7 + s / 2, 4.92, "STUDIO", ha="center", va="top", fontsize=13.5, fontweight="bold", color=STUDIO)
    _framed(fig, W - 0.7 - s, 5.2, s, _square(_img(SS_MANIFEST, "coriander")), WILD)
    ax.text(W - 0.7 - s / 2, 4.92, "IN THE WILD", ha="center", va="top", fontsize=13.5, fontweight="bold", color=WILD)

    xl, xr, xc = 0.7 + s + 0.25, W - 0.7 - s - 0.25, W / 2
    # collapse arc (studio -> wild): thick, arcs up, pure visual
    ax.add_patch(FancyArrowPatch((xl, 6.55), (xr, 6.55), connectionstyle="arc3,rad=-0.25",
                 arrowstyle="-|>", mutation_scale=32, lw=6.5, color=COLLAPSE, zorder=2))
    # free arc (wild -> studio): hairline, arcs down, pure visual
    ax.add_patch(FancyArrowPatch((xr, 5.55), (xl, 5.55), connectionstyle="arc3,rad=-0.25",
                 arrowstyle="-|>", mutation_scale=18, lw=1.8, color=FREE, zorder=2))
    # centre anchor: the headline number, filling the void between the arcs
    mono(xc, 6.28, "− 37.8 pp", 30, COLLAPSE)
    mono(xc, 5.82, "100%  →  62%      COLLAPSE", 12.5, COLLAPSE)
    ax.text(xc, 7.68, "train STUDIO   →   test WILD", ha="center", fontsize=11, fontweight="bold", color=COLLAPSE)
    ax.text(xc, 4.55, "train WILD  →  test STUDIO", ha="center", fontsize=10.5, fontweight="bold", color=FREE)
    mono(xc, 4.2, "− 0.5 pp   free", 11.5, FREE)

    # ── mechanism: SHOW where each model looks (3 panels, input in centre) ───
    ax.plot([0.7, 12.3], [3.95, 3.95], color=RULE, lw=1.2, zorder=1)
    ax.text(0.7, 3.68, "W H Y", ha="left", va="center", fontsize=11, color=INK, fontweight="bold")
    ax.text(1.7, 3.68, "— where each model looks, on the same in-the-wild coriander (HiResCAM)",
            ha="left", va="center", fontsize=11.5, color=MUTED)

    ms = 2.35
    lefts = [1.35, xc - ms / 2, W - 1.35 - ms]
    imgs = [(ov_wild, FREE, "WILD-TRAINED", "the spice   ✓  coriander"),
            (disp, MUTED, "same wild image", ""),
            (ov_studio, COLLAPSE, "STUDIO-TRAINED", "the background   ✗  green_cardamom")]
    for x0, (im, col, top, bot) in zip(lefts, imgs):
        _framed(fig, x0, 0.95, ms, im, col if col != MUTED else RULE, 3.2 if col != MUTED else 2.0)
        ax.text(x0 + ms / 2, 3.5, top, ha="center", va="center", fontsize=12,
                fontweight="bold", color=col, style="italic" if col == MUTED else "normal")
        if bot:
            ax.text(x0 + ms / 2, 0.62, bot, ha="center", va="center", fontsize=11,
                    fontweight="bold", color=col)

    figstyle.save(fig, str(OUT), also_tiff=True)


if __name__ == "__main__":
    main()