VIDEO_CLIP_EVALUATION.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "01be7f60-1cd4-4afb-8829-342980047596",
   "metadata": {
    "id": "01be7f60-1cd4-4afb-8829-342980047596"
   },
   "source": [
    "### Video prompt evaluation\n",
    "\n",
    "In the second block, choose the clip model, prompt(s) and video file to be used. Check out the functions below to evaluate the CLIP similarity and generate a resulting gif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2ee8f09a-40c5-4e52-bb46-75dc38d96a05",
   "metadata": {
    "id": "2ee8f09a-40c5-4e52-bb46-75dc38d96a05"
   },
   "outputs": [],
   "source": [
    "import open_clip\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "import cv2\n",
    "from PIL import Image\n",
    "import torch\n",
    "import gif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "_g7_iRMTHA_h",
   "metadata": {
    "id": "_g7_iRMTHA_h"
   },
   "outputs": [],
   "source": [
    "MODEL = 'RN50'\n",
    "PRETRAINED = 'yfcc15m'\n",
    "\n",
    "PROMPT = 'homemade rocket flying towards the sky'\n",
    "PROMPTS = ['rocket taking off', 'homemade science experiment']\n",
    "FILENAME = './to the moon.mp4'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ghvm7jYCG9om",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 190
    },
    "id": "ghvm7jYCG9om",
    "outputId": "7c1a7529-1541-4147-deb1-7c0ad21c17ee"
   },
   "outputs": [],
   "source": [
    "model, _, preprocess = open_clip.create_model_and_transforms(MODEL, pretrained=PRETRAINED)\n",
    "model = model.cuda()\n",
    "tokenizer = open_clip.get_tokenizer(MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "sECKnL1OGSAx",
   "metadata": {
    "id": "sECKnL1OGSAx"
   },
   "outputs": [],
   "source": [
    "def run_video_through_clip(filename, prompt, model, preprocess, tokenizer):\n",
    "  text_features = model.encode_text(tokenizer([prompt]).cuda())\n",
    "  text_features /= text_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "  clip_rewards = []\n",
    "  frames = []\n",
    "  vidcap = cv2.VideoCapture(filename)\n",
    "  success, image = vidcap.read()\n",
    "  while success:\n",
    "    frames.append(image)\n",
    "\n",
    "    image = preprocess(Image.fromarray(np.uint8(image))).unsqueeze(0)\n",
    "    with torch.no_grad(), torch.cuda.amp.autocast():\n",
    "      image = image.cuda()\n",
    "      image_features = model.encode_image(image)\n",
    "      image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "      sim = (image_features @ text_features.T)\n",
    "      reward = sim.cpu().detach().numpy()[0][0]\n",
    "    \n",
    "    clip_rewards.append(reward)\n",
    "    success, image = vidcap.read()\n",
    "    \n",
    "  return np.array(frames), np.array(clip_rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "797ef82b",
   "metadata": {},
   "outputs": [],
   "source": [
    "frames, clip_rewards = run_video_through_clip(FILENAME, PROMPT, model, preprocess, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "fc57ffdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_video_through_clip_prompt_reward_averaging(filename, prompts, model, preprocess, tokenizer):\n",
    "  text_feature_list = []\n",
    "  clip_rewards = []\n",
    "  for text in prompts:\n",
    "    text_features = model.encode_text(tokenizer([text]).cuda())\n",
    "    text_features /= text_features.norm(dim=-1, keepdim=True)\n",
    "    text_feature_list.append(text_features)\n",
    "    clip_rewards.append([])\n",
    "    \n",
    "  frames = []\n",
    "  vidcap = cv2.VideoCapture(filename)\n",
    "  success, image = vidcap.read()\n",
    "  while success:\n",
    "    frames.append(image)\n",
    "\n",
    "    image = preprocess(Image.fromarray(np.uint8(image))).unsqueeze(0)\n",
    "    with torch.no_grad(), torch.cuda.amp.autocast():\n",
    "      image = image.cuda()\n",
    "      image_features = model.encode_image(image)\n",
    "      image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "      for j, text_features in enumerate(text_feature_list):\n",
    "          sim = (image_features @ text_features.T)\n",
    "          reward = sim.cpu().detach().numpy()[0][0]\n",
    "          clip_rewards[j].append(reward)\n",
    "    \n",
    "    success, image = vidcap.read()\n",
    "\n",
    "  clip_rewards = np.array(clip_rewards)\n",
    "  return np.array(frames), clip_rewards, np.average(clip_rewards, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "27be4f95",
   "metadata": {},
   "outputs": [],
   "source": [
    "frames, clip_rewards, reward_average = run_video_through_clip_prompt_averaging(FILENAME, PROMPTS, model, preprocess, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "17ed9945",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_video_through_clip_prompt_enconding_averaging(filename, prompts, model, preprocess, tokenizer):\n",
    "  text_feature_list = torch.empty((len(prompts),1, 1024)).cuda()\n",
    "  clip_rewards = []\n",
    "  for i, text in enumerate(prompts):\n",
    "    text_features = model.encode_text(tokenizer([text]).cuda())\n",
    "    text_features /= text_features.norm(dim=-1, keepdim=True)\n",
    "    text_feature_list[i] = text_features\n",
    "  text_features = text_feature_list.mean(dim=0)\n",
    "    \n",
    "  frames = []\n",
    "  vidcap = cv2.VideoCapture(filename)\n",
    "  success, image = vidcap.read()\n",
    "  while success:\n",
    "    frames.append(image)\n",
    "\n",
    "    image = preprocess(Image.fromarray(np.uint8(image))).unsqueeze(0)\n",
    "    with torch.no_grad(), torch.cuda.amp.autocast():\n",
    "      image = image.cuda()\n",
    "      image_features = model.encode_image(image)\n",
    "      image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "      sim = (image_features @ text_features.T)\n",
    "      reward = sim.cpu().detach().numpy()[0][0]\n",
    "      clip_rewards.append(reward)\n",
    "    \n",
    "    success, image = vidcap.read()\n",
    "\n",
    "  return np.array(frames), np.array(clip_rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "0498d6e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "frames, clip_rewards = run_video_through_clip_prompt_enconding_averaging(FILENAME, PROMPTS, model, preprocess, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "LZB2nuSJFpZF",
   "metadata": {
    "id": "LZB2nuSJFpZF"
   },
   "outputs": [],
   "source": [
    "@gif.frame\n",
    "def plot_rewards_gif(i, frame, clip_rewards):\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(20, 10))\n",
    "    ax[0].plot(np.arange(0, i), clip_rewards[:i])\n",
    "    ax[0].set_title('CLIP reward')\n",
    "    ax[1].imshow(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "68c2aa45",
   "metadata": {},
   "outputs": [],
   "source": [
    "@gif.frame\n",
    "def plot_prompt_reward_average_gif(i, frame, clip_rewards, reward_average, prompts):\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(20, 10))\n",
    "    for j in range(clip_rewards.shape[0]):\n",
    "        ax[0].plot(np.arange(0, i), clip_rewards[j, :i], ':', label=prompts[j])\n",
    "    ax[0].plot(np.arange(0, i), reward_average[:i], '-', color='r', label='Average')\n",
    "    ax[0].set_title('CLIP reward')\n",
    "    ax[0].legend()\n",
    "    ax[1].imshow(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "ffdce986",
   "metadata": {},
   "outputs": [],
   "source": [
    "@gif.frame\n",
    "def plot_prompt_encoding_average_gif(i, frame, clip_rewards):\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(20, 10))\n",
    "    ax[0].plot(np.arange(0, i), clip_rewards[:i])\n",
    "    ax[0].set_title('CLIP reward')\n",
    "    ax[1].imshow(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "02045f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_gif(frames, clip_rewards):\n",
    "    gif_frames = []\n",
    "    for i in tqdm(range(len(frames))):\n",
    "        gif_frames.append(plot_rewards_gif(i, frames[i], clip_rewards))\n",
    "    gif.save(gif_frames, PROMPT + '_' + FILENAME + '.gif', duration=33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "dCiGtYi-IOUN",
   "metadata": {
    "id": "dCiGtYi-IOUN"
   },
   "outputs": [],
   "source": [
    "def generate_prompt_reward_average_gif(frames, clip_rewards, reward_average, prompts, filename):\n",
    "    gif_frames = []\n",
    "    for i in tqdm(range(len(frames))):\n",
    "        gif_frames.append(plot_rewards_ensemble_gif(i, frames[i], clip_rewards, reward_average, prompts))\n",
    "    gif.save(gif_frames, filename, duration=33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "22fdf7e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_prompt_encoding_average_gif(frames, clip_rewards, filename):\n",
    "    gif_frames = []\n",
    "    for i in tqdm(range(len(frames))):\n",
    "        gif_frames.append(plot_prompt_encoding_average_gif(i, frames[i], clip_rewards))\n",
    "    gif.save(gif_frames, filename, duration=33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "1fe4ef97",
   "metadata": {
    "id": "1fe4ef97"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1235dbd4062f43c7a0c63f6776f7733a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/319 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_prompt_encoding_average_gif(frames, clip_rewards, 'Prompt_encodnig_average.gif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5545b1a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "ani_env_gif",
   "language": "python",
   "name": "ani_env_gif"
  },
  "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.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "af2b1c113fb30453ab0f38cb28889d45a8255f9a09eebc90dd674ef9d0f3b315"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "282008c3005b4707a161a4364208bf1b": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e4b42024029146538544a7d68c33b930": {
     "model_module": "@jupyter-widgets/output",
     "model_module_version": "1.0.0",
     "model_name": "OutputModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/output",
      "_model_module_version": "1.0.0",
      "_model_name": "OutputModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/output",
      "_view_module_version": "1.0.0",
      "_view_name": "OutputView",
      "layout": "IPY_MODEL_282008c3005b4707a161a4364208bf1b",
      "msg_id": "",
      "outputs": [
       {
        "data": {
         "text/html": "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080\">   3%</span> <span style=\"color: #f92672; text-decoration-color: #f92672\">━━</span><span style=\"color: #3a3a3a; text-decoration-color: #3a3a3a\">╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━</span> <span style=\"color: #008000; text-decoration-color: #008000\">6,694/200,000 </span> [ <span style=\"color: #808000; text-decoration-color: #808000\">0:03:33</span> &lt; <span style=\"color: #008080; text-decoration-color: #008080\">1:57:59</span> , <span style=\"color: #800000; text-decoration-color: #800000\">27 it/s</span> ]\n</pre>\n",
         "text/plain": "\u001b[35m   3%\u001b[0m \u001b[38;2;249;38;114m━━\u001b[0m\u001b[38;5;237m╺\u001b[0m\u001b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6,694/200,000 \u001b[0m [ \u001b[33m0:03:33\u001b[0m < \u001b[36m1:57:59\u001b[0m , \u001b[31m27 it/s\u001b[0m ]\n"
        },
        "metadata": {},
        "output_type": "display_data"
       }
      ]
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}