Vid2Music.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tld21HLvAv47",
        "outputId": "0b3de7c5-c4ee-4844-8723-b8f1c629ec18"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.28.1)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.1)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.13.3)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.22.4)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.27.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.65.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.14.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2022.10.31)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (2023.4.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (4.5.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2022.12.7)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.12)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.15)\n"
          ]
        }
      ],
      "source": [
        "!pip install transformers"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZD9RJvKCNGlB",
        "outputId": "ab9de606-5753-4307-b0e3-67caf4ba4598"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import sys\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "from torchvision import transforms, models\n",
        "from PIL import Image\n",
        "from scipy.spatial import KDTree\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import torch\n",
        "from transformers import BertTokenizer, BertModel"
      ],
      "metadata": {
        "id": "41uYaCoyTT0y"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class LanguageModel():\n",
        "  def __init__(self):\n",
        "    self.model = BertModel.from_pretrained('bert-base-uncased',\n",
        "                                    output_hidden_states = True, # Whether the model returns all hidden-states.\n",
        "                                    )\n",
        "    self.model.eval()\n",
        "    self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
        "  \n",
        "  def __tokenize__(self, text):\n",
        "    start_text = \"[CLS] \"\n",
        "    end_text = \" [SEP]\"\n",
        "\n",
        "    marked_text = start_text + text + end_text\n",
        "\n",
        "    tokenized_text = self.tokenizer.tokenize(marked_text)\n",
        "    # Map the token strings to their vocabulary indeces.\n",
        "    indexed_tokens = self.tokenizer.convert_tokens_to_ids(tokenized_text)\n",
        "    segments_ids = [1] * len(tokenized_text)\n",
        "\n",
        "    # Convert inputs to PyTorch tensors\n",
        "    tokens_tensor = torch.tensor([indexed_tokens])\n",
        "    segments_tensors = torch.tensor([segments_ids])\n",
        "\n",
        "    \n",
        "\n",
        "    return tokens_tensor, segments_tensors\n",
        "  \n",
        "  def __tokensToEmbeddings__(self, tokens, segments):\n",
        "    with torch.no_grad():\n",
        "\n",
        "      outputs = self.model(tokens, segments)\n",
        "      hidden_states = outputs[2]\n",
        "  \n",
        "    token_embeddings = torch.stack(hidden_states, dim=0)\n",
        "    token_embeddings = torch.squeeze(token_embeddings, dim=1)\n",
        "    token_embeddings = token_embeddings.permute(1,0,2)\n",
        "\n",
        "    token_vecs = hidden_states[-2][0]\n",
        "\n",
        "    sentence_embedding = torch.mean(token_vecs, dim=0)\n",
        "\n",
        "    return sentence_embedding\n",
        "  \n",
        "  def generateEmbeddings(self, text):\n",
        "    tokens, segments = self.__tokenize__(text)\n",
        "    embedding = self.__tokensToEmbeddings__(tokens, segments)\n",
        "\n",
        "    return embedding\n",
        "\n",
        "\n",
        "class RecommenderSystem():\n",
        "  def __init__(self, songDF, LanguageModel, EmotionModel):\n",
        "    self.emotionModel = EmotionModel\n",
        "    self.LM = LanguageModel\n",
        "    self.songDF = songDF\n",
        "    self.songKDTree = None\n",
        "    self.songLookup = None\n",
        "    self.__initSongData__()\n",
        "  \n",
        "  def __initSongData__(self):\n",
        "    songsLyrics = self.songDF['caption'].values\n",
        "    songDF_id = self.songDF['ytid']\n",
        "    \n",
        "    songEmbeddings = []\n",
        "\n",
        "    songLookup = {}\n",
        "    # i=0\n",
        "    for (i, song) in enumerate(songsLyrics):\n",
        "      if i % 100 == 0:\n",
        "        print(\"Indexed \", i)\n",
        "      songEmb = self.LM.generateEmbeddings(song)\n",
        "      songEmbeddings.append(np.array(songEmb))\n",
        "\n",
        "      songLookup[tuple(songEmb.tolist())] = songDF_id[i]\n",
        "    \n",
        "    songEmbeddings = np.array(songEmbeddings)\n",
        "    self.songKDTree = KDTree(songEmbeddings)\n",
        "    self.songLookup = songLookup\n",
        "\n",
        "  def __imgToEmotion__(self, img):\n",
        "    return self.emotionModelImg(img)\n",
        "  \n",
        "  def __vidToEmotion__(self, vid_path):\n",
        "    return self.emotionModel.predict_emotions(vid_path)\n",
        "  \n",
        "  def __emotionToSongRecs__(self, emotion, n):\n",
        "    emotionEmb = self.LM.generateEmbeddings(emotion)\n",
        "    dist, ind = self.songKDTree.query(np.array(emotionEmb), n)\n",
        "\n",
        "    recs = []\n",
        "\n",
        "    if n == 1:\n",
        "      recs.append(self.songLookup[tuple(self.songKDTree.data[ind])])\n",
        "    else:\n",
        "      recs = [self.songLookup[tuple(self.songKDTree.data[index])] for index in ind]\n",
        "    \n",
        "    return recs\n",
        "  \n",
        "  def getRecsFromImage(self, img, n):\n",
        "    emotion = self.__imgToEmotion__(img)\n",
        "    songRecs = self.__emotionToSongRecs__(emotion, n)\n",
        "    \n",
        "    return songRecs\n",
        "  \n",
        "  def getRecsFromVid(self, vid, n):\n",
        "    emotion = self.__vidToEmotion__(vid)\n",
        "    songRecs = []\n",
        "\n",
        "    for emotionPred, _ in emotion:\n",
        "      songRec = self.__emotionToSongRecs__(emotionPred, n)\n",
        "      songRecs += songRec  \n",
        "\n",
        "    return songRecs\n",
        "\n",
        "\n",
        "class EmotionModel():\n",
        "  def __init__(self, weights_path, emotion_labels, device):\n",
        "    self.device = device\n",
        "    self.emotion_labels = emotion_labels\n",
        "    self.model = models.vgg19(weights='VGG19_Weights.DEFAULT')\n",
        "    self.num_features = self.model.classifier[-1].in_features # num of features in last fully connected model layer\n",
        "    self.model.classifier[-1] = nn.Linear(self.num_features, 7) # 7 unique emotions in both train and test folders resp\n",
        "    self.model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))\n",
        "\n",
        "  \n",
        "  def predict_emotions(self, video_file_path):\n",
        "    # Load the video\n",
        "    cap = cv2.VideoCapture(video_file_path)\n",
        "\n",
        "    # Get the frame rate of the video\n",
        "    fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
        "\n",
        "    # Define the list of emotions\n",
        "    emotions = list(self.emotion_labels.keys())\n",
        " \n",
        "    idx_to_class = {v: k for k, v in self.emotion_labels.items()}\n",
        "    \n",
        "\n",
        "    # Initialize the emotion count dictionary\n",
        "    emotion_count = {emotion: 0 for emotion in emotions}\n",
        "\n",
        "    self.model.to(self.device)\n",
        "\n",
        "    # Loop over the frames of the video\n",
        "    while cap.isOpened():\n",
        "        # Read the frame\n",
        "        ret, frame = cap.read()\n",
        "\n",
        "        # Check if the frame was successfully read\n",
        "        if not ret:\n",
        "            break\n",
        "\n",
        "        # Preprocess the frame\n",
        "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
        "        frame = Image.fromarray(frame)\n",
        "        # Preprocess the frame\n",
        "        frame = transforms.Grayscale(num_output_channels=3)(frame) # Pre-trained model expects RGB\n",
        "        frame = transforms.Resize((224, 224))(frame) # Want a small standard image size to reduce load on GPU\n",
        "        frame = transforms.ToTensor()(frame)\n",
        "        frame = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])(frame) # Normalize(mean_pixel_val, std_pixel_val)\n",
        "        frame = frame.unsqueeze(0)\n",
        "\n",
        "        frame = frame.to(self.device)\n",
        "\n",
        "        # Pass the frame through the model to obtain the predicted emotion\n",
        "        with torch.no_grad():\n",
        "            outputs = self.model(frame)\n",
        "            _, predicted = torch.max(outputs.data, 1)\n",
        "            predicted_emotion = idx_to_class[predicted.item()]\n",
        "\n",
        "        # Increment the count for the predicted emotion\n",
        "        emotion_count[predicted_emotion] += 1\n",
        "\n",
        "        # Wait for the specified amount of time before processing the next frame\n",
        "        cv2.waitKey(1000 // fps)\n",
        "\n",
        "    # Release the video capture\n",
        "    cap.release()\n",
        "\n",
        "    # Get the top 3 predicted emotions\n",
        "    top_3_emotions = sorted(emotion_count.items(), key=lambda x: x[1], reverse=True)[:3]\n",
        "\n",
        "    return top_3_emotions\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def main(video_path, n):\n",
        "  device = 'cpu'\n",
        "  emotion_dict = {'angry': 0, 'disgust': 1, 'fear': 2, 'happy': 3, 'neutral': 4, 'sad': 5, 'surprise': 6}\n",
        "  songDF = pd.read_csv(\"/content/drive/Shareddrives/AML Project/musiccaps-public.csv\")\n",
        "  LM = LanguageModel()\n",
        "  EM = EmotionModel('/content/drive/Shareddrives/AML Project/CV_PART/vgg_modified.pth', emotion_dict, device)\n",
        "  RS = RecommenderSystem(songDF, LanguageModel= LM, EmotionModel= EM)\n",
        "\n",
        "  print(RS.getRecsFromVid(video_path, n=n))"
      ],
      "metadata": {
        "id": "aN_I5-luJQyN"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if __name__ == '__main__':\n",
        "  vid_path = input(\"video path: \")\n",
        "  n = input(\"number of recommendations: \")\n",
        "  main(vid_path, n)"
      ],
      "metadata": {
        "id": "efm-R109zH1n"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}