retrieval_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os, sys\n",
    "import importlib\n",
    "from tqdm.notebook import tqdm, trange\n",
    "import pandas as pd\n",
    "import json\n",
    "import itertools\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### VLM and Parameter Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV = os.environ[\"CONDA_DEFAULT_ENV\"]\n",
    "\n",
    "if ENV == \"videoclip\":\n",
    "    from video_clip.video_clip import VideoClipVLM as VLM\n",
    "    vlm_params = {\n",
    "        \"num_seconds\": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],\n",
    "        \"sample_strat\": [\"center\", \"start\", \"spread\"]\n",
    "    }\n",
    "\n",
    "elif ENV == \"VLM_MILES\":\n",
    "    from MILES.wrapper import MILES_SimilarityVLM as VLM\n",
    "    vlm_params = {}\n",
    "\n",
    "elif ENV == \"VLM_CLIP\":\n",
    "    from CLIP.CLIPVLM import ClipVLM as VLM\n",
    "    vlm_params = {\n",
    "        \"num_frames\": [1, 2, 3, 4, 6, 8, 10, 20, 50, 100]\n",
    "    }\n",
    "\n",
    "else:\n",
    "    ValueError(ENV)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_FOLDER = \"/home/datasets/500p/10s_dataset\"\n",
    "MP4_FOLDER = f\"{DATA_FOLDER}/10s_clips\"\n",
    "TXT_FOLDER = f\"{DATA_FOLDER}/10s_kaldi_texts\"\n",
    "\n",
    "video_names = [name[:-4] for name in os.listdir(MP4_FOLDER) if name.endswith(\".mp4\")]\n",
    "video_names.sort()\n",
    "text_names = [name[:-4] for name in os.listdir(TXT_FOLDER) if name.endswith(\".txt\")]\n",
    "text_names.sort()\n",
    "assert video_names == text_names\n",
    "pair_names = video_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_ID = \"kinetics_100\"\n",
    "\n",
    "if DATASET_ID == \"500p\":\n",
    "    N = len(pair_names)\n",
    "    vid_paths = np.array([f\"{MP4_FOLDER}/{name}.mp4\" for name in pair_names])\n",
    "    vid_text = []\n",
    "    for name in pair_names:\n",
    "        text_path = f\"{TXT_FOLDER}/{name}.txt\"\n",
    "        with open(text_path, \"r\") as fp:\n",
    "            vid_text.append(fp.read().lower().strip())\n",
    "    vid_text = np.array(vid_text)\n",
    "\n",
    "elif DATASET_ID == \"500p.text_overflow_1_1\":\n",
    "    N = len(pair_names)\n",
    "    vid_paths = np.array([f\"{MP4_FOLDER}/{name}.mp4\" for name in pair_names])\n",
    "    raw_vid_text = []\n",
    "    for name in pair_names:\n",
    "        text_path = f\"{TXT_FOLDER}/{name}.txt\"\n",
    "        with open(text_path, \"r\") as fp:\n",
    "            raw_vid_text.append(fp.read().lower().strip())\n",
    "    vid_text = []\n",
    "    for i in range(len(raw_vid_text)):\n",
    "        vid_text.append(\" \".join(raw_vid_text[max(0, i - 1) : min(len(raw_vid_text), i + 2)]))\n",
    "    vid_text = np.array(vid_text)\n",
    "    \n",
    "elif DATASET_ID == \"kinetics_100\":\n",
    "    from dataset import DatasetHandler\n",
    "    dataset = DatasetHandler(\"kinetics_100\", split=\"all\")\n",
    "    N = dataset.category_count()\n",
    "    vid_paths, vid_text = next(iter(dataset.few_shot(1, N, 0, 1)))\n",
    "    vid_paths = vid_paths[:, 0]\n",
    "\n",
    "else:\n",
    "    raise ValueError(DATASET_ID)\n",
    "\n",
    "print(DATASET_ID)\n",
    "print(f\"{N} Video-Text Pairs\")\n",
    "for i in range(N):\n",
    "    print(f\"{i:>3}: {vid_paths[i].split('/')[-1]:>20}\")\n",
    "    print(vid_text[i])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieval Tester Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(\"retrieval_results.csv\"):\n",
    "    results = pd.read_csv(\"retrieval_results.csv\")\n",
    "else:\n",
    "    results = pd.DataFrame(columns=[\"vlm_class\", \"vlm_params\", \"dataset\",\n",
    "                                    \"R@1 (of 82)\", \"R@5 (of 82)\", \"R@10 (of 82)\", \"R@20 (of 82)\", \"R@50 (of 82)\",\n",
    "                                    \"Mean R\", \"Med R\",\n",
    "                                    \"Ranked Text Indices\", \"Vid-Text Choice Ranks\",\n",
    "                                    \"Correct Vid-Text Pair Ranks\", \"Vid-Text Pairs (Easy->Hard)\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieval_test(vlm):    \n",
    "    pbar = trange(2 * N, leave=False)\n",
    "    vid_embeds = []\n",
    "    text_embeds = []\n",
    "    for path in vid_paths:\n",
    "        vid_embeds.append(vlm.get_video_embeds(path))\n",
    "        pbar.update(1)\n",
    "    for text in vid_text:\n",
    "        text_embeds.append(vlm.get_text_embeds(text))\n",
    "        pbar.update(1)\n",
    "    vid_embeds = np.array(vid_embeds)\n",
    "    text_embeds = np.array(text_embeds)\n",
    "    \n",
    "    similarity = vlm.default_similarity_metric()(vid_embeds, text_embeds)\n",
    "    \n",
    "    sorted_text_choice_indices = np.argsort(-similarity, axis=1) # i, j = text index which is the jth best match to video index i\n",
    "    pair_ranks = np.argsort(sorted_text_choice_indices, axis=1) # i, j = rank position of text j for vid i out of all text options (0 = best choice, 81 = worst)\n",
    "    correct_pair_ranks = pair_ranks[np.arange(N), np.arange(N)] # i = rank position of correct pair (vid i - text i) out of all text options (0 = best, 81 = worst)\n",
    "    sorted_pair_indices = np.argsort(correct_pair_ranks) # vid-text pair index with best rank to vid-text pair index with worst rank\n",
    "    \n",
    "    \n",
    "    # R@1\n",
    "    R1_sum = np.sum(correct_pair_ranks < 1)\n",
    "    R1 = R1_sum / N\n",
    "    \n",
    "    # R@5\n",
    "    R5_sum = np.sum(correct_pair_ranks < 5)\n",
    "    R5 = R5_sum / N\n",
    "    \n",
    "    # R@10\n",
    "    R10_sum = np.sum(correct_pair_ranks < 10)\n",
    "    R10 = R10_sum / N\n",
    "    \n",
    "    # R@20\n",
    "    R20_sum = np.sum(correct_pair_ranks < 20)\n",
    "    R20 = R20_sum / N\n",
    "    \n",
    "    # R@50\n",
    "    R50_sum = np.sum(correct_pair_ranks < 50)\n",
    "    R50 = R50_sum / N\n",
    "    \n",
    "    # Mean/Med Rank\n",
    "    mean_R = np.mean(correct_pair_ranks) + 1\n",
    "    med_R = np.median(correct_pair_ranks) + 1\n",
    "    \n",
    "    print(f\"{vlm.__class__.__name__}   {json.dumps(vlm.params())}\")\n",
    "    print(f\"{'R1':>10}{'R5':>10}{'R10':>10}{'R20':>10}{'R50':>10}{'MeanR':>10}{'MedR':>10}\")\n",
    "    print(f\"{R1:10.3f}{R5:10.3f}{R10:10.3f}{R20:10.3f}{R50:10.3f}{mean_R:10.3f}{med_R:10.3f}\")\n",
    "    print()\n",
    "    \n",
    "    \n",
    "    # Save results\n",
    "    # Remove any previously-saved versions of this same test\n",
    "    global results\n",
    "    test_spec = {\n",
    "        \"vlm_class\": vlm.__class__.__name__,\n",
    "        \"vlm_params\": vlm.params(),\n",
    "        \"dataset\": DATASET_ID\n",
    "    }\n",
    "    prev_matching_tests = (results[list(test_spec.keys())] == pd.Series(test_spec)).all(axis=1)\n",
    "    if np.any(prev_matching_tests):\n",
    "        results = results[~prev_matching_tests].reset_index(drop=True)\n",
    "        \n",
    "    results.loc[len(results)] = [\n",
    "        test_spec[\"vlm_class\"],\n",
    "        test_spec[\"vlm_params\"],\n",
    "        test_spec[\"dataset\"],\n",
    "        R1_sum, R5_sum, R10_sum, R20_sum, R50_sum,\n",
    "        mean_R, med_R,\n",
    "        sorted_text_choice_indices, pair_ranks,\n",
    "        correct_pair_ranks, sorted_pair_indices\n",
    "    ]\n",
    "    results.to_csv(\"retrieval_results.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test over all vlm params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dynamically display most recent test results\n",
    "disp = display(display_id=True)\n",
    "disp.update(results.tail(5))\n",
    "\n",
    "vlm = None\n",
    "if len(vlm_params):\n",
    "    param_list = tqdm(list(itertools.product(*vlm_params.values())))\n",
    "else:\n",
    "    param_list = tqdm([[]])\n",
    "    \n",
    "for params in param_list:\n",
    "    params = dict(zip(vlm_params.keys(), params))\n",
    "    param_list.set_postfix(params)\n",
    "    \n",
    "    vlm = VLM(**params)\n",
    "    \n",
    "    retrieval_test(vlm)\n",
    "    disp.update(results.tail(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = []\n",
    "for x in results[\"Correct Vid-Text Pair Ranks\"][[3, 30, 40]].values:\n",
    "    test.append([int(rank) for rank in x.replace(\"[\", \"\").replace(\"]\", \"\").replace(\"\\n\", \"\").split()])\n",
    "test = np.array(test)\n",
    "test = np.mean(test, axis=0)\n",
    "print(test)\n",
    "print(np.argsort(test))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('videoclip')",
   "language": "python",
   "name": "python3"
  },
  "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.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "576597d045da71b18680487735a015f28433d9aa438b3c061008c825d6c37722"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}