AIF/utils_vis.py at main · kaist-ami/AIF · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import os

import torch
from PIL import Image, ImageDraw
import matplotlib.cm as cm
import numpy as np


def overlay_mask(images, token_mask_map, att_shape):
    # Mask out the image regions where token_mask_map is non-zero
    in_img = images[0]
    mask_map = token_mask_map.reshape(att_shape).to(torch.float32).detach().cpu().numpy()
    mask_map = (mask_map - mask_map.min()) / (mask_map.max() - mask_map.min())
    mask_img = (np.repeat(mask_map[:, :, None], 3, axis=-1) * 255).astype(np.uint8)
    mask_img_re = Image.fromarray(mask_img).resize(in_img.size, Image.NEAREST)
    masked_img = Image.fromarray((np.array(mask_img_re) == 0) * np.array(in_img))
    return masked_img


def overlay_att_to_img(rel_att, att_shape, image, resize_type=Image.NEAREST):
    # Blend a per-token relevance map onto an image as a jet-colormap heatmap
    if isinstance(rel_att, torch.Tensor):
        rel_att = rel_att.to(torch.float32).detach().cpu().numpy()
    rel_att_map = rel_att.reshape(att_shape)
    rel_att_map_norm = (rel_att_map - rel_att_map.min()) / (rel_att_map.max() - rel_att_map.min())

    heatmap = (cm.jet(rel_att_map_norm)[:, :, :3] * 255).astype(np.uint8)
    heatmap_image = Image.fromarray(heatmap).resize(image.size, resize_type)
    overlay_image = Image.blend(image, heatmap_image, alpha=0.6)
    return overlay_image, heatmap_image


def make_panel(labeled_images, n_cols=2, label_height=24):
    # Arrange labeled images into a grid panel with text labels
    size = labeled_images[0][1].size
    n_rows = -(-len(labeled_images) // n_cols)  # ceil division
    cell_w, cell_h = size[0], size[1] + label_height
    panel = Image.new("RGB", (cell_w * n_cols, cell_h * n_rows), "white")
    draw = ImageDraw.Draw(panel)
    for i, (label, img) in enumerate(labeled_images):
        row, col = divmod(i, n_cols)
        x, y = col * cell_w, row * cell_h
        panel.paste(img.resize(size), (x, y + label_height))
        draw.text((x + 4, y + 4), label, fill="black")
    return panel


def save_aif_artifacts(output_dir, image, img2text_att_max, token_entropy, token_mask_map, att_shape):
    # Save a combined attention/entropy/mask visualization panel to output_dir
    os.makedirs(output_dir, exist_ok=True)

    att_mu = img2text_att_max.mean(dim=0)[0]  # mean over layers -> [num_img_tokens]
    att_overlay, _ = overlay_att_to_img(att_mu, att_shape, image)
    entropy_overlay, _ = overlay_att_to_img(token_entropy, att_shape, image)
    masked_img = overlay_mask([image], token_mask_map, att_shape)

    panel = make_panel(
        [
            ("image", image),
            ("attention", att_overlay),
            ("token_entropy", entropy_overlay),
            ("masked_tokens", masked_img),
        ]
    )
    panel.save(os.path.join(output_dir, "aif_visualization.png"))
    print(f"[aif] saved combined attention/entropy/mask visualization to {output_dir}")