notebooks/im2clue-gptj.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb0c959-b9fb-496b-aeb2-7394720f2f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%autosave 60\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e35c9f5-f56b-4a6f-b7af-3f52493fde30",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pickle\n",
    "from collections import Counter, OrderedDict\n",
    "from copy import deepcopy\n",
    "from pathlib import Path\n",
    "from typing import Any, Dict, List, Optional, Tuple, Union, cast\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import PIL.Image as pil_img\n",
    "import seaborn as sns\n",
    "import sklearn as skl\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from IPython.display import Image, display\n",
    "from matplotlib.patches import Rectangle\n",
    "from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
    "from torch.nn import functional as F\n",
    "from tqdm.contrib import tenumerate, tmap, tzip\n",
    "from tqdm.contrib.bells import tqdm, trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8103a868-7e93-49f2-a31a-256bec5cea4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option(\"display.max_colwidth\", None)\n",
    "pd.set_option(\"display.max_columns\", 15)\n",
    "pd.set_option(\"display.max_rows\", 50)\n",
    "# Suitable default display for floats\n",
    "pd.options.display.float_format = \"{:,.2f}\".format\n",
    "plt.rcParams[\"figure.figsize\"] = (12, 10)\n",
    "\n",
    "# This one is optional -- change graphs to SVG only use if you don't have a\n",
    "# lot of points/lines in your graphs. Can also just use ['retina'] if you\n",
    "# don't want SVG.\n",
    "%config InlineBackend.figure_formats = [\"retina\"]\n",
    "set_matplotlib_formats(\"pdf\", \"png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f308e74c-7696-4e46-b81c-f0d7200445c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sims(a: torch.tensor, b: torch.tensor, batch_size=1000):\n",
    "    sims = {i: (0, -1) for i in range(a.shape[0])}\n",
    "    for j in tqdm(range(0, b.shape[0], batch_size)):\n",
    "        batch_sims = a @ b[j : j + batch_size].T\n",
    "        batch_sims = torch.from_numpy(batch_sims)\n",
    "        values, idxs = batch_sims.max(dim=-1)\n",
    "        # Append an index and sim every iteration\n",
    "        for i in range(a.shape[0]):\n",
    "            offset = j + idxs[i].item()\n",
    "            if values[i] > sims[i][0]:\n",
    "                sims[i] = (values[i].item(), offset)\n",
    "    return sims\n",
    "\n",
    "\n",
    "def normalize_rows(mat: torch.tensor) -> None:\n",
    "    for i in range(len(mat)):\n",
    "        mat[i] /= mat[i].norm(p=2, dim=-1, keepdim=True)\n",
    "    return mat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87e91062-cac0-4b28-a028-3958b94c29a8",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Geoguessr In-game Frames -> GPT-J Text Lookup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85eaa884-1d66-438e-b24f-8d5ee3a74d71",
   "metadata": {},
   "outputs": [],
   "source": [
    "split = \"test\"\n",
    "# Load query image embeddings:\n",
    "image_embs = pickle.load(\n",
    "    open(\n",
    "        f\"/shared/gbiamby/geo/models/clip_ft/vit-b32/geoframes_clip_samples_fixed_window_{split}_img.pkl\",\n",
    "        \"rb\",\n",
    "    )\n",
    ")\n",
    "text_embs = pickle.load(\n",
    "    open(f\"/shared/gbiamby/geo/models/clip_ft/vit-b32/gptj_clues_text.pkl\", \"rb\")\n",
    ")\n",
    "# load target captions:\n",
    "gpt_caps = json.load(open(\"/shared/g-luo/geoguessr/data/data/guidebook/kb/v3/cleaned_clues.json\"))[\"clues\"]\n",
    "\n",
    "# # Append embeddings\n",
    "# for caption in gpt_caps:\n",
    "#     caption[\"clip_emb\"] = image_embs[caption[\"file_path\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f5c79ec-7ea4-4b6f-992b-c7bafb4e4aba",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_image_embs = normalize_rows(\n",
    "    torch.stack([torch.tensor(emb) for emb in image_embs.values()]).to(\"cuda\")\n",
    ")\n",
    "t_text_embs = normalize_rows(\n",
    "    torch.stack([torch.tensor(emb) for emb in text_embs.values()]).to(\"cuda\")\n",
    ")\n",
    "\n",
    "print(f\"image_embs.shape: {t_image_embs.shape}, text_emb.shape: {t_text_embs.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a5acdff-0b50-4a1b-8d4f-84bd140b9492",
   "metadata": {},
   "outputs": [],
   "source": [
    "sims = torch.mm(t_image_embs, t_text_embs.T)\n",
    "max_sim_scores, max_sim_idxs = sims.max(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4fae438-bac1-43e1-b604-d218aef0605c",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts_unique = list(text_embs.keys())\n",
    "img_to_text_sims = {}\n",
    "for i, img_path in enumerate(image_embs.keys()):\n",
    "    img_to_text_sims[img_path] = {\n",
    "        \"best_match_text\": texts_unique[max_sim_idxs[i]],\n",
    "        \"best_sim_score\": max_sim_scores[i].tolist(),\n",
    "        \"best_sim_idx\": max_sim_idxs[i].tolist(),\n",
    "        \"sims_all\": sims[i].tolist(),\n",
    "        \"file_path\": img_path,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d56c187b-1cc4-409e-83ef-b5dfb61601d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(t_image_embs.shape, t_text_embs.shape, sims.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a91a7da5-95fe-428f-81aa-4b9412c910e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# list(img_to_text_sims.items())[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f19d72d-7405-452a-acfb-dad923bf082b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import HTML, Markdown\n",
    "\n",
    "\n",
    "def show_samples(img_to_text_sims, n_samples: int = 20):\n",
    "    df_random = np.random.choice(range(len(img_to_text_sims)), n_samples, replace=False)\n",
    "\n",
    "    for i in df_random:\n",
    "        img_row = deepcopy(img_to_text_sims[i])\n",
    "        del img_row[\"sims_all\"]\n",
    "        print(\"=\" * 180)\n",
    "        # print(img_row.keys())\n",
    "        # print(img_row)\n",
    "        display(pd.DataFrame({k: [v] for k, v in img_row.items()}).T)\n",
    "        img = pil_img.open(img_row[\"file_path\"])\n",
    "        img.thumbnail((1080, 640), pil_img.NEAREST)\n",
    "        display(img)\n",
    "        print(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec14fe6c-3ac3-406e-bf10-11439af45bd4",
   "metadata": {},
   "source": [
    "### Choose Images w/ Highest Img/Text Similarity Scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f157ce-028b-478f-9c5d-9c47756b74a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "show_samples(\n",
    "    sorted(img_to_text_sims.values(), key=lambda x: x[\"best_sim_score\"], reverse=True)[:100],\n",
    "    n_samples=20,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19ca1cd1-c4f4-4e69-b59a-d7a5738ce08e",
   "metadata": {},
   "source": [
    "### Choose Random Images, show best match"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "187abe47-dba7-493e-9875-8b62136b09fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_samples(list(img_to_text_sims.values()), n_samples=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0597e40c-00ee-4cc3-9703-3aa37caf5695",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geoscreens",
   "language": "python",
   "name": "geoscreens"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}