{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f7f10b6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']\n"
     ]
    }
   ],
   "source": [
    "import os \n",
    "os.environ[\"OPENBLAS_NUM_THREADS\"] = \"10\"\n",
    "\n",
    "import json \n",
    "import cv2\n",
    "import re\n",
    "import random\n",
    "import time\n",
    "from tqdm.notebook import tqdm\n",
    "import pandas as pd \n",
    "import numpy as np \n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from collections import OrderedDict\n",
    "from pattern.text.en import singularize, pluralize\n",
    "\n",
    "import torch\n",
    "import clip\n",
    "from PIL import Image\n",
    "\n",
    "import nengo.spa as spa\n",
    "from utils import *\n",
    "from dataset import GQADataset\n",
    "\n",
    "RANDOM_SEED = 17\n",
    "\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "random.seed(RANDOM_SEED)\n",
    "\n",
    "DATA_PATH = '/scratch/penzkofer/GQA'\n",
    "    \n",
    "print(clip.available_models())\n",
    "\n",
    "CUDA_DEVICE = 7\n",
    "torch.cuda.set_device(CUDA_DEVICE)\n",
    "\n",
    "device = torch.device(\"cuda:\" + str(CUDA_DEVICE))\n",
    "clip_model, preprocess = clip.load(\"ViT-B/32\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6c6c9d4c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of vector space: 10000x100x512\n",
      "x-axis resolution = 100, y-axis resolution = 100\n",
      "width resolution = 10, height resolution = 10\n",
      "(100, 100, 10, 10, 512)\n",
      "CPU times: user 7min 11s, sys: 2.2 s, total: 7min 13s\n",
      "Wall time: 7min 13s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# define SSP vector space \n",
    "res = 100\n",
    "dim = 1024\n",
    "new_size = (25, 25)  # size should be smaller than resolution\n",
    "\n",
    "xs = np.linspace(0, new_size[1], res)\n",
    "ys = np.linspace(0, new_size[0], res)\n",
    "ws = np.linspace(1, 10, 10)\n",
    "hs = np.linspace(1, 10, 10)\n",
    "\n",
    "rng = np.random.RandomState(seed=RANDOM_SEED)\n",
    "x_axis = make_good_unitary(dim, rng=rng)\n",
    "y_axis = make_good_unitary(dim, rng=rng)\n",
    "w_axis = make_good_unitary(dim, rng=rng)\n",
    "h_axis = make_good_unitary(dim, rng=rng)\n",
    "\n",
    "print(f'Size of vector space: {res**2}x{10**2}x{dim}')\n",
    "print(f'x-axis resolution = {len(xs)}, y-axis resolution = {len(ys)}')\n",
    "print(f'width resolution = {len(ws)}, height resolution = {len(hs)}')\n",
    "\n",
    "# precompute all vectors -- this will take some time \n",
    "VECTORS = get_heatmap_vectors_multidim(xs, ys, ws, hs, x_axis, y_axis, w_axis, h_axis)\n",
    "print(VECTORS.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d6ef10a",
   "metadata": {},
   "source": [
    "### Image to SSP memory "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "dde1bbd1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "943000\n",
      "0 02930152 Is the sky dark? 2354786\n",
      "1 07333408 What is on the white wall? 2375429\n",
      "2 07333405 Is that pipe red? 2375429\n",
      "3 15736264 Is the tall clock small or large? 2368326\n",
      "4 111007521 Who is wearing a shirt? 2331819\n"
     ]
    }
   ],
   "source": [
    "TEST = False\n",
    "\n",
    "# load questions, programs and scenegraphs\n",
    "if TEST: \n",
    "    questions_path = 'val_balanced_questions.json'\n",
    "    programs_path = 'programs/trainval_balanced_programs.json'\n",
    "    scene_path = 'val_sceneGraphs.json'\n",
    "else: \n",
    "    questions_path = 'train_balanced_questions.json'\n",
    "    programs_path = 'programs/trainval_balanced_programs.json'\n",
    "    scene_path = 'train_sceneGraphs.json'\n",
    "\n",
    "with open(os.path.join(DATA_PATH, questions_path), 'r') as f:\n",
    "    questions = json.load(f)\n",
    "\n",
    "with open(os.path.join(DATA_PATH, programs_path), 'r') as f:\n",
    "    programs = json.load(f)\n",
    "\n",
    "with open(os.path.join(DATA_PATH, scene_path), 'r') as f:\n",
    "    scenegraphs = json.load(f)\n",
    "\n",
    "\n",
    "columns = ['semantic', 'entailed', 'equivalent', 'question', 'imageId', 'isBalanced', 'groups', \n",
    "                   'answer', 'semanticStr', 'annotations', 'types', 'fullAnswer']\n",
    "\n",
    "questions = pd.DataFrame.from_dict(questions, orient='index', columns=columns)\n",
    "questions = questions.reset_index()\n",
    "questions = questions.rename(columns={\"index\": \"questionID\"}, errors=\"raise\")\n",
    "            \n",
    "columns = ['imageID', 'question', 'program', 'questionID', 'answer']\n",
    "programs = pd.DataFrame(programs, columns=columns)\n",
    "\n",
    "DATA = GQADataset(questions, programs, scenegraphs, vectors=VECTORS, path=os.path.join(DATA_PATH, 'images/images/'),\n",
    "                  axes=[x_axis, y_axis, w_axis, h_axis], linspace=[xs, ys, ws, hs], visualize=True)\n",
    "\n",
    "print(f'Length of {\"test\" if TEST else \"train\"} data set: ', len(DATA))\n",
    "\n",
    "for i in range(5):\n",
    "    print(i, questions.iloc[i].questionID, questions.iloc[i].question, questions.iloc[i].imageId)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "62842206",
   "metadata": {},
   "outputs": [],
   "source": [
    "from run_programs import *\n",
    "\n",
    "def run_program(data, img, info, counter, memory, dim, verbose=0):\n",
    "    \"\"\" run program for question on given image: \n",
    "    for each step in program select appropriate function\n",
    "    \"\"\"\n",
    "    scale, w_scale, h_scale = info['scales']\n",
    "    img_size = img.shape[:2]\n",
    "\n",
    "    results = []\n",
    "    last_step = False\n",
    "    last_func = None\n",
    "    \n",
    "    for i, step in enumerate(info['program']): \n",
    "        \n",
    "        if i+1 == len(info['program']): \n",
    "            last_step = True\n",
    "    \n",
    "        _, func = step.split('=')\n",
    "        attr = func.split('(')[-1].split(')')[0].split(',')\n",
    "        \n",
    "        if verbose > 0: \n",
    "            print(f'{i+1}. step: \\t {func}')\n",
    "        \n",
    "        if 'select' in func: \n",
    "            obj = attr[0].strip()\n",
    "            res = select_func(data, img, obj, info, counter, memory, visualize=VISUALIZE)\n",
    "            \n",
    "            results.append(res)\n",
    "            \n",
    "            if res is None:\n",
    "                if verbose > 1: \n",
    "                    print(f'Could not find {obj}')\n",
    "                \n",
    "        \n",
    "        elif 'relate' in func: \n",
    "        \n",
    "            found_rel_obj = relate_func(data, func, attr, results, info, dim, memory, visualize=VISUALIZE)\n",
    "\n",
    "            if found_rel_obj is not None:\n",
    "                assert found_rel_obj in info['encoded_ssps'], f'Result of {func}: {found_rel_obj} is not encoded'\n",
    "                \n",
    "                selected_ssp = info['encoded_ssps'][found_rel_obj]     \n",
    "                _, selected_pos = select_ssp(found_rel_obj, memory, info['encoded_ssps'], data.ssp_vectors, data.linspace)\n",
    "                results.append([found_rel_obj, selected_ssp, selected_pos])\n",
    "\n",
    "                if last_step:\n",
    "                    return 'yes' \n",
    "            else:\n",
    "                results.append(None)\n",
    "                \n",
    "                if last_step:\n",
    "                    return 'no'  \n",
    "                \n",
    "        elif 'filter' in func: \n",
    "            \n",
    "            last_filter = filter_func(func, img, attr, img_size, results, info, visualize=VISUALIZE)\n",
    "\n",
    "            if last_filter: \n",
    "                results.append(results[-1])\n",
    "            else:\n",
    "                if results[-1] is None: \n",
    "                    results.append(None)\n",
    "                elif results[-1][0].split(\"_\")[0] + \"_\" + str(counter+1) in info['encoded_ssps'].keys():\n",
    "                    counter += 1\n",
    "                    return None\n",
    "                else: \n",
    "                    last_filter = False\n",
    "                    results.append(results[-1])\n",
    "                    \n",
    "                \n",
    "        elif 'verify' in func: \n",
    "            pred = verify_func(data, img, attr, results, info, dim, memory, verbose=verbose, visualize=VISUALIZE)\n",
    "            \n",
    "            if 'verify_relation_name' in func or 'verify_relation_inv_name' in func: \n",
    "                results.append(results[-1] if pred else None)\n",
    "            else:\n",
    "                results.append(pred)\n",
    "\n",
    "            if last_step:\n",
    "                return 'yes' if pred else 'no'\n",
    "                \n",
    "            \n",
    "        elif 'query' in func: \n",
    "\n",
    "            return query_func(func, img, attr, results, info, img_size, dim, verbose=verbose, visualize=VISUALIZE)\n",
    "        \n",
    "        elif 'exist' in func: \n",
    "\n",
    "            num = int(re.findall(r'\\d+', attr[0])[0])\n",
    "\n",
    "            if last_step:\n",
    "                return 'yes' if results[num] is not None else 'no'\n",
    "            else:\n",
    "                if results[num] is not None and 'filter' not in last_func: \n",
    "                    results.append(True)\n",
    "                elif results[num] is not None and last_filter:\n",
    "                    results.append(True)\n",
    "                else:\n",
    "                    results.append(False)\n",
    "\n",
    "        elif 'or(' in func: \n",
    "            attr1 = int(re.findall(r'\\d+', attr[0])[0])\n",
    "            attr2 = int(re.findall(r'\\d+', attr[1])[0])\n",
    "            return 'yes' if results[attr1] or results[attr2] else 'no'\n",
    "\n",
    "        elif 'and(' in func: \n",
    "            attr1 = int(re.findall(r'\\d+', attr[0])[0])\n",
    "            attr2 = int(re.findall(r'\\d+', attr[1])[0])\n",
    "            return 'yes' if results[attr1] and results[attr2] else 'no'\n",
    "        \n",
    "        elif 'different' in func: \n",
    "            \n",
    "            if len(attr) == 1:\n",
    "                print('[WARNING]'+f'{func} cannot be computed')\n",
    "                return None \n",
    "            \n",
    "            else:    \n",
    "                pred_attr1 = query_func(f'query_{attr[2].strip()}', img, [attr[0], attr[2]], results, info, img_size, dim)\n",
    "                pred_attr2 = query_func(f'query_{attr[2].strip()}', img, [attr[1], attr[2]], results, info, img_size, dim)\n",
    "\n",
    "                if pred_attr1 != pred_attr2:\n",
    "                    return 'yes'\n",
    "                else:\n",
    "                    return 'no'\n",
    "            \n",
    "        elif 'same' in func: \n",
    "            \n",
    "            if len(attr) == 1:\n",
    "                print('[WARNING]'+f'{func} cannot be computed')\n",
    "                return None \n",
    "            \n",
    "            pred_attr1 = query_func(f'query_{attr[2].strip()}', img, [attr[0], attr[2]], results, info, img_size, dim)\n",
    "            pred_attr2 = query_func(f'query_{attr[2].strip()}', img, [attr[1], attr[2]], results, info, img_size, dim)\n",
    "            \n",
    "            if pred_attr1 == pred_attr2:\n",
    "                return 'yes'\n",
    "            else:\n",
    "                return 'no'\n",
    "        \n",
    "        elif 'choose' in func: \n",
    "            return choose_func(data, img, func, attr, img_size, results, info, dim, memory, visualize=VISUALIZE)\n",
    "            \n",
    "        else: \n",
    "            print('[WARNING]'+f'{func} not implemented')\n",
    "            return -1\n",
    "        \n",
    "        last_func = func\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "40084203",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "425457\n",
      "Question #16402843: \n",
      "Are there any knives or glasses that are black?\n",
      "[yes] Yes, the knife is black.\n",
      "\n",
      "Program:\n",
      "0. select(knife)\n",
      "1. filter([0], black)\n",
      "2. exist([1])\n",
      "3. select(glass)\n",
      "4. filter([3], black)\n",
      "5. exist([4])\n",
      "6. or([2],[5])\n",
      "\n",
      "Average mean-squared error of 2D locations: 41.5312\n",
      "Average IoU of 4D bounding boxes: 0.67\n",
      "Correct items: 12 / 17\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 442720 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. step: \t select(knife)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 426676 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2. step: \t filter([0], black)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 426696 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CLIP\n",
      "[0.5888671875, 0.4111328125]\n",
      "The knife is  black\n",
      "3. step: \t exist([1])\n",
      "4. step: \t select(glass)\n",
      "Could not find glass\n",
      "5. step: \t filter([3], black)\n",
      "6. step: \t exist([4])\n",
      "7. step: \t or([2],[5])\n",
      "\n",
      "Question: Are there any knives or glasses that are black?\n",
      "Answer: yes\n",
      "Answer \"yes\" is correct\n",
      "\n",
      "🎊 🎊 🎊 \n",
      "\n",
      "Took 1.90 seconds\n"
     ]
    }
   ],
   "source": [
    "IDX = np.random.randint(len(DATA))\n",
    "VISUALIZE = True\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "print(IDX)\n",
    "DATA.print_item(IDX)\n",
    "DATA.set_visualize(False)\n",
    "DATA.set_verbose(1)\n",
    "img, info, memory = DATA.encode_item(IDX, dim=dim)\n",
    "avg_mse, avg_iou, correct_items = DATA.decode_item(img, info, memory)\n",
    "print()\n",
    "\n",
    "img_size = np.array(img.shape[:2])\n",
    "fig, ax = plt.subplots(1,1)\n",
    "\n",
    "ax.imshow(img, interpolation='none', origin='upper', vmin=0, vmax=1,\n",
    "          extent=[0, img_size[1] / info['scales'][0], img_size[0] / info['scales'][0], 0], cmap='gray')\n",
    "\n",
    "for i, obj in enumerate(info['encoded_items'].items()):\n",
    "    name, [x, y, width, height] = obj\n",
    "    name = name.split('_')[0]\n",
    "        \n",
    "    width, height = (width * info['scales'][1]) / info['scales'][0], (height * info['scales'][2]) / info['scales'][0] \n",
    "    rect = patches.Rectangle((x, y),\n",
    "                             width, height,\n",
    "                             linewidth = 2,\n",
    "                             label = name,\n",
    "                             edgecolor = RGB_COLORS[i],\n",
    "                             facecolor = 'none')\n",
    "    ax.add_patch(rect)\n",
    "    plt.text(x+0.2, y+1, name, fontsize='medium', color=RGB_COLORS[i])\n",
    "\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "answer = run_program(DATA, img, info, counter=1, memory=memory, dim=dim, verbose=2)\n",
    "\n",
    "time_in_sec = time.time() - start\n",
    "correct = answer == info[\"answer\"]\n",
    "\n",
    "print()\n",
    "print(f'Question: {info[\"question\"]}')\n",
    "print(f'Answer: {answer}')\n",
    "print(f'Answer \"{answer}\" is {\"correct\" if correct else \"incorrect\"}\\n')\n",
    "if correct:\n",
    "    print('\\U0001F38A \\U0001F38A \\U0001F38A \\n')\n",
    "    \n",
    "print(f'Took {time_in_sec:.2f} seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "187f1a83",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ssp_env",
   "language": "python",
   "name": "ssp_env"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
