# coding=utf-8 from __future__ import absolute_import, division, print_function import argparse import math import os import numpy as np import torch from PIL import Image, ImageDraw, ImageFilter from torchvision import transforms from models.modeling import CONFIGS, VisionTransformer IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] def build_model(args): config = CONFIGS[args.model_type] config.split = args.split config.slide_step = args.slide_step model = VisionTransformer( config, args.img_size, zero_head=True, num_classes=args.num_classes, smoothing_value=0.0, grad_checkpoint=False, ) if args.pretrained_dir: model.load_from(np.load(args.pretrained_dir)) if args.checkpoint: checkpoint = torch.load(args.checkpoint, map_location="cpu") state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint model.load_state_dict(state_dict, strict=False) model.to(args.device) model.eval() return model, config def load_center_crop(image_path, img_size): image = Image.open(image_path).convert("RGB") resized = image.resize((600, 600), Image.BILINEAR) left = (600 - img_size) // 2 top = (600 - img_size) // 2 cropped = resized.crop((left, top, left + img_size, top + img_size)) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), ]) tensor = transform(cropped).unsqueeze(0) return cropped, tensor def token_to_box(token_index, grid_size, patch_size, slide_step): row = int(token_index) // grid_size col = int(token_index) % grid_size x0 = col * slide_step y0 = row * slide_step return x0, y0, x0 + patch_size, y0 + patch_size def draw_boxes(image, boxes, output_path): canvas = image.copy() draw = ImageDraw.Draw(canvas) for i, box in enumerate(boxes): color = "red" if i == 0 else "deepskyblue" draw.rectangle(box, outline=color, width=3) canvas.save(output_path) def build_attention_mask(attention_map, grid_size, img_size, mode, gamma, blur_radius): if mode == "max": heat = attention_map.max(axis=0) else: heat = attention_map.mean(axis=0) heat = heat.reshape(grid_size, grid_size) heat = heat - heat.min() heat = heat / (heat.max() + 1e-8) heat = np.power(heat, gamma) heat_img = Image.fromarray(np.uint8(heat * 255), mode="L") heat_img = heat_img.resize((img_size, img_size), Image.BICUBIC) if blur_radius > 0: heat_img = heat_img.filter(ImageFilter.GaussianBlur(radius=blur_radius)) heat = np.asarray(heat_img).astype(np.float32) / 255.0 heat = heat - heat.min() heat = heat / (heat.max() + 1e-8) return heat def save_focus_image(image, mask, output_path, background_strength): image_np = np.asarray(image).astype(np.float32) mask_3c = mask[..., None] multiplier = background_strength + (1.0 - background_strength) * mask_3c focused = image_np * multiplier focused = np.clip(focused, 0, 255).astype(np.uint8) Image.fromarray(focused).save(output_path) def save_attention_heatmap(mask, output_path): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt plt.figure(figsize=(5, 5)) plt.imshow(mask, cmap="jet") plt.axis("off") plt.tight_layout(pad=0) plt.savefig(output_path, dpi=200, bbox_inches="tight", pad_inches=0) plt.close() def save_paper_style_pair(original, focus, output_path): width, height = original.size pair = Image.new("RGB", (width, height * 2 + 6), "white") pair.paste(original, (0, 0)) pair.paste(focus, (0, height + 6)) pair.save(output_path) def make_contact_sheet(parts, output_path, scale): if not parts: return scaled = [part.resize((part.width * scale, part.height * scale), Image.NEAREST) for part in parts] cols = min(6, len(scaled)) rows = math.ceil(len(scaled) / cols) cell_w, cell_h = scaled[0].size sheet = Image.new("RGB", (cols * cell_w, rows * cell_h), "white") for i, part in enumerate(scaled): x = (i % cols) * cell_w y = (i // cols) * cell_h sheet.paste(part, (x, y)) sheet.save(output_path) def make_output_path(output_dir, image_type, filename): type_dir = os.path.join(output_dir, image_type) os.makedirs(type_dir, exist_ok=True) return os.path.join(type_dir, filename) def main(): parser = argparse.ArgumentParser(description="Visualize TransFG part selection tokens.") parser.add_argument("--image", required=True, help="Path to a CUB image.") parser.add_argument("--pretrained_dir", default="pretrained/ViT-B_16.npz", help="Path to ViT-B_16.npz.") parser.add_argument("--checkpoint", default=None, help="Optional trained TransFG checkpoint.") parser.add_argument("--output_dir", default="visual_outputs", help="Directory for saved PNG files.") parser.add_argument("--model_type", default="ViT-B_16", choices=["ViT-B_16", "ViT-B_32", "ViT-L_16", "ViT-L_32", "ViT-H_14"]) parser.add_argument("--num_classes", type=int, default=200) parser.add_argument("--img_size", type=int, default=448) parser.add_argument("--split", default="overlap", choices=["overlap", "non-overlap"]) parser.add_argument("--slide_step", type=int, default=12) parser.add_argument("--part_rank", type=int, default=0, help="Which ranked selected token to crop as the center part.") parser.add_argument("--scale", type=int, default=10, help="Upscale factor for selected token crops.") parser.add_argument("--mask_mode", default="max", choices=["max", "mean"], help="How to merge head attention maps.") parser.add_argument("--mask_gamma", type=float, default=0.55, help="Lower values expand the highlighted area.") parser.add_argument("--mask_blur", type=float, default=10.0, help="Gaussian blur radius for the attention mask.") parser.add_argument("--background_strength", type=float, default=0.16, help="Brightness retained in low-attention background.") parser.add_argument("--cpu", action="store_true", help="Force CPU inference.") args = parser.parse_args() args.device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu") os.makedirs(args.output_dir, exist_ok=True) model, config = build_model(args) crop_image, tensor = load_center_crop(args.image, args.img_size) tensor = tensor.to(args.device) with torch.no_grad(): _ = model(tensor) patch_size = config.patches["size"][0] if args.split == "overlap": grid_size = (args.img_size - patch_size) // args.slide_step + 1 else: grid_size = args.img_size // patch_size indices = model.transformer.encoder.last_part_patch_indices[0].cpu().numpy() scores = model.transformer.encoder.last_part_scores[0].cpu().numpy() attention_map = model.transformer.encoder.last_part_attention_map[0].cpu().numpy() order = np.argsort(-scores) ranked_indices = indices[order] ranked_scores = scores[order] boxes = [token_to_box(index, grid_size, patch_size, args.slide_step) for index in ranked_indices] parts = [crop_image.crop(box) for box in boxes] stem = os.path.splitext(os.path.basename(args.image))[0] overlay_path = make_output_path(args.output_dir, "part_selection_overlays", f"{stem}_part_selection_overlay.png") center_path = make_output_path(args.output_dir, "selected_part_crops", f"{stem}_selected_part_rank{args.part_rank}.png") sheet_path = make_output_path(args.output_dir, "selected_parts_sheets", f"{stem}_selected_parts_sheet.png") focus_path = make_output_path(args.output_dir, "attention_focus", f"{stem}_attention_focus.png") heatmap_path = make_output_path(args.output_dir, "attention_heatmaps", f"{stem}_attention_heatmap.png") pair_path = make_output_path(args.output_dir, "paper_style_pairs", f"{stem}_paper_style_pair.png") draw_boxes(crop_image, boxes, overlay_path) mask = build_attention_mask( attention_map, grid_size=grid_size, img_size=args.img_size, mode=args.mask_mode, gamma=args.mask_gamma, blur_radius=args.mask_blur, ) save_focus_image(crop_image, mask, focus_path, args.background_strength) save_attention_heatmap(mask, heatmap_path) focus_image = Image.open(focus_path).convert("RGB") save_paper_style_pair(crop_image, focus_image, pair_path) rank = min(max(args.part_rank, 0), len(parts) - 1) center = parts[rank].resize((parts[rank].width * args.scale, parts[rank].height * args.scale), Image.NEAREST) center.save(center_path) make_contact_sheet(parts, sheet_path, args.scale) print("Saved overlay:", overlay_path) print("Saved attention focus:", focus_path) print("Saved attention heatmap:", heatmap_path) print("Saved paper-style pair:", pair_path) print("Saved center selected part:", center_path) print("Saved all selected parts:", sheet_path) print("Top selected token index:", int(ranked_indices[rank])) print("Top selected token score:", float(ranked_scores[rank])) if __name__ == "__main__": main()