natural-language-youtube-search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MlDWup4TdWki"
   },
   "source": [
    "# Search YouTube videos using natural language\n",
    "\n",
    "You can use this notebook to play with [OpenAI's CLIP](https://openai.com/blog/clip/) neural network for searching through YouTube videos using natural language."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IU9bfB_GeYjM"
   },
   "source": [
    "## How it works\n",
    "\n",
    "1. Download the YouTube video\n",
    "2. Extract every N-th frame \n",
    "3. Encode all frames using using CLIP\n",
    "4. Encode a natural language search query using CLIP\n",
    "5. Find the images that best match the search query"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bYWJ4Vhte9H8"
   },
   "source": [
    "## Select a YouTube video\n",
    "\n",
    "Paste a link to a YouTube video or choose one of the examples. \n",
    "\n",
    "Choose a value for `N` which defines how many frames should be skipped. `N = 30` is usually about 1 frame every second."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mrDn4qnmCQCC"
   },
   "source": [
    "# Dashcam video of driving around San Francisco\n",
    "video_url = \"https://www.youtube.com/watch?v=PGMu_Z89Ao8\"  \n",
    "\n",
    "# How much frames to skip\n",
    "N = 120"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OQEncpKceUUy"
   },
   "source": [
    "## Setup the environment\n",
    "\n",
    "Install CLIP and install its dependencies."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "LQ_le-BO30LB"
   },
   "source": [
    "# pytube is used to download videos from YouTube\n",
    "!pip install pytube\n",
    "\n",
    "!pip install opencv-python\n",
    "!pip install pandas\n",
    "!pip install numpy==1.23.2\n",
    "\n",
    "# Intall a newer version of plotly\n",
    "!pip install plotly==4.14.3\n",
    "\n",
    "# Install CLIP from the GitHub repo\n",
    "!pip install git+https://github.com/openai/CLIP.git\n",
    "\n",
    "# Install torch 1.7.1 with GPU support\n",
    "!pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 -f https://download.pytorch.org/whl/torch_stable.html"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8snSpx68fgPK"
   },
   "source": [
    "## Download the video and process with CLIP\n",
    "\n",
    "In this section the video is downloaded, every N-th frame is extracted and each frame is processed using CLIP."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gdQCw7S0gbPp"
   },
   "source": [
    "Download the video locallyat 360p resoultion using `pytube`. A bigger resolution is not needed because CLIP scales down the images anyway."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Uz8HXiTL4T5Z"
   },
   "source": [
    "from pytubefix import YouTube\n",
    "\n",
    "# Choose a video stream with resolution of 360p\n",
    "streams = YouTube(video_url).streams.filter(adaptive=True, subtype=\"mp4\", resolution=\"360p\", only_video=True)\n",
    "\n",
    "# Check if there is a valid stream\n",
    "if len(streams) == 0:\n",
    "  raise \"No suitable stream found for this YouTube video!\"\n",
    "\n",
    "# Download the video as video.mp4\n",
    "print(\"Downloading...\")\n",
    "streams[0].download(filename=\"video.mp4\")\n",
    "print(\"Download completed.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hfG5b7_Ig_CO"
   },
   "source": [
    "Extract every `N-th` frame of the video."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tVmQBxqpCAX7",
    "outputId": "1945414f-10e8-4b8e-a252-1e5a32102aaa"
   },
   "source": [
    "import cv2\n",
    "from PIL import Image\n",
    "\n",
    "# The frame images will be stored in video_frames\n",
    "video_frames = []\n",
    "\n",
    "# Open the video file\n",
    "capture = cv2.VideoCapture('video.mp4')\n",
    "fps = capture.get(cv2.CAP_PROP_FPS)\n",
    "\n",
    "current_frame = 0\n",
    "while capture.isOpened():\n",
    "  # Read the current frame\n",
    "  ret, frame = capture.read()\n",
    "\n",
    "  # Convert it to a PIL image (required for CLIP) and store it\n",
    "  if ret == True:\n",
    "    video_frames.append(Image.fromarray(frame[:, :, ::-1]))\n",
    "  else:\n",
    "    break\n",
    "\n",
    "  # Skip N frames\n",
    "  current_frame += N\n",
    "  capture.set(cv2.CAP_PROP_POS_FRAMES, current_frame)\n",
    "\n",
    "# Print some statistics\n",
    "print(f\"Frames extracted: {len(video_frames)}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dw6oW1-sh50o"
   },
   "source": [
    "Load the public CLIP model."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ztgBXLiqC5M-"
   },
   "source": [
    "import clip\n",
    "import torch\n",
    "\n",
    "# Load the open CLIP model\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model, preprocess = clip.load(\"ViT-B/32\", device=device)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "75fR84yEibsx"
   },
   "source": [
    "Encode all frames using CLIP. The encoding is done in batches for a better efficiency."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "RfLC2iEiE3yL"
   },
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# You can try tuning the batch size for very large videos, but it should usually be OK\n",
    "batch_size = 256\n",
    "batches = math.ceil(len(video_frames) / batch_size)\n",
    "\n",
    "# The encoded features will bs stored in video_features\n",
    "video_features = torch.empty([0, 512], dtype=torch.float16).to(device)\n",
    "\n",
    "# Process each batch\n",
    "for i in range(batches):\n",
    "  print(f\"Processing batch {i+1}/{batches}\")\n",
    "\n",
    "  # Get the relevant frames\n",
    "  batch_frames = video_frames[i*batch_size : (i+1)*batch_size]\n",
    "  \n",
    "  # Preprocess the images for the batch\n",
    "  batch_preprocessed = torch.stack([preprocess(frame) for frame in batch_frames]).to(device)\n",
    "  \n",
    "  # Encode with CLIP and normalize\n",
    "  with torch.no_grad():\n",
    "    batch_features = model.encode_image(batch_preprocessed)\n",
    "    batch_features /= batch_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "  # Append the batch to the list containing all features\n",
    "  video_features = torch.cat((video_features, batch_features))\n",
    "\n",
    "# Print some stats\n",
    "print(f\"Features: {video_features.shape}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JB8p4hPijG-G"
   },
   "source": [
    "## Define functions for searching the video\n",
    "\n",
    "This section defines the functions used to search the video."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Ytk8xkb_LZ44"
   },
   "source": [
    "import plotly.express as px\n",
    "import datetime\n",
    "from IPython.core.display import HTML\n",
    "\n",
    "def search_video(search_query, display_heatmap=True, display_results_count=3):\n",
    "\n",
    "  # Encode and normalize the search query using CLIP\n",
    "  with torch.no_grad():\n",
    "    text_features = model.encode_text(clip.tokenize(search_query).to(device))\n",
    "    text_features /= text_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "  # Compute the similarity between the search query and each frame using the Cosine similarity\n",
    "  similarities = (100.0 * video_features @ text_features.T)\n",
    "  values, best_photo_idx = similarities.topk(display_results_count, dim=0)\n",
    "\n",
    "  # Display the heatmap\n",
    "  if display_heatmap:\n",
    "    print(\"Search query heatmap over the frames of the video:\")\n",
    "    fig = px.imshow(similarities.T.cpu().numpy(), height=50, aspect='auto', color_continuous_scale='viridis')\n",
    "    fig.update_layout(coloraxis_showscale=False)\n",
    "    fig.update_xaxes(showticklabels=False)\n",
    "    fig.update_yaxes(showticklabels=False)\n",
    "    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))\n",
    "    fig.show()\n",
    "    print()\n",
    "\n",
    "  # Display the top 3 frames\n",
    "  for frame_id in best_photo_idx:\n",
    "    display(video_frames[frame_id])\n",
    "\n",
    "    # Find the timestamp in the video and display it\n",
    "    seconds = round(frame_id.cpu().numpy()[0] * N / fps)\n",
    "    display(HTML(f\"Found at {str(datetime.timedelta(seconds=seconds))} (<a target=\\\"_blank\\\" href=\\\"{video_url}&t={seconds}\\\">link</a>)\"))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LLA89mYqjpot"
   },
   "source": [
    "## Search queries\n",
    "\n",
    "You can run the example search queries for some interesting results or you can try out your own.\n",
    "\n",
    "On top of eah result you will also find a heatmap that shows how likely the search query is at all frames from the video."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hLKKPzjXleHq"
   },
   "source": [
    "### \"A fire truck\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "w5iz0gIRj4hL",
    "outputId": "dd365bca-157c-4710-c909-b4659e6535ed"
   },
   "source": [
    "search_video(\"a fire truck\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IDjuAxt7li5y"
   },
   "source": [
    "### \"Road works\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "cwyQ6mYGksWq",
    "outputId": "aace95cc-854f-4366-b09e-2bb648faa3f0"
   },
   "source": [
    "search_video(\"road works\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i0x5MPh2ll-1"
   },
   "source": [
    "### \"People crossing the street\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "O5IuxlcLkxlH",
    "outputId": "7bb22d64-82a5-4683-e498-fab6506e2149"
   },
   "source": [
    "search_video(\"people crossing the street\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9lZ4tWROlpzm"
   },
   "source": [
    "### \"The Embarcadero\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "Rd9Jw9m3k1pV",
    "outputId": "af38c58e-f1da-41d2-8be9-64c0a3238fc7"
   },
   "source": [
    "search_video(\"the Embarcadero\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "isII2s3AlwnW"
   },
   "source": [
    "### \"Waiting at the red light\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "z3TzNk-DlDeC",
    "outputId": "b0c91cf2-7ac3-4085-8346-2e752fb4ca65"
   },
   "source": [
    "search_video(\"waiting at the red light\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kvRBsv7Alzcl"
   },
   "source": [
    "### \"Green bike lane\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "LxTh_S_KlJL6",
    "outputId": "37f73410-2d66-4468-e5cf-fe9f1f69a409"
   },
   "source": [
    "search_video(\"green bike lane\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vFZois6fl1wq"
   },
   "source": [
    "### \"A street with tram tracks\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "NL0yyNCOlLmi",
    "outputId": "2f19a3e2-27b9-4a34-8a3c-984bdc77bb3e"
   },
   "source": [
    "search_video(\"a street with tram tracks\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X4UGP2Ndl5oE"
   },
   "source": [
    "### \"The Transamerica Pyramid\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "HR4LJoRElRD6",
    "outputId": "318da3e6-3529-4122-c4aa-31b4c953af82"
   },
   "source": [
    "search_video(\"the Transamerica Pyramid\")"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "OQEncpKceUUy",
    "8snSpx68fgPK",
    "JB8p4hPijG-G"
   ],
   "name": "natural-language-youtube-search.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}