visualize_part_selection.py

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