visualize_attn.py

# %%
from functools import partial

import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
from matplotlib.patches import Rectangle
from PIL import Image

from configs import consts as C
from configs import paths as P
from data.ofa_dataset import OFADataset
from fairseq.data.dictionary import Dictionary
from fairseq.data.encoders.gpt2_bpe import GPT2BPE, GPT2BPEConfig
from utils import img_fmt as IF

# %%
row_id = 1
attn_layer_id = 3
batch_id = 0
iot = 'which region does the text " {question} " describe?'
max_src_length = 80
exp_tag = "example"
arch = "tiny"
trainP = "Instruct-2"
valP = "Base"
testP = "Instruct-2"
split = f"test_private-P{testP}"
pivot_words = [" what", " region"]
df = pd.read_csv(P.DATASETS_DIR / exp_tag / f"{split}.csv", names=C.OFFICIAL_COLUMNS)

# %%
BPE_DIR = P.ROOT / "utils" / "BPE"
tgt_dict = Dictionary.load(str(BPE_DIR / "dict.txt"))
bpe = GPT2BPE(
    GPT2BPEConfig(
        gpt2_encoder_json=BPE_DIR / "encoder.json", gpt2_vocab_bpe=BPE_DIR / "vocab.bpe"
    )
)


# %%
def encode_text(
    text: str,
    bpe: GPT2BPE,
    tgt_dict: Dictionary,
    length: int | None = None,
    append_bos: bool | None = False,
    append_eos: bool | None = False,
    use_bpe: bool | None = True,
) -> str:
    return tgt_dict.encode_line(
        line=bpe.encode(text) if use_bpe else text,
        add_if_not_exist=False,
        append_eos=False,
    ).long()


default_encode_text = partial(
    encode_text,
    bpe=bpe,
    tgt_dict=tgt_dict,
)

# %%
# Load the prediction generated by your model
predict = (
    pd.read_csv(
        P.RESULTS_DIR
        / arch
        / exp_tag
        / f"train-P{trainP}"
        / f"val-P{valP}"
        / f"private_test-P{testP}_predict.csv",
        names=["id", "box"],
    )["box"]
    .map(eval)
    .map(lambda x: np.array(x).astype(int))
)
assert len(predict) == len(df), "answer length not match: {} vs {}".format(
    len(predict), len(df)
)

# merge the prediction into the dataframe
df["predict"] = predict

# format the annotaion answer
df["answer"] = df.apply(
    lambda x: np.array(
        [x["left"], x["top"], x["left"] + x["width"], x["top"] + x["height"]]
    ),
    axis=1,
)
df.drop(["left", "top", "width", "height"], axis=1, inplace=True)

# format the image from url to numpy array
df["pix"] = df["image"].map(
    lambda x: IF.filepath_to_np_array(str(P.IMAGESS_DIR / x.split("/")[-1]))
)

# format the question to the input format
df["question_input"] = (
    df["question"]
    .apply(lambda q: OFADataset.pre_caption(q, max_src_length))
    .map(lambda q: iot.format(question=q))
)

df["question_tokens"] = (
    df["question_input"]
    .apply(lambda q: default_encode_text(text=q))
    .map(
        lambda q: torch.cat(
            [torch.LongTensor([tgt_dict.bos()]), q, torch.LongTensor([tgt_dict.eos()])],
        )
    )
)

# %%

attn = np.load(
    P.RESULTS_DIR
    / "attn"
    / arch
    / exp_tag
    / f"train-P{trainP}"
    / f"val-P{valP}"
    / f"test-P{testP}"
    / f"attn_record-{row_id}.npy",
    allow_pickle=True,
)

# enc_layer, seq_len, batch_size, hidden_size
assert (
    len(attn.shape) == 4
    and attn.shape[0] == C.OFA_ARCH_ENC_LAYERS_MAPPING[arch]
    and attn.shape[-1] == C.OFA_ARCH_HIDDEN_SIZE_MAPPING[arch]
), (
    f"attn shape not match, your: `{attn.shape}` "
    "expected: `("
    f"{C.OFA_ARCH_ENC_LAYERS_MAPPING[arch]}, "
    "seq_len, "
    "batch_size, "
    f"{C.OFA_ARCH_HIDDEN_SIZE_MAPPING[arch]})`"
)
# %%
row = df.iloc[row_id]
# img_attn shape: (1024, hidden_size)
img_attn = attn[attn_layer_id, :1024, batch_id, :]
# txt_attn shape: (seq_len-1024, hidden_size)
txt_attn = attn[attn_layer_id, 1024:, batch_id, :]

assert (
    img_attn.shape[0] == 1024
    and img_attn.shape[-1] == C.OFA_ARCH_HIDDEN_SIZE_MAPPING[arch]
), (
    f"img_attn shape not match, yours: `{img_attn.shape}` "
    f"expected: `(1024, {C.OFA_ARCH_HIDDEN_SIZE_MAPPING[arch]})`"
)
assert len(txt_attn) == len(row["question_tokens"]), (
    f"txt_attn shape not match, yours: `{txt_attn.shape}` "
    f"expected: `({len(row['question_tokens'])}, {C.OFA_ARCH_HIDDEN_SIZE_MAPPING[arch]})`"
)

# %%
pivot_words_pos = []

# %%
for i in range(len(row["question_tokens"])):
    word = bpe.decode(
        tgt_dict.string(row["question_tokens"][None, i])
        .replace("<BOS>", "")
        .replace("<EOS>", "")
    )
    if word in pivot_words:
        pivot_words_pos.append(i)

print(f"pivot_words: {pivot_words}\npivot_words_pos: {pivot_words_pos}")

# %%
box = row["predict"]
pivot_word_attn = txt_attn[pivot_words_pos, ...]
vis_attn = img_attn @ pivot_word_attn.T
vis_attn = vis_attn.squeeze()
# average all pivot words' attention
vis_attn = np.mean(vis_attn, axis=-1)

# %% ================================================================================
# word to image attention
# ================================================================================
fig, (ax2, ax1) = plt.subplots(1, 2, figsize=(9.1, 5))
fig.suptitle(f"{row['question']}")
sns.heatmap(
    vis_attn.reshape(32, 32),
    annot=False,
    fmt="g",
    ax=ax1,
    cbar=False,
    xticklabels=[],
    yticklabels=[],
)
img = Image.fromarray(
    cv2.rectangle(
        np.array(row["pix"]),
        (row["predict"][0], row["predict"][1]),
        (row["predict"][2], row["predict"][3]),
        (255, 0, 0),
        3,
    )
)
ax2.imshow(img.resize((256, 256)))
ax2.set_xticks([])
ax2.set_yticks([])
plt.tight_layout()
# %% ================================================================================


# %% ================================================================================
# patch to sentence attention
# ================================================================================

txt_pox = 18
what = txt[txt_pox, ...]
fig, (ax2, ax1) = plt.subplots(1, 2, figsize=(9.1, 5))
fig.suptitle(f"{otxt}")
vis_attn = vis @ what
# vis_attn = -vis_attn
sns.heatmap(
    vis_attn.reshape(32, 32),
    annot=False,
    fmt="g",
    ax=ax1,
    cbar=False,
    xticklabels=[],
    yticklabels=[],
)
w, h = img.size
# reshape to 32x32
patch = Rectangle(
    (int((res[2] + res[0]) / 2 / w * 32) - 1, int((res[1] + res[3]) / 2 / h * 32) - 2),
    int((res[2] - res[0]) / 2 / w * 32) + 5,
    int((res[1] - res[3]) / 2 / h * 32) + 5,
    edgecolor="lime",
    facecolor="none",
    lw=2,
)
ax1.add_patch(patch)
# img = Image.fromarray(cv2.rectangle(np.array(img), (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2))
img = Image.fromarray(
    cv2.rectangle(np.array(img), (res[0], res[1]), (res[2], res[3]), (255, 0, 0), 3)
)
# img = Image.fromarray(cv2.drawMarker(
#     np.array(img),
#     ((res[2]+res[0])//2, (res[3]+res[1])//2),
#     (255, 0, 0),
#     markerType=cv2.MARKER_CROSS,
#     markerSize=50,
#     thickness=2,
# ))

ax2.imshow(img.resize((256, 256)))
ax2.set_xticks([])
ax2.set_yticks([])
plt.tight_layout()

#########################################
# img to txt
#########################################
que_t = que[7:-5]
txt_t = txt[8 : 8 + len(que_t)]

box_r = (res / max_box * 32).astype(int)

coor = ((box_r[2:] + box_r[:2]) // 2)[::-1].tolist()
coor
vis_attn_tmp = vis_attn.copy()
# vis_attn.reshape(32, 32)[coor[0], coor[1]] = 0

plt.figure(figsize=(5, 5))
sns.heatmap(
    vis_attn.reshape(32, 32),
    annot=False,
    fmt="g",
    cbar=False,
    xticklabels=[],
    yticklabels=[],
)
plt.tight_layout()

vis_attn = vis_attn_tmp

what_t = vis.reshape((32, 32, 1024))[coor[0], coor[1], ...]
txt_attn_t = txt_t @ what_t
txt_attn_t = txt_attn_t / np.sum(txt_attn_t)
plt.figure(figsize=(10, 5))
plt.bar(np.arange(len(txt_attn_t)), txt_attn_t)
plt.ylim(txt_attn_t.min() - 0.001, txt_attn_t.max() + 0.001)
_ = plt.xticks(np.arange(len(txt_attn_t)), que_t, rotation=40)

# %%