Lesson 50: Visual Question Answering¶

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

SIZE = 32

QUESTION_VOCAB = ['<pad>', 'what', 'shape', 'is', 'this', 'how', 'many', 'shapes', 'are', 'there',
                   'a', 'circle', 'plus', 'in', 'the', 'image']
Q2ID = {w: i for i, w in enumerate(QUESTION_VOCAB)}
ANSWER_VOCAB = ['plus', 'circle', '0', '1', '2', '3', 'yes', 'no']
A2ID = {w: i for i, w in enumerate(ANSWER_VOCAB)}

def make_shape(img, shape_type, cx, cy, r=4):
    yy, xx = np.mgrid[0:img.shape[0], 0:img.shape[1]]
    if shape_type == 'plus':
        img[cy-1:cy+2, cx-r:cx+r+1] = 1.0
        img[cy-r:cy+r+1, cx-1:cx+2] = 1.0
    else:
        img[((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2] = 1.0
    return img

def make_scene(rng, size=SIZE, max_shapes=3):
    n = rng.integers(0, max_shapes + 1)
    img = np.zeros((size, size), dtype=np.float32)
    shapes = []
    tries = 0
    while len(shapes) < n and tries < 30:
        tries += 1
        cx, cy = rng.integers(6, size - 6), rng.integers(6, size - 6)
        if any(abs(cx - sx) < 9 and abs(cy - sy) < 9 for sx, sy, _ in shapes):
            continue
        st = rng.choice(['plus', 'circle'])
        make_shape(img, st, cx, cy)
        shapes.append((cx, cy, st))
    img = np.clip(img + rng.normal(0, 0.03, img.shape), 0, 1).astype(np.float32)
    return img, shapes

def make_qa(rng, shapes):
    n_circle = sum(1 for _, _, t in shapes if t == 'circle')
    n_plus = sum(1 for _, _, t in shapes if t == 'plus')
    qtype = rng.choice(['what_shape', 'count', 'is_there_circle', 'is_there_plus'])
    if qtype == 'what_shape':
        words = ['what', 'shape', 'is', 'in', 'the', 'image']
        answer = shapes[0][2] if len(shapes) == 1 else str(min(len(shapes), 3))
    elif qtype == 'count':
        words = ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image']
        answer = str(len(shapes))
    elif qtype == 'is_there_circle':
        words = ['is', 'there', 'a', 'circle', 'in', 'the', 'image']
        answer = 'yes' if n_circle > 0 else 'no'
    else:
        words = ['is', 'there', 'a', 'plus', 'in', 'the', 'image']
        answer = 'yes' if n_plus > 0 else 'no'
    return words, [Q2ID[w] for w in words], A2ID[answer]

def pad_question(ids, max_len=8):
    return ids + [0] * (max_len - len(ids))  # 0 = <pad>, distinct from every real word

rng = np.random.default_rng(4)
N = 1000
imgs, questions, answers, question_words = [], [], [], []
for _ in range(N):
    img, shapes = make_scene(rng)
    words, q_ids, a_id = make_qa(rng, shapes)
    imgs.append(img)
    questions.append(pad_question(q_ids))
    answers.append(a_id)
    question_words.append(words)
imgs = np.array(imgs, dtype=np.float32)
questions = np.array(questions, dtype=np.int64)
answers = np.array(answers, dtype=np.int64)

split = int(0.85 * N)
Xtr, Qtr, Atr = imgs[:split], questions[:split], answers[:split]
Xte, Qte, Ate = imgs[split:], questions[split:], answers[split:]

fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))
for ax, im, words, a in zip(axes, imgs[:4], question_words[:4], answers[:4]):
    ax.imshow(im, cmap='gray'); ax.axis('off')
    ax.set_title(f'{" ".join(words)}?\n-> {ANSWER_VOCAB[a]}', fontsize=7)
plt.show()

class ImageEncoder(nn.Module):
    def __init__(self, out_dim=32):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1))
        self.proj = nn.Linear(32, out_dim)

    def forward(self, x):
        return self.proj(self.conv(x).flatten(1))

class QuestionEncoder(nn.Module):
    def __init__(self, vocab_size, out_dim=32, embed_dim=16):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.proj = nn.Linear(embed_dim, out_dim)

    def forward(self, ids):
        mask = (ids != 0).float().unsqueeze(-1)  # ignore <pad> tokens when pooling
        emb = self.embed(ids) * mask
        pooled = emb.sum(1) / mask.sum(1).clamp(min=1)
        return self.proj(pooled)

class VQAModel(nn.Module):
    def __init__(self, vocab_size, n_answers, dim=32):
        super().__init__()
        self.img_enc = ImageEncoder(dim)
        self.q_enc = QuestionEncoder(vocab_size, dim)
        self.head = nn.Sequential(nn.Linear(dim, 32), nn.ReLU(), nn.Linear(32, n_answers))

    def forward(self, img, q_ids):
        fused = self.img_enc(img) * self.q_enc(q_ids)
        return self.head(fused)

torch.manual_seed(0)
model = VQAModel(len(QUESTION_VOCAB), len(ANSWER_VOCAB))
opt = torch.optim.Adam(model.parameters(), lr=0.01)
Xt = torch.tensor(Xtr).unsqueeze(1); Qt = torch.tensor(Qtr); At = torch.tensor(Atr)
for _ in range(300):
    opt.zero_grad()
    loss = F.cross_entropy(model(Xt, Qt), At)
    loss.backward()
    opt.step()

with torch.no_grad():
    preds = model(torch.tensor(Xte).unsqueeze(1), torch.tensor(Qte)).argmax(1).numpy()
acc = (preds == Ate).mean()

most_common = np.bincount(Atr).argmax()
baseline_acc = (Ate == most_common).mean()

print(f'VQA test accuracy: {acc:.1%}')
print(f'baseline (always predict the most common training answer, "{ANSWER_VOCAB[most_common]}"): {baseline_acc:.1%}')

test_rng = np.random.default_rng(123)
demo_img, demo_shapes = make_scene(test_rng)
while len(demo_shapes) != 2 or demo_shapes[0][2] == demo_shapes[1][2]:
    demo_img, demo_shapes = make_scene(test_rng)  # find a scene with one of each shape

demo_questions = [
    ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image'],
    ['is', 'there', 'a', 'circle', 'in', 'the', 'image'],
    ['is', 'there', 'a', 'plus', 'in', 'the', 'image'],
]
img_t = torch.tensor(demo_img[None, None])
print(f'true scene contents: {[s[2] for s in demo_shapes]}')
with torch.no_grad():
    for words in demo_questions:
        q_ids = torch.tensor([pad_question([Q2ID[w] for w in words])])
        pred = model(img_t, q_ids).argmax(1).item()
        print(f'  "{" ".join(words)}?" -> {ANSWER_VOCAB[pred]}')