{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 50: Visual Question Answering\n",
    "\n",
    "Lesson 49's CLIP matched a *fixed* set of text prompts against an image. **Visual Question Answering (VQA)** asks for something more flexible: given an image and an arbitrary natural-language question about it, produce an answer. This is the mechanism behind modern multimodal assistants that can look at a picture and answer questions about it — a generative-VLM capability, distinct from CLIP's contrastive matching. This lesson builds the classic (and still widely used) simplification: treat VQA as classification over a small, closed vocabulary of possible answers, fusing an image encoder and a question encoder before the final prediction."
   ]
  },
  {
   "cell_type": "code",
   "id": "f80c1b03",
   "source": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "0c4fd57c",
   "source": "## Scenes with (question, answer) pairs\n\nScenes contain 0-3 plus/circle shapes. Each scene gets one of four question templates: \"what shape is in the image\" (only meaningful when exactly one shape is present), \"how many shapes are there\", \"is there a circle in the image\", \"is there a plus in the image\" \u2014 with the true answer computed directly from the scene's ground truth. The answer space is a small, fixed vocabulary: shape names, counts, and yes/no.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "aaa2ce35",
   "source": "SIZE = 32\n\nQUESTION_VOCAB = ['<pad>', 'what', 'shape', 'is', 'this', 'how', 'many', 'shapes', 'are', 'there',\n                   'a', 'circle', 'plus', 'in', 'the', 'image']\nQ2ID = {w: i for i, w in enumerate(QUESTION_VOCAB)}\nANSWER_VOCAB = ['plus', 'circle', '0', '1', '2', '3', 'yes', 'no']\nA2ID = {w: i for i, w in enumerate(ANSWER_VOCAB)}\n\ndef make_shape(img, shape_type, cx, cy, r=4):\n    yy, xx = np.mgrid[0:img.shape[0], 0:img.shape[1]]\n    if shape_type == 'plus':\n        img[cy-1:cy+2, cx-r:cx+r+1] = 1.0\n        img[cy-r:cy+r+1, cx-1:cx+2] = 1.0\n    else:\n        img[((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2] = 1.0\n    return img\n\ndef make_scene(rng, size=SIZE, max_shapes=3):\n    n = rng.integers(0, max_shapes + 1)\n    img = np.zeros((size, size), dtype=np.float32)\n    shapes = []\n    tries = 0\n    while len(shapes) < n and tries < 30:\n        tries += 1\n        cx, cy = rng.integers(6, size - 6), rng.integers(6, size - 6)\n        if any(abs(cx - sx) < 9 and abs(cy - sy) < 9 for sx, sy, _ in shapes):\n            continue\n        st = rng.choice(['plus', 'circle'])\n        make_shape(img, st, cx, cy)\n        shapes.append((cx, cy, st))\n    img = np.clip(img + rng.normal(0, 0.03, img.shape), 0, 1).astype(np.float32)\n    return img, shapes\n\ndef make_qa(rng, shapes):\n    n_circle = sum(1 for _, _, t in shapes if t == 'circle')\n    n_plus = sum(1 for _, _, t in shapes if t == 'plus')\n    qtype = rng.choice(['what_shape', 'count', 'is_there_circle', 'is_there_plus'])\n    if qtype == 'what_shape':\n        words = ['what', 'shape', 'is', 'in', 'the', 'image']\n        answer = shapes[0][2] if len(shapes) == 1 else str(min(len(shapes), 3))\n    elif qtype == 'count':\n        words = ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image']\n        answer = str(len(shapes))\n    elif qtype == 'is_there_circle':\n        words = ['is', 'there', 'a', 'circle', 'in', 'the', 'image']\n        answer = 'yes' if n_circle > 0 else 'no'\n    else:\n        words = ['is', 'there', 'a', 'plus', 'in', 'the', 'image']\n        answer = 'yes' if n_plus > 0 else 'no'\n    return words, [Q2ID[w] for w in words], A2ID[answer]\n\ndef pad_question(ids, max_len=8):\n    return ids + [0] * (max_len - len(ids))  # 0 = <pad>, distinct from every real word\n\nrng = np.random.default_rng(4)\nN = 1000\nimgs, questions, answers, question_words = [], [], [], []\nfor _ in range(N):\n    img, shapes = make_scene(rng)\n    words, q_ids, a_id = make_qa(rng, shapes)\n    imgs.append(img)\n    questions.append(pad_question(q_ids))\n    answers.append(a_id)\n    question_words.append(words)\nimgs = np.array(imgs, dtype=np.float32)\nquestions = np.array(questions, dtype=np.int64)\nanswers = np.array(answers, dtype=np.int64)\n\nsplit = int(0.85 * N)\nXtr, Qtr, Atr = imgs[:split], questions[:split], answers[:split]\nXte, Qte, Ate = imgs[split:], questions[split:], answers[split:]\n\nfig, axes = plt.subplots(1, 4, figsize=(9, 2.5))\nfor ax, im, words, a in zip(axes, imgs[:4], question_words[:4], answers[:4]):\n    ax.imshow(im, cmap='gray'); ax.axis('off')\n    ax.set_title(f'{\" \".join(words)}?\\n-> {ANSWER_VOCAB[a]}', fontsize=7)\nplt.show()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "304afe47",
   "source": "## Fusing an image encoder and a question encoder\n\nBoth encoders map to the same-size vector, same pattern as Lesson 49's CLIP. The difference is what happens next: CLIP compared image and text embeddings with a dot product to check *if they match*. VQA needs to actually *combine* them into one joint representation and classify \u2014 a standard, simple fusion is elementwise multiplication (the image vector and the question vector gate each other), followed by a small classifier head over the fixed answer vocabulary.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "ecddd06d",
   "source": "class ImageEncoder(nn.Module):\n    def __init__(self, out_dim=32):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1))\n        self.proj = nn.Linear(32, out_dim)\n\n    def forward(self, x):\n        return self.proj(self.conv(x).flatten(1))\n\nclass QuestionEncoder(nn.Module):\n    def __init__(self, vocab_size, out_dim=32, embed_dim=16):\n        super().__init__()\n        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n        self.proj = nn.Linear(embed_dim, out_dim)\n\n    def forward(self, ids):\n        mask = (ids != 0).float().unsqueeze(-1)  # ignore <pad> tokens when pooling\n        emb = self.embed(ids) * mask\n        pooled = emb.sum(1) / mask.sum(1).clamp(min=1)\n        return self.proj(pooled)\n\nclass VQAModel(nn.Module):\n    def __init__(self, vocab_size, n_answers, dim=32):\n        super().__init__()\n        self.img_enc = ImageEncoder(dim)\n        self.q_enc = QuestionEncoder(vocab_size, dim)\n        self.head = nn.Sequential(nn.Linear(dim, 32), nn.ReLU(), nn.Linear(32, n_answers))\n\n    def forward(self, img, q_ids):\n        fused = self.img_enc(img) * self.q_enc(q_ids)\n        return self.head(fused)\n\ntorch.manual_seed(0)\nmodel = VQAModel(len(QUESTION_VOCAB), len(ANSWER_VOCAB))\nopt = torch.optim.Adam(model.parameters(), lr=0.01)\nXt = torch.tensor(Xtr).unsqueeze(1); Qt = torch.tensor(Qtr); At = torch.tensor(Atr)\nfor _ in range(300):\n    opt.zero_grad()\n    loss = F.cross_entropy(model(Xt, Qt), At)\n    loss.backward()\n    opt.step()\n\nwith torch.no_grad():\n    preds = model(torch.tensor(Xte).unsqueeze(1), torch.tensor(Qte)).argmax(1).numpy()\nacc = (preds == Ate).mean()\n\nmost_common = np.bincount(Atr).argmax()\nbaseline_acc = (Ate == most_common).mean()\n\nprint(f'VQA test accuracy: {acc:.1%}')\nprint(f'baseline (always predict the most common training answer, \"{ANSWER_VOCAB[most_common]}\"): {baseline_acc:.1%}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "d27a4875",
   "source": "## The real test: does the same image get different answers to different questions?\n\nA model that's secretly ignoring the question and just pattern-matching the image would give the same answer regardless of what's asked. Pose all four question types against one fixed test image and check.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "06acc6c9",
   "source": "test_rng = np.random.default_rng(123)\ndemo_img, demo_shapes = make_scene(test_rng)\nwhile len(demo_shapes) != 2 or demo_shapes[0][2] == demo_shapes[1][2]:\n    demo_img, demo_shapes = make_scene(test_rng)  # find a scene with one of each shape\n\ndemo_questions = [\n    ['how', 'many', 'shapes', 'are', 'there', 'in', 'the', 'image'],\n    ['is', 'there', 'a', 'circle', 'in', 'the', 'image'],\n    ['is', 'there', 'a', 'plus', 'in', 'the', 'image'],\n]\nimg_t = torch.tensor(demo_img[None, None])\nprint(f'true scene contents: {[s[2] for s in demo_shapes]}')\nwith torch.no_grad():\n    for words in demo_questions:\n        q_ids = torch.tensor([pad_question([Q2ID[w] for w in words])])\n        pred = model(img_t, q_ids).argmax(1).item()\n        print(f'  \"{\" \".join(words)}?\" -> {ANSWER_VOCAB[pred]}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "e1b89946",
   "source": "The same fixed image produces three different, individually correct answers depending on what's asked \u2014 the fusion step is genuinely reading both inputs, not just memorizing image-to-answer shortcuts.\n\n## From closed-vocabulary classification to real generative VQA\n\nTreating VQA as classification over a small fixed answer set (as this lesson did) was the dominant approach for years and is still used when the answer space genuinely is small and known in advance. It breaks down the moment an answer needs to be open-ended text \u2014 \"describe what's unusual about this image,\" or an answer word that never appeared in training. Modern VQA and image-captioning systems solve this the way Lesson 45's decoder does: replace the classification head with an autoregressive text decoder (causal self-attention over previously generated words, cross-attention into the image features) that generates the answer one token at a time, exactly the mechanism behind large multimodal assistants that can describe, question, and reason about an image in free-form language rather than picking from a fixed list.\n\n### Exercise\n\n1. Add a fifth question template asking about *position* (e.g. \"is there a shape in the top half of the image\", answer yes/no based on `cy < SIZE // 2`). Retrain and check whether accuracy on this new question type matches the others \u2014 does the image encoder's global max pooling (Lesson 33) make position-based questions structurally harder to answer than presence/count questions?\n2. Replace elementwise multiplication fusion (`img_enc(img) * q_enc(q_ids)`) with concatenation followed by a linear layer (`nn.Linear(2 * dim, dim)` on `torch.cat([img_feat, q_feat], dim=-1)`). Does accuracy change noticeably, and which fusion method converges faster during training?\n3. This lesson's question encoder mean-pools word embeddings, so `\"is there a circle\"` and `\"there a circle is\"` encode identically (Lesson 49's exercise made the same point about captions). Construct a question pair for this dataset where that word-order-blindness would actually cause an answer error, if such a pair exists \u2014 or explain why this dataset's question templates are simple enough that it never matters.",
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}