Use_Case_3_Video_Descriptions.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "# Copyright 2024 Reddit, Inc.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "metadata": {
        "id": "pT4bpoMxzPsm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Use Case 3. Video Descriptions\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/reddit/kdd2024-tutorial-breaking-barriers/blob/master/Use_Case_3_Video_Descriptions.ipynb)"
      ],
      "metadata": {
        "id": "OsbsT6Um0EgO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Overview\n",
        "\n",
        "This notebook guides participants through the\n",
        "steps of designing and implementing a pipeline to generate video\n",
        "descriptions, combining keyframe extraction, image captioning,\n",
        "audio transcript and summarization using LLMs. We will explore\n",
        "the challenges and advantages for different types of video content."
      ],
      "metadata": {
        "id": "c9iZrO4Rda8G"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Setting Up Google Colab\n",
        "Google Colab provides a convenient platform to run Python code in the cloud, with access to powerful computing resources, including GPUs. Again, for this tutorial, it is recommended to enable GPU acceleration:\n",
        "\n",
        "1.   Click on *Runtime* in the top menu.\n",
        "2.   Select *Change runtime type*.\n",
        "3.   In the dialog that appears, under *Hardware accelerator*, choose **T4 GPU** (or any other GPU that you may have access to) if it is not already enabled.\n",
        "4.   Click *Save*."
      ],
      "metadata": {
        "id": "WhueG38H2CLc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Requirements\n",
        "\n",
        "\n",
        "Run the following cell to install the required Python packages."
      ],
      "metadata": {
        "id": "VuI4RoF82OEh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -U openai-whisper transformers==4.41.2 bitsandbytes accelerate flash_attn"
      ],
      "metadata": {
        "id": "VlJkwhI73Mm4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Settings\n",
        "\n",
        "Run the following cells to make some convenient settings."
      ],
      "metadata": {
        "id": "231Wh_HUUgDf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Disable Transformer warnings\n",
        "import logging\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "\n",
        "import transformers\n",
        "transformers.logging.set_verbosity_error()\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "# Set GPU device\n",
        "import torch\n",
        "torch.set_default_device(\"cuda\") # or \"cpu\" is GPU is not available"
      ],
      "metadata": {
        "id": "-Z4jG0074gfl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the following cell to get the run time on every cell execution:"
      ],
      "metadata": {
        "id": "HomTygaHjI46"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install ipython-autotime\n",
        "%load_ext autotime"
      ],
      "metadata": {
        "id": "9P6DMHXhrkSw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the following cell to enable wrap when printing long strings:"
      ],
      "metadata": {
        "id": "tCmvlRdPspjG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML(\"\"\"\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  \"\"\"))\n",
        "get_ipython().events.register(\"pre_run_cell\", set_css)"
      ],
      "metadata": {
        "id": "nJJcC9czsX0K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Some other useful imports:"
      ],
      "metadata": {
        "id": "SWFc2ZhgVK9H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from PIL import Image\n",
        "from tqdm import tqdm\n",
        "from glob import glob\n",
        "import re"
      ],
      "metadata": {
        "id": "6EzYxAVcVMhX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Test Videos\n",
        "\n",
        "We will use the first video in the tutorial dataset."
      ],
      "metadata": {
        "id": "DBshtSyB2aH8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget \"https://raw.githubusercontent.com/reddit/kdd2024-tutorial-breaking-barriers/main/media/video1.mp4\""
      ],
      "metadata": {
        "id": "YHH8BLUoVEz5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import moviepy.editor\n",
        "moviepy.editor.ipython_display(\"video1.mp4\", height=400, maxduration=300)"
      ],
      "metadata": {
        "id": "hQAnP8sdVGaX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Use Case 1. Videos with Speech\n",
        "\n",
        "The first use case is focused on videos that have a speech track.\n",
        "\n",
        "For these videos, the best approach is to:\n",
        "\n",
        "1. Generate the audio transcript (with a speech-to-text LLM such as Whisper) and then\n",
        "\n",
        "2. Use an LLM to generate a summary of the transcript."
      ],
      "metadata": {
        "id": "sOqtvvnk2eaK"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1) Extract Audio Transcript"
      ],
      "metadata": {
        "id": "IjTmtjfQ3B0c"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Import package\n",
        "import whisper\n",
        "\n",
        "# Load model - choose the most appropriate\n",
        "stt_model = whisper.load_model(\"base\")\n",
        "\n",
        "# Transcribe audio/video\n",
        "result = stt_model.transcribe(\"video1.mp4\")\n",
        "transcript = result[\"text\"]\n",
        "\n",
        "# Print transcript\n",
        "print(transcript)"
      ],
      "metadata": {
        "id": "lmI8MrND3Vkf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2) Generate Description based on Transcript\n",
        "\n",
        "In this case we are going to choose the same `imp-v1-3b` model that we used for image captions."
      ],
      "metadata": {
        "id": "u30ENY403k70"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load model\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "vision_model = AutoModelForCausalLM.from_pretrained(\n",
        "    \"MILVLG/imp-v1-3b\",\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True)\n",
        "vision_tokenizer = AutoTokenizer.from_pretrained(\n",
        "    \"MILVLG/imp-v1-3b\",\n",
        "    trust_remote_code=True)"
      ],
      "metadata": {
        "id": "rj08pXHJ6UVx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Text prompt\n",
        "prompt = f\"Write a short summary for a video that has the following transcript:\\n{transcript}\"\n",
        "\n",
        "# Process text\n",
        "text = vision_tokenizer.apply_chat_template(\n",
        "    [{\"role\": \"user\", \"content\": prompt}],\n",
        "    tokenize=False,\n",
        "    add_generation_prompt=True\n",
        ")\n",
        "input_ids = vision_tokenizer(text, return_tensors=\"pt\").input_ids\n",
        "image_tensor = None\n",
        "\n",
        "# Generate the answer\n",
        "output_ids = vision_model.generate(\n",
        "    input_ids,\n",
        "    max_new_tokens=256,\n",
        "    images=image_tensor,\n",
        "    use_cache=True)[0]\n",
        "description = vision_tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True)\n",
        "\n",
        "# Display description\n",
        "print(description)"
      ],
      "metadata": {
        "id": "o8HaazoC6h5u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Improved Implementation"
      ],
      "metadata": {
        "id": "zaRdnJ4SAUmm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def prompt(vision_tokenizer, vision_model, prompt, image = None):\n",
        "  # Process text and optionally image\n",
        "  text = vision_tokenizer.apply_chat_template(\n",
        "      [{\"role\": \"user\", \"content\": prompt}],\n",
        "      tokenize=False,\n",
        "      add_generation_prompt=True\n",
        "  )\n",
        "  input_ids = vision_tokenizer(text, return_tensors=\"pt\").input_ids\n",
        "  if image:\n",
        "    image_tensor = vision_model.image_preprocess(image)\n",
        "  else:\n",
        "    image_tensor = None\n",
        "\n",
        "  # Generate the answer\n",
        "  output_ids = vision_model.generate(\n",
        "      input_ids,\n",
        "      max_new_tokens=256,\n",
        "      images=image_tensor,\n",
        "      use_cache=True)[0]\n",
        "  result = vision_tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True)\n",
        "  result = re.sub(r\"<\\|im_end\\|>.*$\", \"\", result.replace(\"\\n\", \"\"))\n",
        "\n",
        "  return result\n",
        "\n",
        "# Run\n",
        "description = prompt(vision_tokenizer,\n",
        "                     vision_model,\n",
        "                     f\"Write a short summary for a video that has the following transcript:\\n{transcript}\")\n",
        "print(description)"
      ],
      "metadata": {
        "id": "5ztxtIwUAXaX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Use Case 2. Videos without Speech\n",
        "\n",
        "In this case, the only available information are the video frames, so the approach is:\n",
        "\n",
        "1. Extract the key frames of the video, those that best represent its contents.\n",
        "2. Generate a caption for each frame using a multimodal LLM.\n",
        "3. Use an LLM to generate a summary based on all the frame captions."
      ],
      "metadata": {
        "id": "hcMiry4L8Qg-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1) Extract Key Frames\n",
        "\n",
        "We will extract the key frames using `ffmpeg` tool to the `keyframes` folder:"
      ],
      "metadata": {
        "id": "eop2rzyf8xjV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir keyframes\n",
        "!ffmpeg -skip_frame nokey -i \"video1.mp4\" -vsync vfr \"keyframes/frame-%2d.jpg\" -hide_banner -loglevel error\n",
        "!ls keyframes"
      ],
      "metadata": {
        "id": "StysV4TW8-4p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "all_frames = sorted(glob(f\"keyframes/*jpg\"))\n",
        "\n",
        "frames_to_process = 10\n",
        "\n",
        "frames = all_frames[0::int(len(all_frames)/frames_to_process-1)]\n",
        "for frame in frames:\n",
        "  display(frame, Image.open(frame))"
      ],
      "metadata": {
        "id": "rgpLZZhz99Gf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Programmatically:"
      ],
      "metadata": {
        "id": "6xpoBEfF9WbC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2) Generate Frame Captions"
      ],
      "metadata": {
        "id": "qMpyTEC1AKCk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "captions = []\n",
        "for frame in tqdm(frames, total=len(frames)):\n",
        "  image = Image.open(frame)\n",
        "  caption = prompt(vision_tokenizer,\n",
        "                   vision_model,\n",
        "                   \"<image>\\nWrite a very short caption for the image with less than 20 words\",\n",
        "                   Image.open(frame))\n",
        "  captions.append(caption)\n",
        "captions"
      ],
      "metadata": {
        "id": "9lDSK4_vA_gz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3) Generate Description based on Frame Captions"
      ],
      "metadata": {
        "id": "1XDvoe24CDQo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "caption_list = \"\\n-\".join(captions)\n",
        "description = prompt(vision_tokenizer,\n",
        "                     vision_model,\n",
        "                     f\"Write a short summary for a video that has the following frames:\\n{caption_list}\")\n",
        "print(description)"
      ],
      "metadata": {
        "id": "q4-SdVPjCJiJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "## Complete Implementation (MLLMv2)"
      ],
      "metadata": {
        "id": "S17j6ZtmCgSE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import whisper\n",
        "from PIL import Image\n",
        "import tempfile\n",
        "from glob import glob\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "import re\n",
        "\n",
        "class MLLMv2:\n",
        "\n",
        "  def __init__(self):\n",
        "    torch.set_default_device(\"cuda\")\n",
        "    self.vision_model = AutoModelForCausalLM.from_pretrained(\n",
        "      \"MILVLG/imp-v1-3b\",\n",
        "      torch_dtype=torch.float16,\n",
        "      device_map=\"auto\",\n",
        "      trust_remote_code=True)\n",
        "    self.vision_tokenizer = AutoTokenizer.from_pretrained(\n",
        "        \"MILVLG/imp-v1-3b\",\n",
        "        trust_remote_code=True)\n",
        "    self.stt_model = whisper.load_model(\"base\")\n",
        "\n",
        "  def get_image_caption(self,\n",
        "                        image_file: Image,\n",
        "                        base_prompt=\"Write a very short caption for the image with less than 20 words\") -> str:\n",
        "    return self.prompt_llm(image_file, base_prompt)\n",
        "\n",
        "  def get_image_description(self,\n",
        "                            image_file: Image,\n",
        "                            base_prompt=\"Write a short description for the image\") -> str:\n",
        "    return self.prompt_llm(image_file, base_prompt)\n",
        "\n",
        "  def get_video_transcript(self,\n",
        "                           video_file: str):\n",
        "    result = self.stt_model.transcribe(video_file)\n",
        "    transcript = result[\"text\"]\n",
        "    return transcript\n",
        "\n",
        "  def get_video_frame_captions(self,\n",
        "                               video_file: str,\n",
        "                               base_prompt=\"Write a very short caption for the image with less than 20 words\",\n",
        "                               frames_to_process=5):\n",
        "    # Extract frames\n",
        "    with tempfile.TemporaryDirectory() as tmpdirname:\n",
        "      os.system(f\"ffmpeg -skip_frame nokey -i '{video_file}' -vsync vfr '{tmpdirname}/frame-%2d.jpg' -hide_banner -loglevel error\")\n",
        "      all_frames = sorted(glob(f\"{tmpdirname}/*jpg\"))\n",
        "      print(f\"All frames: {len(all_frames)}\")\n",
        "      frames = all_frames[0::int(len(all_frames)/frames_to_process-1)]\n",
        "      print(f\"Selected frames: {len(frames)}\")\n",
        "      # Generate captions\n",
        "      captions = []\n",
        "      for frame in tqdm(frames, total=len(frames)):\n",
        "        image = Image.open(frame)\n",
        "        caption = self.prompt_llm(image, base_prompt)\n",
        "        if caption:\n",
        "          captions.append(caption)\n",
        "    return {\"total_frames\": len(all_frames),\n",
        "            \"frames\": frames,\n",
        "            \"captions\": captions}\n",
        "\n",
        "  def get_video_description(self,\n",
        "                            video_file: str,\n",
        "                            include_audio: bool,\n",
        "                            include_frames: bool):\n",
        "    transcript = None\n",
        "    captions = None\n",
        "    if include_audio:\n",
        "      transcript = self.get_video_transcript(video_file)\n",
        "    if include_frames:\n",
        "      frame_captions = self.get_video_frame_captions(video_file)\n",
        "      captions = frame_captions[\"captions\"]\n",
        "    if transcript and captions:\n",
        "      prompt = f\"Write a short summary for a video that has the following transcript and frames:\\n\"\n",
        "      prompt += f\"## Transcript:\\n{transcript}\"\n",
        "      caption_list = \"\\n-\".join(captions)\n",
        "      prompt += f\"## Frames:\\n-{caption_list}\"\n",
        "    elif transcript:\n",
        "      prompt = f\"Write a short summary for a video that has the following transcript:\\n{transcript}\"\n",
        "    elif captions:\n",
        "      caption_list = \"\\n-\".join(captions)\n",
        "      prompt = f\"Write a short summary for a video that has the following frames:\\n-{caption_list}\"\n",
        "    result = self.prompt_llm(None, prompt) if prompt else None\n",
        "    return {\"transcript\": transcript,\n",
        "            \"captions\": captions,\n",
        "            \"description\": result}\n",
        "\n",
        "  def prompt_llm(self,\n",
        "                 image: Image,\n",
        "                 prompt: str,\n",
        "                 max_new_tokens: int = 256,\n",
        "                 temperature: float = 0.9,\n",
        "                 top_k: int = 50,\n",
        "                 top_p: float = 0.95) -> str:\n",
        "    if image:\n",
        "      text = self.vision_tokenizer.apply_chat_template(\n",
        "          [{\"role\": \"user\", \"content\": f\"<image>\\n{prompt}\"}],\n",
        "          tokenize=False,\n",
        "          add_generation_prompt=True\n",
        "      )\n",
        "      image_tensor = self.vision_model.image_preprocess(image)\n",
        "    else:\n",
        "      text = self.vision_tokenizer.apply_chat_template(\n",
        "          [{\"role\": \"user\", \"content\": f\"{prompt}\"}],\n",
        "          tokenize=False,\n",
        "          add_generation_prompt=True\n",
        "      )\n",
        "      image_tensor = None\n",
        "    input_ids = self.vision_tokenizer(text, return_tensors=\"pt\").input_ids\n",
        "    output_ids = self.vision_model.generate(\n",
        "      input_ids,\n",
        "      max_new_tokens=max_new_tokens,\n",
        "      images=image_tensor,\n",
        "      temperature=temperature,\n",
        "      do_sample=True,\n",
        "      top_k=top_k,\n",
        "      top_p=top_p,\n",
        "      use_cache=True)[0]\n",
        "    response = self.vision_tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True)\n",
        "    response = response.replace(\"\\n\", \" \").replace(\"  \", \" \").strip()\n",
        "    response = re.sub(r\"<\\|im_end\\|>.*$\", \"\", response)\n",
        "    return response"
      ],
      "metadata": {
        "id": "lhKnewImCjkL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load model\n",
        "vc = MLLMv2()"
      ],
      "metadata": {
        "id": "wJkxR9svl90e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Speech Transcript\n",
        "vc.get_video_transcript(\"video1.mp4\")"
      ],
      "metadata": {
        "id": "ZRSYykiCeR9i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Frame Captions\n",
        "vc.get_video_frame_captions(\"video1.mp4\")"
      ],
      "metadata": {
        "id": "1nKc_aHoeUlc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Video Description with Speech Transcript\n",
        "vc.get_video_description(\"video1.mp4\", True, False)"
      ],
      "metadata": {
        "id": "A5IC8QzdeXGT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Video Description with Frame Captions\n",
        "vc.get_video_description(\"video1.mp4\", False, True)"
      ],
      "metadata": {
        "id": "dKT-mPlJeaTD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Video Description with Speech Transcript and Frame Captions\n",
        "vc.get_video_description(\"video1.mp4\", True, True)"
      ],
      "metadata": {
        "id": "T88eRVL4eb4J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "# Discussion: How Video Description with Multimodal LLMs can improve Accessibility in Social Media\n",
        "\n",
        "- **Accessibility for all**: By making videos more accessible, multimodal LLMs can encourage greater participation and engagement from users with disabilities, fostering a more inclusive online community.\n",
        "- **Automatic tagging and categorization**: LLMs can analyze the generated captions and speech transcripts to automatically assign relevant tags and categories to videos, making it easier for users with disabilities to find relevant content.\n",
        "- **Multiple language support**: LLMs can be trained to generate captions and tags in multiple languages, improving accessibility for a wider range of users.\n",
        "- **Speech-to-text integration**: Utilizing speech transcripts alongside video frames, multimodal LLMs can generate more accurate and detailed captions for videos, even if the audio is unclear or contains background noise.\n",
        "- **Key frame captions**: For videos without speech, LLMs can analyze key frames in videos and generate descriptive captions for each, providing a textual summary for users who cannot watch the video.\n",
        "- **Understanding user engagement**: By analyzing user interactions with captions and tags, LLMs can provide valuable data on the accessibility of different content formats, enabling platform owners to improve user experience."
      ],
      "metadata": {
        "id": "RFUlDNa3HoLD"
      }
    }
  ]
}