"""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()