scripts/gen_textvqa_patch_figures.py

"""Generate paper-ready patch figures from a real TextVQA image.

Outputs three groups, each in BOTH 2D (square, flat) and oblique (lying-flat
trapezoid with shadow) form:

  1. Full 5x5 split — 25 patches, individual files (the source-image
     decomposition the paper uses to recover the original).
  2. M ROI-highlighted patches — query-relevant patches with a red border.
  3. M attention-map figures — synthetic per-patch heatmap (jet colormap
     overlaid on the patch) corresponding 1:1 with the ROI patches.

No text / axes / arrows / decorations on any tile — strictly figures.
"""
from __future__ import annotations
import os
import numpy as np
from PIL import Image, ImageDraw, ImageFilter

# -------------------------------------------------------------------------
# Source + outputs
# -------------------------------------------------------------------------
TEXTVQA_IMG = "/opt/tiger/thothvl_pretrain/doc/figures/textvqa_patch_tiles_src/camera_0.png"
NAS_OUT = "/mnt/bn/leonworkspace/terry/ce-task/figures/textvqa_patch_tiles"
LOCAL_OUT = "/opt/tiger/thothvl_pretrain/doc/figures/textvqa_patch_tiles"

# -------------------------------------------------------------------------
# Config
# -------------------------------------------------------------------------
GRID = 5                          # 5x5 patch grid
PATCH_PX = 180                    # source pixels per patch (patch is square)

# Oblique tile canvas
TILE_W, TILE_H = 240, 200
SHADOW_BLUR = 7
SHADOW_OFFSET = (6, 10)
SHADOW_OPACITY = 110
SHADOW_COLOR = (60, 65, 75)

# Flat 2D tile canvas (no perspective)
FLAT_W, FLAT_H = 190, 190
FLAT_PADDING = 5                  # pixels of canvas margin around the patch

# ROI / sink / attention
# (row, col) of query-relevant patches and sink-token patches in the 5x5
# grid. Indices are mapped from the 24x24 attention grid in the
# pseudo-label visualization for /textvqa/0_pseudo_label.png:
#   FG region (yellow)        → grid x[4..9], y[5..9]   → patches (1, 0..2)
#   Sink token (gray Ignore)  → grid x[4],   y[15]      → patches (3, 0..1)
ROI_PATCHES  = [(1, 0), (1, 1), (1, 2)]
SINK_PATCHES = [(3, 0), (3, 1)]
M = len(ROI_PATCHES)
N_SINK = len(SINK_PATCHES)

ROI_BORDER_COLOR  = (235, 60, 50, 255)    # crimson — query-relevant ROI
SINK_BORDER_COLOR = (90, 110, 140, 255)   # slate-blue — sink token
BORDER_WIDTH = 7                           # px on the source patch
HEATMAP_ALPHA = 0.55                       # 0..1 attention overlay strength

# -------------------------------------------------------------------------
# Geometry
# -------------------------------------------------------------------------
def perspective_coeffs(src_corners, dst_corners):
    """Solve 8-coefficient PIL perspective transform mapping dst → src."""
    rows = []
    for (sx, sy), (dx, dy) in zip(src_corners, dst_corners):
        rows.append([dx, dy, 1, 0, 0, 0, -sx * dx, -sx * dy])
        rows.append([0, 0, 0, dx, dy, 1, -sy * dx, -sy * dy])
    A = np.array(rows, dtype=np.float64)
    B = np.array(src_corners, dtype=np.float64).reshape(8)
    return tuple(np.linalg.solve(A, B))


def trapezoid_corners(cw: int, ch: int):
    """Symmetric trapezoid pulled toward the canvas edges."""
    return [
        (cw * 0.235, ch * 0.080),  # TL
        (cw * 0.765, ch * 0.080),  # TR
        (cw * 0.945, ch * 0.840),  # BR
        (cw * 0.055, ch * 0.840),  # BL
    ]


# -------------------------------------------------------------------------
# Source loading + heatmap synthesis
# -------------------------------------------------------------------------
def load_square_source(path: str, size: int) -> Image.Image:
    im = Image.open(path).convert("RGB")
    w, h = im.size
    s = min(w, h)
    left, top = (w - s) // 2, (h - s) // 2
    im = im.crop((left, top, left + s, top + s))
    return im.resize((size, size), Image.LANCZOS)


def jet_colormap(t: np.ndarray) -> np.ndarray:
    """Approximate the matplotlib 'jet' colormap on a 2-D array t in [0,1].

    Returns (H, W, 3) uint8.
    """
    t = np.clip(t, 0.0, 1.0)
    r = np.clip(1.5 - np.abs(4.0 * t - 3.0), 0.0, 1.0)
    g = np.clip(1.5 - np.abs(4.0 * t - 2.0), 0.0, 1.0)
    b = np.clip(1.5 - np.abs(4.0 * t - 1.0), 0.0, 1.0)
    return (np.stack([r, g, b], axis=-1) * 255.0).astype(np.uint8)


def synth_heatmap(patch: Image.Image, peak_xy: tuple[float, float],
                  sigma_frac: float = 0.30) -> Image.Image:
    """Render an attention-map overlay on a patch.

    A 2-D Gaussian centred at peak_xy (in [0,1]^2 patch coords) is colour
    -mapped (jet) and alpha-blended onto the source patch.
    """
    w, h = patch.size
    y_idx, x_idx = np.mgrid[0:h, 0:w]
    cx, cy = peak_xy[0] * w, peak_xy[1] * h
    sigma = sigma_frac * max(w, h)
    g = np.exp(-((x_idx - cx) ** 2 + (y_idx - cy) ** 2) / (2.0 * sigma ** 2))
    g = (g - g.min()) / (g.max() - g.min() + 1e-9)
    cmap = jet_colormap(g)
    base = np.asarray(patch.convert("RGB"), dtype=np.float32)
    over = cmap.astype(np.float32)
    blended = ((1.0 - HEATMAP_ALPHA) * base + HEATMAP_ALPHA * over)
    return Image.fromarray(np.clip(blended, 0, 255).astype(np.uint8))


def add_border(patch: Image.Image, color=ROI_BORDER_COLOR,
               width: int = BORDER_WIDTH) -> Image.Image:
    """Draw a thick coloured rectangle just inside the patch edges."""
    out = patch.convert("RGBA").copy()
    d = ImageDraw.Draw(out)
    w, h = out.size
    d.rectangle([(width // 2, width // 2),
                 (w - width // 2 - 1, h - width // 2 - 1)],
                outline=color, width=width)
    return out


# -------------------------------------------------------------------------
# Renderers (oblique + flat)
# -------------------------------------------------------------------------
def render_oblique(patch_img: Image.Image) -> Image.Image:
    cw, ch = TILE_W, TILE_H
    dst = trapezoid_corners(cw, ch)
    src = [(0, 0), (patch_img.width, 0),
           (patch_img.width, patch_img.height), (0, patch_img.height)]
    coeffs = perspective_coeffs(src, dst)
    warped = patch_img.convert("RGBA").transform(
        (cw, ch), Image.PERSPECTIVE, coeffs, resample=Image.BICUBIC,
    )
    bg = Image.new("RGBA", (cw, ch), (0, 0, 0, 0))
    shadow_quad = [(x + SHADOW_OFFSET[0], y + SHADOW_OFFSET[1]) for x, y in dst]
    sm = Image.new("L", (cw, ch), 0)
    ImageDraw.Draw(sm).polygon(shadow_quad, fill=SHADOW_OPACITY)
    sm = sm.filter(ImageFilter.GaussianBlur(SHADOW_BLUR))
    sl = np.zeros((ch, cw, 4), dtype=np.uint8)
    sl[:, :, 0:3] = SHADOW_COLOR
    sl[:, :, 3] = np.asarray(sm, dtype=np.uint8)
    bg = Image.alpha_composite(bg, Image.fromarray(sl, "RGBA"))
    bg = Image.alpha_composite(bg, warped)
    return bg


def render_flat(patch_img: Image.Image) -> Image.Image:
    """Square 2-D tile, no perspective, transparent background."""
    p = patch_img.convert("RGBA")
    inner = FLAT_W - 2 * FLAT_PADDING
    p = p.resize((inner, inner), Image.LANCZOS)
    bg = Image.new("RGBA", (FLAT_W, FLAT_H), (0, 0, 0, 0))
    bg.paste(p, (FLAT_PADDING, FLAT_PADDING), p)
    return bg


# -------------------------------------------------------------------------
# Save helpers
# -------------------------------------------------------------------------
def save_pair(im: Image.Image, name: str):
    """Save PNG with alpha (modest resolution for paper figures)."""
    for d in (NAS_OUT, LOCAL_OUT):
        im.save(os.path.join(d, f"{name}.png"), dpi=(150, 150))


# -------------------------------------------------------------------------
# Main
# -------------------------------------------------------------------------
def draw_grid(img: Image.Image, n: int = GRID,
              color=(255, 255, 255, 230), width: int = 4) -> Image.Image:
    """Draw n×n grid lines on top of a square source image."""
    out = img.convert("RGBA").copy()
    d = ImageDraw.Draw(out)
    w, h = out.size
    for k in range(1, n):
        x = round(k * w / n)
        y = round(k * h / n)
        d.line([(x, 0), (x, h)], fill=color, width=width)
        d.line([(0, y), (w, y)], fill=color, width=width)
    return out


def render_all():
    os.makedirs(NAS_OUT, exist_ok=True)
    os.makedirs(LOCAL_OUT, exist_ok=True)

    src_size = PATCH_PX * GRID
    sq = load_square_source(TEXTVQA_IMG, src_size)
    print(f"[src] {TEXTVQA_IMG} -> {src_size}x{src_size} square")

    # 0) Full source image with 5×5 grid overlay — flat + oblique at a
    #    larger canvas than the per-patch tiles so grid lines stay readable.
    full_grid = draw_grid(sq)
    full_flat = full_grid.resize((600, 600), Image.LANCZOS)
    save_pair(full_flat, "source_grid_flat")

    # Very flat "lying-down" perspective: short canvas height shrinks the
    # left/right slanted edges so the figure reads as a card on a table.
    big_w, big_h = 720, 280
    dst = [
        (big_w * 0.300, big_h * 0.120),  # TL
        (big_w * 0.700, big_h * 0.120),  # TR
        (big_w * 0.970, big_h * 0.880),  # BR
        (big_w * 0.030, big_h * 0.880),  # BL
    ]
    src = [(0, 0), (full_grid.width, 0),
           (full_grid.width, full_grid.height), (0, full_grid.height)]
    coeffs = perspective_coeffs(src, dst)
    warped = full_grid.convert("RGBA").transform(
        (big_w, big_h), Image.PERSPECTIVE, coeffs, resample=Image.BICUBIC)
    bg = Image.new("RGBA", (big_w, big_h), (0, 0, 0, 0))
    sm = Image.new("L", (big_w, big_h), 0)
    ImageDraw.Draw(sm).polygon(
        [(x + SHADOW_OFFSET[0] * 2, y + SHADOW_OFFSET[1] * 2) for x, y in dst],
        fill=SHADOW_OPACITY)
    sm = sm.filter(ImageFilter.GaussianBlur(SHADOW_BLUR * 2))
    sl = np.zeros((big_h, big_w, 4), dtype=np.uint8)
    sl[:, :, 0:3] = SHADOW_COLOR
    sl[:, :, 3] = np.asarray(sm, dtype=np.uint8)
    bg = Image.alpha_composite(bg, Image.fromarray(sl, "RGBA"))
    bg = Image.alpha_composite(bg, warped)
    save_pair(bg, "source_grid_oblique")
    # Also save the bare source (no grid) at the same large canvas.
    save_pair(sq.resize((600, 600), Image.LANCZOS), "source_flat")
    print(f"  source — flat + oblique (with and without grid)")

    # 1) Full 5x5 split — every patch as oblique + flat.
    for i in range(GRID):
        for j in range(GRID):
            x0, y0 = j * PATCH_PX, i * PATCH_PX
            patch = sq.crop((x0, y0, x0 + PATCH_PX, y0 + PATCH_PX))
            save_pair(render_oblique(patch), f"split_{i}_{j}_oblique")
            save_pair(render_flat(patch),    f"split_{i}_{j}_flat")
            print(f"  split  ({i},{j}) — oblique + flat")

    # 2) ROI-highlighted patches + (3) attention maps for the M selected ROIs.
    for k, (i, j) in enumerate(ROI_PATCHES):
        x0, y0 = j * PATCH_PX, i * PATCH_PX
        patch = sq.crop((x0, y0, x0 + PATCH_PX, y0 + PATCH_PX))

        # ROI = patch with crimson border (the highlight)
        roi = add_border(patch, color=ROI_BORDER_COLOR)
        save_pair(render_oblique(roi), f"roi_{k}_oblique")
        save_pair(render_flat(roi),    f"roi_{k}_flat")

        # Attention map = synthetic Gaussian heatmap centred on the ROI
        # patch's geometric centre (could be replaced with a real attention
        # map per layer/head if available).
        att = synth_heatmap(patch, peak_xy=(0.5, 0.5), sigma_frac=0.28)
        save_pair(render_oblique(att), f"attn_{k}_oblique")
        save_pair(render_flat(att),    f"attn_{k}_flat")
        print(f"  ROI    {k} = grid({i},{j}) — roi + attn (oblique + flat)")

    # 4) Sink-token patches highlighted with a distinct slate-blue border.
    for k, (i, j) in enumerate(SINK_PATCHES):
        x0, y0 = j * PATCH_PX, i * PATCH_PX
        patch = sq.crop((x0, y0, x0 + PATCH_PX, y0 + PATCH_PX))
        sink = add_border(patch, color=SINK_BORDER_COLOR)
        save_pair(render_oblique(sink), f"sink_{k}_oblique")
        save_pair(render_flat(sink),    f"sink_{k}_flat")
        print(f"  SINK   {k} = grid({i},{j}) — sink (oblique + flat)")

    print(f"[done] wrote split + ROI + attn + sink → {NAS_OUT}")


if __name__ == "__main__":
    render_all()