Natural_Language_YouTube_Searcher.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Natural Language YouTube Searcher",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tg-bomze/collection-of-notebooks/blob/master/Natural_Language_YouTube_Searcher.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PFQ_7sRn5fUl"
      },
      "source": [
        "<b><font color=\"black\" size=\"+4\">Natural Language YouTube Searcher</font></b>\r\n",
        "\r\n",
        "<b><font color=\"black\" size=\"+2\">Based on:</font></b>\r\n",
        "\r\n",
        "**GitHub repository**: [CLIP](https://github.com/openai/CLIP)\r\n",
        "\r\n",
        "Article: [Learning Transferable Visual Models From Natural Language Supervision](https://cdn.openai.com/papers/Learning_Transferable_Visual_Models_From_Natural_Language_Supervision.pdf)\r\n",
        "\r\n",
        "Creator: **[OpenAI](https://openai.com/)**\r\n",
        "\r\n",
        "Colab based on **[Vladimir Haltakov's](https://haltakov.net/)** [notebook](https://colab.research.google.com/github/haltakov/natural-language-youtube-search/blob/main/natural-language-youtube-search.ipynb).\r\n",
        "\r\n",
        "<b><font color=\"black\" size=\"+2\">Colab created by:</font></b>\r\n",
        "\r\n",
        "GitHub: [@tg-bomze](https://github.com/tg-bomze),\r\n",
        "Telegram: [@bomze](https://t.me/bomze),\r\n",
        "Twitter: [@tg_bomze](https://twitter.com/tg_bomze).\r\n",
        "\r\n",
        "---\r\n",
        "\r\n",
        "\r\n",
        "```\r\n",
        "(ENG) To get started, click on the button (where the red arrow indicates). After clicking, wait until the execution is complete.\r\n",
        "```\r\n",
        "```\r\n",
        "(RUS) Чтобы начать, поочередно нажимайте на кнопки (куда указывают красные стрелки), дожидаясь завершения выполнения каждого блока.\r\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IJ0O4uS2wZ2L",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Install all necessary libraries</font></b>\r\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Установить все необходимые библиотеки</font></b>\r\n",
        "\r\n",
        "!nvidia-smi -L\r\n",
        "\r\n",
        "try: \r\n",
        "  !pip3 install googletrans==3.1.0a0\r\n",
        "  from googletrans import Translator, constants\r\n",
        "  from pprint import pprint\r\n",
        "  translator = Translator()\r\n",
        "except: pass\r\n",
        "\r\n",
        "import subprocess\r\n",
        "from IPython.display import clear_output\r\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"UTF-8\").split(\", \") if s.startswith(\"release\")][0].split(\" \")[-1]\r\n",
        "print(\"CUDA version:\", CUDA_version)\r\n",
        "\r\n",
        "if CUDA_version == \"10.0\":\r\n",
        "    torch_version_suffix = \"+cu100\"\r\n",
        "elif CUDA_version == \"10.1\":\r\n",
        "    torch_version_suffix = \"+cu101\"\r\n",
        "elif CUDA_version == \"10.2\":\r\n",
        "    torch_version_suffix = \"\"\r\n",
        "else:\r\n",
        "    torch_version_suffix = \"+cu110\"\r\n",
        "\r\n",
        "# pytube is used to download videos from YouTube\r\n",
        "!pip install pytube\r\n",
        "\r\n",
        "# Intall a newer version of plotly\r\n",
        "!pip install plotly==4.14.3\r\n",
        "\r\n",
        "# Install CLIP from the GitHub repo\r\n",
        "!pip install git+https://github.com/openai/CLIP.git\r\n",
        "\r\n",
        "# Install torch 1.7.1 with GPU support\r\n",
        "!pip install torch==1.7.1{torch_version_suffix} torchvision==0.8.2{torch_version_suffix} -f https://download.pytorch.org/whl/torch_stable.html ftfy regex\r\n",
        "\r\n",
        "import clip\r\n",
        "import torch\r\n",
        "from pytube import YouTube\r\n",
        "import cv2\r\n",
        "from PIL import Image\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "import os\r\n",
        "import math\r\n",
        "import numpy as np\r\n",
        "import plotly.express as px\r\n",
        "\r\n",
        "# Load the open CLIP model\r\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\r\n",
        "model, preprocess = clip.load(\"ViT-B/32\", device=device)\r\n",
        "\r\n",
        "def search_video(search_query, display_heatmap=True, display_results_count=3):\r\n",
        "\r\n",
        "  # Encode and normalize the search query using CLIP\r\n",
        "  with torch.no_grad():\r\n",
        "    text_features = model.encode_text(clip.tokenize(search_query).to(device))\r\n",
        "    text_features /= text_features.norm(dim=-1, keepdim=True)\r\n",
        "\r\n",
        "  # Compute the similarity between the search query and each frame using the Cosine similarity\r\n",
        "  similarities = (100.0 * video_features @ text_features.T)\r\n",
        "  values, best_photo_idx = similarities.topk(display_results_count, dim=0)\r\n",
        "\r\n",
        "  # Display the heatmap\r\n",
        "  if display_heatmap:\r\n",
        "    print(\"Search query heatmap over the frames of the video:\")\r\n",
        "    fig = px.imshow(similarities.T.cpu().numpy(), height=50, aspect='auto', color_continuous_scale='viridis')\r\n",
        "    fig.update_layout(coloraxis_showscale=False)\r\n",
        "    fig.update_xaxes(showticklabels=False)\r\n",
        "    fig.update_yaxes(showticklabels=False)\r\n",
        "    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))\r\n",
        "    fig.show()\r\n",
        "    print()\r\n",
        "\r\n",
        "  # Display the top N frames\r\n",
        "  for frame_id in best_photo_idx:\r\n",
        "    display(video_frames[frame_id])\r\n",
        "    try:\r\n",
        "      in_sec = int((int(frame_id)*frames_to_skip)/fps_of_video)\r\n",
        "      print(f'This frame is at {in_sec} seconds of the video.\\n')\r\n",
        "    except: print()\r\n",
        "\r\n",
        "clear_output()\r\n",
        "print('Done!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uz8HXiTL4T5Z",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Get Video</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Получить видеофайл</font></b>\n",
        "\n",
        "# Dashcam video of driving around San Francisco\n",
        "video_url = \"https://www.youtube.com/watch?v=odM92ap8_c0\" #@param {type:'string'}  \n",
        "\n",
        "# How much frames to skip\n",
        "frames_to_skip = 60 #@param {type:\"integer\"}\n",
        "#@markdown *(ENG) The more \"frames_to_skip\", the faster the neural network works, but less accurate.*\n",
        "\n",
        "#@markdown *(RUS) Чем больше значение \"frames_to_skip\", тем быстрее работает нейросеть, но менее точно.*\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\")\n",
        "print(\"Download completed.\")\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_of_video = int(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 += frames_to_skip\n",
        "  capture.set(cv2.CAP_PROP_POS_FRAMES, current_frame)\n",
        "\n",
        "# Print some statistics\n",
        "print(f\"Frames extracted: {len(video_frames)}\")\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}\")\n",
        "print('Done!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w5iz0gIRj4hL",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Start Searching</font></b>\r\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Начать поиск</font></b>\r\n",
        "search_text = 'king kong' #@param {type:'string'}\r\n",
        "try: prompt = translator.translate(search_text).text\r\n",
        "except: prompt = search_text\r\n",
        "\r\n",
        "max_num_results = 3 #@param {type:\"slider\", min:1, max:12, step:1}\r\n",
        "\r\n",
        "search_video(prompt, display_results_count=max_num_results)\r\n",
        "#@markdown *(ENG) The more \"max_num_results\", the more search results, but less accurate.*\r\n",
        "\r\n",
        "#@markdown *(RUS) Чем больше значение \"max_num_results\", тем больше результатов поиска, но меньше точность.*\r\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}