sample/whisper_at_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cy9y5cdzM15J"
   },
   "source": [
    "## [Whisper-AT: Noise-Robust Automatic Speech Recognizers are Also Strong Audio Event Taggers]()\n",
    "\n",
    "This colab script contains a step-by-step tutorial on how to use Whisper-AT for joint automatic speech recognitiom (ASR) and audio tagging (AT).\n",
    "\n",
    "Please cite our paper if you find this repository useful.\n",
    "\n",
    "```\n",
    "@inproceedings{gong_whisperat,\n",
    "  author={Gong, Yuan and Khurana, Sameer and Karlinsky, Leonid and Glass, James},\n",
    "  title={Whisper-AT: Noise-Robust Automatic Speech Recognizers are Also Strong Audio Event Taggers},\n",
    "  year=2023,\n",
    "  booktitle={Proc. Interspeech 2023}\n",
    "}\n",
    "```\n",
    "For more information, please check https://github.com/YuanGongND/whisper-at"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6EA_0A7eOn7a"
   },
   "source": [
    "### Step 1. Install Whisper-AT Package\n",
    "\n",
    "We intentionally do not any additional dependencies to the original Whisper. So if your environment can run the original Whisper, it must can also run Whisper-AT. Note that following original Whisper, it also requires the command-line tool [`ffmpeg`](https://ffmpeg.org/) to be installed on your system. Please check OpenAI Whisper repo for details.\n",
    "\n",
    "\n",
    "Whisper-AT can be installed simply by `pip install whisper-at`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "R5P5xbXwRjf3",
    "outputId": "db524cca-53ac-4aea-9835-b05a53fc1c65"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Collecting whisper-at\n",
      "  Downloading whisper_at-0.2-py3-none-any.whl (1.2 MB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m1.2/1.2 MB\u001B[0m \u001B[31m4.3 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25hRequirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from whisper-at) (2.0.0)\n",
      "Requirement already satisfied: numba in /usr/local/lib/python3.10/dist-packages (from whisper-at) (0.56.4)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from whisper-at) (1.22.4)\n",
      "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from whisper-at) (2.0.1+cu118)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from whisper-at) (4.65.0)\n",
      "Requirement already satisfied: more-itertools in /usr/local/lib/python3.10/dist-packages (from whisper-at) (9.1.0)\n",
      "Collecting tiktoken==0.3.3 (from whisper-at)\n",
      "  Downloading tiktoken-0.3.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.7 MB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m1.7/1.7 MB\u001B[0m \u001B[31m85.1 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25hRequirement already satisfied: regex>=2022.1.18 in /usr/local/lib/python3.10/dist-packages (from tiktoken==0.3.3->whisper-at) (2022.10.31)\n",
      "Requirement already satisfied: requests>=2.26.0 in /usr/local/lib/python3.10/dist-packages (from tiktoken==0.3.3->whisper-at) (2.27.1)\n",
      "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->whisper-at) (3.25.2)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->whisper-at) (3.12.2)\n",
      "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->whisper-at) (16.0.6)\n",
      "Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /usr/local/lib/python3.10/dist-packages (from numba->whisper-at) (0.39.1)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from numba->whisper-at) (67.7.2)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch->whisper-at) (4.6.3)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->whisper-at) (1.11.1)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->whisper-at) (3.1)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->whisper-at) (3.1.2)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken==0.3.3->whisper-at) (1.26.16)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken==0.3.3->whisper-at) (2023.5.7)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken==0.3.3->whisper-at) (2.0.12)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken==0.3.3->whisper-at) (3.4)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->whisper-at) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->whisper-at) (1.3.0)\n",
      "Installing collected packages: tiktoken, whisper-at\n",
      "Successfully installed tiktoken-0.3.3 whisper-at-0.2\n"
     ]
    }
   ],
   "source": [
    "!pip install whisper-at"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xN9r0mMoPJQl"
   },
   "source": [
    "### Step 2. Use as the Original Whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 76
    },
    "id": "r8Or_Q7sTywM",
    "outputId": "1a6e84bb-a7fa-44f9-b6ed-3dd358e39e9b"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  Preparing metadata (setup.py) ... \u001B[?25l\u001B[?25hdone\n",
      "  Building wheel for wget (setup.py) ... \u001B[?25l\u001B[?25hdone\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'/content/sample_audio.flac'"
      ],
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      }
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "source": [
    "# download a sample audio\n",
    "!pip -q install wget\n",
    "import wget,IPython\n",
    "wget.download('https://www.dropbox.com/s/7eznyazmc1pmw9h/case_closed.wav?dl=1', '/content/sample_audio.flac')\n",
    "#IPython.display.Audio('/content/sample_audio.flac')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GKAjB8AZUKcF",
    "outputId": "b00ed09b-8a2f-49e0-aaaa-5f95c3e3e53b"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|█████████████████████████████████████| 2.87G/2.87G [01:04<00:00, 47.8MiB/s]\n",
      "100%|███████████████████████████████████████| 153M/153M [00:05<00:00, 27.4MiB/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "0.0 s- 5.0 s  My name's Jimmy Kudo, and I've always wanted to be a great detective.\n",
      "5.0 s- 10.0 s  Over the years I've worked at perfecting the craft, impressing almost everyone, except maybe Rachel.\n",
      "10.0 s- 14.0 s  But then some bad guys ambushed me, slipping me some kind of poison.\n",
      "14.0 s- 16.0 s  Now what's gonna happen?\n",
      "16.0 s- 30.0 s  It was the first new century in one hundred years.\n",
      "30.0 s- 36.0 s  And when I felt like I should cry, I laughed away my tears.\n",
      "36.0 s- 48.0 s  At the end of a millennium, we waited a long, long time,\n",
      "48.0 s- 54.0 s  To see the brave new world, and the mountains we would climb.\n",
      "54.0 s- 60.0 s  Things I tried to comprehend as a child were made of mystery.\n",
      "60.0 s- 67.0 s  There's nothing I need to defend, there's nothing great about me.\n",
      "67.0 s- 73.0 s  All I will ever believe is the pounding of my heart, though.\n",
      "73.0 s- 79.0 s  It doesn't answer questions, that's just the way it goes.\n",
      "79.0 s- 84.0 s  All I will ever have bacon is the beating in my heart.\n",
      "90.0 s- 95.0 s  Thanks for watching!\n"
     ]
    }
   ],
   "source": [
    "# note this is whisper\"_\"at not whisper-at\n",
    "import whisper_at as whisper\n",
    "\n",
    "# the only new thing in whisper-at\n",
    "# specify the temporal resolution for audio tagging, 10 means Whisper-AT predict audio event every 10 seconds (hop and window=10s).\n",
    "audio_tagging_time_resolution = 10\n",
    "\n",
    "model = whisper.load_model(\"large-v1\")\n",
    "# for large, medium, small models, we provide low-dim proj AT models to save compute.\n",
    "# model = whisper.load_model(\"large-v1\", at_low_compute=Ture)\n",
    "result = model.transcribe(\"/content/sample_audio.flac\", at_time_res=audio_tagging_time_resolution)\n",
    "for segment in result['segments']:\n",
    "  print(segment['start'], 's-', segment['end'], 's', segment['text'])\n",
    "\n",
    "# # translation task is also supported\n",
    "# result = model.transcribe(\"/content/sample_audio.flac\", task='translate', at_time_res=audio_tagging_time_resolution)\n",
    "# print(result[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r-YoQNujPb50"
   },
   "source": [
    "`result[\"text\"]` is the ASR output transcripts, it will be identical to that of the original Whisper and is not impacted by `at_time_res`, the ASR function still follows Whisper's 30 second window. `at_time_res` is only related to audio tagging.\n",
    "\n",
    "Compared to the original Whisper, the only new thing is `at_time_res`, which is the hop and window size for Whisper-AT to predict audio events. For example, for a 60-second audio, setting `at_time_res = 10` means the audio will be segmented to 6 10-second segments, and Whisper-AT will predict audio tags based on each 10-second segment,\n",
    "a total of 6 audio event predictions will be made. **Note `at_time_res` must be an integer multiple of 0.4, e.g., 0.4, 0.8, ...**, the default value is 10.0, which is the value we use to train the model and should lead to best performance.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lNFynV7CRYHk"
   },
   "source": [
    "### Step 3. Get the Audio Tagging Output\n",
    "\n",
    "Compared with the original Whisper, `result` contains a new entry called `audio_tag`. `result['audio_tag']` is a torch tensor of shape [⌈`audio_length`/`at_time_res`⌉, 527]. For example, for a 60-second audio and `at_time_res = 10`, `result['audio_tag']` is a tensor of shape [6, 527]. 527 is the size of the [AudioSet label set](), `result['audio_tag'][i,j]` is the (unnormalised) logits of class `j` of the `i`th segment.\n",
    "\n",
    "If you are familiar with audio tagging and AudioSet, you can take raw `result['audio_tag']` for your usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TA5-Sw8zRwsB",
    "outputId": "454b5c30-54cb-42f0-d89e-ebab20a74c58"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Audio length is 90.05, at time resolution is 10.0, Whisper-AT output in shape torch.Size([10, 527])\n"
     ]
    }
   ],
   "source": [
    "import torchaudio\n",
    "audio, sr = torchaudio.load('/content/sample_audio.flac')\n",
    "audio_len = audio.shape[1] / sr\n",
    "print('Audio length is {:.2f}, at time resolution is {:.1f}, Whisper-AT output in shape'.format(audio_len, audio_tagging_time_resolution), result['audio_tag'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cskn2R_7RuSw"
   },
   "source": [
    "But we also provide a tool to make it easier.\n",
    "You can feed the `result` to `whisepr.parse_at_label` and get readable results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EtkoNNKHUXel",
    "outputId": "608202a3-6d37-419e-a927-fdd4bd884d1f"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'time': {'start': 0, 'end': 10}, 'audio tags': [('Music', 1.8221180438995361), ('Speech', 0.932380735874176)]}\n",
      "{'time': {'start': 10, 'end': 20}, 'audio tags': [('Music', 1.3558011054992676), ('Grunge', -1.350265622138977), ('Progressive rock', -1.424497127532959), ('Punk rock', -1.5711394548416138)]}\n",
      "{'time': {'start': 20, 'end': 30}, 'audio tags': [('Music', 0.8049014806747437)]}\n",
      "{'time': {'start': 30, 'end': 40}, 'audio tags': [('Music', 1.323777437210083), ('Grunge', -1.3895658254623413), ('Progressive rock', -1.614871859550476), ('Punk rock', -1.6720657348632812), ('Independent music', -1.8473705053329468)]}\n",
      "{'time': {'start': 40, 'end': 50}, 'audio tags': [('Music', 0.7121016383171082), ('Singing', -1.492790699005127)]}\n",
      "{'time': {'start': 50, 'end': 60}, 'audio tags': [('Music', -1.4463484287261963)]}\n",
      "{'time': {'start': 60, 'end': 70}, 'audio tags': [('Music', 1.2629746198654175), ('Grunge', -0.9662912487983704), ('Pop music', -1.6864861249923706), ('Independent music', -1.7853952646255493)]}\n",
      "{'time': {'start': 70, 'end': 80}, 'audio tags': [('Music', 0.7769227623939514), ('Grunge', -1.5903935432434082), ('Singing', -1.6665356159210205)]}\n",
      "{'time': {'start': 80, 'end': 90}, 'audio tags': [('Speech', -1.0284314155578613), ('Silence', -1.3325611352920532), ('Music', -1.5024266242980957)]}\n",
      "{'time': {'start': 90, 'end': 100}, 'audio tags': [('Silence', 0.9160257577896118)]}\n"
     ]
    }
   ],
   "source": [
    "audio_tag_result = whisper.parse_at_label(result, language='follow_asr', top_k=5, p_threshold=-2, include_class_list=list(range(527)))\n",
    "for segment in audio_tag_result:\n",
    "  print(segment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0RwNK7QdduDu"
   },
   "source": [
    "If you change the audio tagging resolution to 2s, then the output will be more fine-grained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zv8kDJ9uduex",
    "outputId": "7d456ab4-c92f-4dd1-a746-d0c120990edb"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "<ipython-input-6-c709aa3dc946>:2: UserWarning: Current at_time_res is 2.00 second, the audio tagging model is trained with time resolution of 10 seconds. Mismatch time resolution may cause an audio tagging performance drop, but won't impact ASR performance.\n",
      "  result = model.transcribe(\"/content/sample_audio.flac\", at_time_res=audio_tagging_time_resolution)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Audio length is 90.05, at time resolution is 2.0, Whisper-AT output in shape torch.Size([46, 527])\n",
      "{'time': {'start': 0, 'end': 2}, 'audio tags': [('Music', -0.4525575637817383), ('Speech', -0.8882830739021301), ('Silence', -1.8616548776626587)]}\n",
      "{'time': {'start': 2, 'end': 4}, 'audio tags': [('Music', 0.9328102469444275), ('Speech', -0.158194437623024)]}\n",
      "{'time': {'start': 4, 'end': 6}, 'audio tags': [('Music', 0.7153851389884949), ('Speech', 0.38911065459251404)]}\n",
      "{'time': {'start': 6, 'end': 8}, 'audio tags': [('Music', 0.40575844049453735), ('Speech', -0.6977154612541199)]}\n",
      "{'time': {'start': 8, 'end': 10}, 'audio tags': [('Music', 1.6714180707931519), ('Speech', 0.5128814578056335), ('Drum kit', -1.7371801137924194), ('Drum', -1.9375439882278442)]}\n",
      "{'time': {'start': 10, 'end': 12}, 'audio tags': [('Music', 1.3964262008666992), ('Speech', -0.23263752460479736)]}\n",
      "{'time': {'start': 12, 'end': 14}, 'audio tags': [('Music', 0.7910271883010864), ('Speech', -1.9267847537994385)]}\n",
      "{'time': {'start': 14, 'end': 16}, 'audio tags': [('Music', 1.018101692199707), ('Speech', -1.863593339920044)]}\n",
      "{'time': {'start': 16, 'end': 18}, 'audio tags': [('Music', 1.2444162368774414), ('Rock and roll', -1.106768012046814)]}\n",
      "{'time': {'start': 18, 'end': 20}, 'audio tags': [('Music', 1.1749851703643799), ('Rock and roll', -1.696586012840271)]}\n",
      "{'time': {'start': 20, 'end': 22}, 'audio tags': [('Music', 0.884811282157898)]}\n",
      "{'time': {'start': 22, 'end': 24}, 'audio tags': [('Music', 1.3546898365020752), ('Rock and roll', -1.7012099027633667), ('Speech', -1.882798671722412)]}\n",
      "{'time': {'start': 24, 'end': 26}, 'audio tags': [('Music', 1.7713080644607544), ('Rock and roll', -1.4794251918792725), ('Grunge', -1.6708897352218628), ('Rock music', -1.7350000143051147), ('Punk rock', -1.7468372583389282)]}\n",
      "{'time': {'start': 26, 'end': 28}, 'audio tags': [('Music', 1.2942352294921875), ('Punk rock', -1.5345808267593384), ('Progressive rock', -1.5366796255111694), ('Rock and roll', -1.795093059539795), ('Guitar', -1.9063055515289307)]}\n",
      "{'time': {'start': 28, 'end': 30}, 'audio tags': [('Music', 1.0816413164138794), ('Punk rock', -1.1613421440124512), ('Grunge', -1.1674290895462036), ('Progressive rock', -1.6921743154525757), ('Rock music', -1.8889744281768799)]}\n",
      "{'time': {'start': 30, 'end': 32}, 'audio tags': [('Music', 0.7130565047264099), ('Grunge', -1.7496439218521118), ('Psychedelic rock', -1.9967429637908936)]}\n",
      "{'time': {'start': 32, 'end': 34}, 'audio tags': [('Music', 0.9006908535957336)]}\n",
      "{'time': {'start': 34, 'end': 36}, 'audio tags': [('Music', 1.1229232549667358)]}\n",
      "{'time': {'start': 36, 'end': 38}, 'audio tags': [('Music', 1.5627844333648682), ('Rock and roll', -1.616320013999939), ('Pop music', -1.860297441482544)]}\n",
      "{'time': {'start': 38, 'end': 40}, 'audio tags': [('Music', 1.4832371473312378), ('Rock and roll', -1.5093891620635986)]}\n",
      "{'time': {'start': 40, 'end': 42}, 'audio tags': [('Music', 1.1807818412780762), ('Rock and roll', -1.8162614107131958)]}\n",
      "{'time': {'start': 42, 'end': 44}, 'audio tags': [('Music', 1.5886766910552979), ('Rock and roll', -1.6210857629776)]}\n",
      "{'time': {'start': 44, 'end': 46}, 'audio tags': [('Music', 1.3883097171783447), ('Punk rock', -1.4467766284942627), ('Rock and roll', -1.6486458778381348), ('Grunge', -1.8386918306350708)]}\n",
      "{'time': {'start': 46, 'end': 48}, 'audio tags': [('Music', 1.3570588827133179), ('Grunge', -1.6415327787399292), ('Independent music', -1.6770530939102173)]}\n",
      "{'time': {'start': 48, 'end': 50}, 'audio tags': [('Music', 1.2322897911071777), ('Progressive rock', -1.670889973640442), ('Punk rock', -1.7552521228790283), ('Grunge', -1.9851617813110352)]}\n",
      "{'time': {'start': 50, 'end': 52}, 'audio tags': [('Music', 0.9578625559806824), ('Musical instrument', -1.7584046125411987), ('Grunge', -1.7898246049880981), ('Guitar', -1.9311943054199219)]}\n",
      "{'time': {'start': 52, 'end': 54}, 'audio tags': [('Music', 0.40247398614883423)]}\n",
      "{'time': {'start': 54, 'end': 56}, 'audio tags': [('Music', 0.35136619210243225)]}\n",
      "{'time': {'start': 56, 'end': 58}, 'audio tags': [('Music', 1.121218204498291), ('Singing', -0.7474337220191956), ('Guitar', -1.1709648370742798), ('Musical instrument', -1.8700003623962402), ('Plucked string instrument', -1.8718571662902832)]}\n",
      "{'time': {'start': 58, 'end': 60}, 'audio tags': [('Music', 0.48927149176597595), ('Singing', -1.516997218132019)]}\n",
      "{'time': {'start': 60, 'end': 62}, 'audio tags': [('Music', 1.4061903953552246), ('Grunge', -1.6520559787750244), ('Pop music', -1.7633814811706543), ('Singing', -1.8479968309402466)]}\n",
      "{'time': {'start': 62, 'end': 64}, 'audio tags': [('Music', 1.0178853273391724), ('Pop music', -1.6401900053024292), ('Rhythm and blues', -1.9780902862548828)]}\n",
      "{'time': {'start': 64, 'end': 66}, 'audio tags': [('Music', 1.2572810649871826), ('Grunge', -1.1667473316192627)]}\n",
      "{'time': {'start': 66, 'end': 68}, 'audio tags': [('Music', 1.0675106048583984), ('Grunge', -1.7898226976394653)]}\n",
      "{'time': {'start': 68, 'end': 70}, 'audio tags': [('Music', 1.1966230869293213), ('Grunge', -1.2361091375350952), ('Independent music', -1.3939752578735352)]}\n",
      "{'time': {'start': 70, 'end': 72}, 'audio tags': [('Music', 1.193831205368042), ('Song', -1.9588038921356201)]}\n",
      "{'time': {'start': 72, 'end': 74}, 'audio tags': [('Music', 0.8933382034301758), ('Grunge', -1.828943133354187)]}\n",
      "{'time': {'start': 74, 'end': 76}, 'audio tags': [('Music', 1.2897837162017822), ('Singing', -1.9028337001800537)]}\n",
      "{'time': {'start': 76, 'end': 78}, 'audio tags': [('Music', 0.4346665143966675)]}\n",
      "{'time': {'start': 78, 'end': 80}, 'audio tags': [('Music', 0.6879321932792664), ('Grunge', -1.932247281074524)]}\n",
      "{'time': {'start': 80, 'end': 82}, 'audio tags': [('Music', 0.5806633830070496), ('Singing', -1.524978518486023)]}\n",
      "{'time': {'start': 82, 'end': 84}, 'audio tags': [('Music', 0.3592168986797333), ('Speech', -1.9439775943756104)]}\n",
      "{'time': {'start': 84, 'end': 86}, 'audio tags': [('Silence', 0.6775252819061279), ('Speech', -1.3987550735473633)]}\n",
      "{'time': {'start': 86, 'end': 88}, 'audio tags': [('Silence', -0.42543643712997437), ('Speech', -0.682219386100769)]}\n",
      "{'time': {'start': 88, 'end': 90}, 'audio tags': [('Speech', -0.7888064980506897), ('Silence', -1.000313401222229), ('Sigh', -1.4150739908218384), ('Music', -1.4370006322860718), ('Stomach rumble', -1.5969691276550293)]}\n",
      "{'time': {'start': 90, 'end': 92}, 'audio tags': [('Silence', 1.4236705303192139)]}\n"
     ]
    }
   ],
   "source": [
    "audio_tagging_time_resolution = 2\n",
    "result = model.transcribe(\"/content/sample_audio.flac\", at_time_res=audio_tagging_time_resolution)\n",
    "print('Audio length is {:.2f}, at time resolution is {:.1f}, Whisper-AT output in shape'.format(audio_len, audio_tagging_time_resolution), result['audio_tag'].shape)\n",
    "audio_tag_result = whisper.parse_at_label(result, language='follow_asr', top_k=5, p_threshold=-2, include_class_list=list(range(527)))\n",
    "for segment in audio_tag_result:\n",
    "  print(segment)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "qh3i2ta_euUL"
   },
   "outputs": [],
   "source": [
    "# Go back to 10s for better readability\n",
    "audio_tagging_time_resolution = 10\n",
    "result = model.transcribe(\"/content/sample_audio.flac\", at_time_res=audio_tagging_time_resolution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U8yK4hASW165"
   },
   "source": [
    "Let's take a closer look at `whisper.parse_at_label`.\n",
    "\n",
    "First, `top_k` and `p_threshold` controls how many audio tags are output. Specifically, `whisper.parse_at_label` will output up to `k` labels that have unnormalised logits above `p_threshold`.\n",
    "\n",
    "For example, set `top_k` = 1 allows the model to output at most 1 label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "y33gPUHRW1TN",
    "outputId": "670edb75-f547-4f6f-8dfb-d91b6a129531"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'time': {'start': 0, 'end': 10}, 'audio tags': [('Music', 1.8221180438995361)]}\n",
      "{'time': {'start': 10, 'end': 20}, 'audio tags': [('Music', 1.3558011054992676)]}\n",
      "{'time': {'start': 20, 'end': 30}, 'audio tags': [('Music', 0.8049014806747437)]}\n",
      "{'time': {'start': 30, 'end': 40}, 'audio tags': [('Music', 1.323777437210083)]}\n",
      "{'time': {'start': 40, 'end': 50}, 'audio tags': [('Music', 0.7121016383171082)]}\n",
      "{'time': {'start': 50, 'end': 60}, 'audio tags': [('Music', -1.4463484287261963)]}\n",
      "{'time': {'start': 60, 'end': 70}, 'audio tags': [('Music', 1.2629746198654175)]}\n",
      "{'time': {'start': 70, 'end': 80}, 'audio tags': [('Music', 0.7769227623939514)]}\n",
      "{'time': {'start': 80, 'end': 90}, 'audio tags': [('Speech', -1.0284314155578613)]}\n",
      "{'time': {'start': 90, 'end': 100}, 'audio tags': [('Silence', 0.9160257577896118)]}\n"
     ]
    }
   ],
   "source": [
    "audio_tag_result = whisper.parse_at_label(result, language='follow_asr', top_k=1, p_threshold=-2, include_class_list=list(range(527)))\n",
    "for segment in audio_tag_result:\n",
    "  print(segment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5XNWSMJ_XcRN"
   },
   "source": [
    "Setting larger `top_k` and smaller `p_threshold` makes the model more verbose, and vise-versa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-EgkFCP4XmEx",
    "outputId": "d97a3a6f-71cb-43b2-daaa-1bf180ec5b40"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'time': {'start': 0, 'end': 10}, 'audio tags': []}\n",
      "{'time': {'start': 10, 'end': 20}, 'audio tags': []}\n",
      "{'time': {'start': 20, 'end': 30}, 'audio tags': []}\n",
      "{'time': {'start': 30, 'end': 40}, 'audio tags': []}\n",
      "{'time': {'start': 40, 'end': 50}, 'audio tags': []}\n",
      "{'time': {'start': 50, 'end': 60}, 'audio tags': []}\n",
      "{'time': {'start': 60, 'end': 70}, 'audio tags': []}\n",
      "{'time': {'start': 70, 'end': 80}, 'audio tags': []}\n",
      "{'time': {'start': 80, 'end': 90}, 'audio tags': []}\n",
      "{'time': {'start': 90, 'end': 100}, 'audio tags': []}\n"
     ]
    }
   ],
   "source": [
    "audio_tag_result = whisper.parse_at_label(result, language='follow_asr', top_k=10, p_threshold=-5, include_class_list=[47])\n",
    "for segment in audio_tag_result:\n",
    "  print(segment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nA3Imgy1X3Hr"
   },
   "source": [
    "Second, you can also select the classes of interest by inputting a list to `include_class_list`. For the name-to-index mappling, simply let Whisper-AT print it for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EkqfZmDIYS73",
    "outputId": "6a05327f-6562-4bcc-9369-ff41e3f421d3"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "index: 0 : Speech\n",
      "index: 1 : Male speech, man speaking\n",
      "index: 2 : Female speech, woman speaking\n",
      "index: 3 : Child speech, kid speaking\n",
      "index: 4 : Conversation\n",
      "index: 5 : Narration, monologue\n",
      "index: 6 : Babbling\n",
      "index: 7 : Speech synthesizer\n",
      "index: 8 : Shout\n",
      "index: 9 : Bellow\n",
      "index: 10 : Whoop\n",
      "index: 11 : Yell\n",
      "index: 12 : Battle cry\n",
      "index: 13 : Children shouting\n",
      "index: 14 : Screaming\n",
      "index: 15 : Whispering\n",
      "index: 16 : Laughter\n",
      "index: 17 : Baby laughter\n",
      "index: 18 : Giggle\n",
      "index: 19 : Snicker\n",
      "index: 20 : Belly laugh\n",
      "index: 21 : Chuckle, chortle\n",
      "index: 22 : Crying, sobbing\n",
      "index: 23 : Baby cry, infant cry\n",
      "index: 24 : Whimper\n",
      "index: 25 : Wail, moan\n",
      "index: 26 : Sigh\n",
      "index: 27 : Singing\n",
      "index: 28 : Choir\n",
      "index: 29 : Yodeling\n",
      "index: 30 : Chant\n",
      "index: 31 : Mantra\n",
      "index: 32 : Male singing\n",
      "index: 33 : Female singing\n",
      "index: 34 : Child singing\n",
      "index: 35 : Synthetic singing\n",
      "index: 36 : Rapping\n",
      "index: 37 : Humming\n",
      "index: 38 : Groan\n",
      "index: 39 : Grunt\n",
      "index: 40 : Whistling\n",
      "index: 41 : Breathing\n",
      "index: 42 : Wheeze\n",
      "index: 43 : Snoring\n",
      "index: 44 : Gasp\n",
      "index: 45 : Pant\n",
      "index: 46 : Snort\n",
      "index: 47 : Cough\n",
      "index: 48 : Throat clearing\n",
      "index: 49 : Sneeze\n",
      "index: 50 : Sniff\n",
      "index: 51 : Run\n",
      "index: 52 : Shuffle\n",
      "index: 53 : Walk, footsteps\n",
      "index: 54 : Chewing, mastication\n",
      "index: 55 : Biting\n",
      "index: 56 : Gargling\n",
      "index: 57 : Stomach rumble\n",
      "index: 58 : Burping, eructation\n",
      "index: 59 : Hiccup\n",
      "index: 60 : Fart\n",
      "index: 61 : Hands\n",
      "index: 62 : Finger snapping\n",
      "index: 63 : Clapping\n",
      "index: 64 : Heart sounds, heartbeat\n",
      "index: 65 : Heart murmur\n",
      "index: 66 : Cheering\n",
      "index: 67 : Applause\n",
      "index: 68 : Chatter\n",
      "index: 69 : Crowd\n",
      "index: 70 : Hubbub, speech noise, speech babble\n",
      "index: 71 : Children playing\n",
      "index: 72 : Animal\n",
      "index: 73 : Domestic animals, pets\n",
      "index: 74 : Dog\n",
      "index: 75 : Bark\n",
      "index: 76 : Yip\n",
      "index: 77 : Howl\n",
      "index: 78 : Bow-wow\n",
      "index: 79 : Growling\n",
      "index: 80 : Whimper (dog)\n",
      "index: 81 : Cat\n",
      "index: 82 : Purr\n",
      "index: 83 : Meow\n",
      "index: 84 : Hiss\n",
      "index: 85 : Caterwaul\n",
      "index: 86 : Livestock, farm animals, working animals\n",
      "index: 87 : Horse\n",
      "index: 88 : Clip-clop\n",
      "index: 89 : Neigh, whinny\n",
      "index: 90 : Cattle, bovinae\n",
      "index: 91 : Moo\n",
      "index: 92 : Cowbell\n",
      "index: 93 : Pig\n",
      "index: 94 : Oink\n",
      "index: 95 : Goat\n",
      "index: 96 : Bleat\n",
      "index: 97 : Sheep\n",
      "index: 98 : Fowl\n",
      "index: 99 : Chicken, rooster\n",
      "index: 100 : Cluck\n",
      "index: 101 : Crowing, cock-a-doodle-doo\n",
      "index: 102 : Turkey\n",
      "index: 103 : Gobble\n",
      "index: 104 : Duck\n",
      "index: 105 : Quack\n",
      "index: 106 : Goose\n",
      "index: 107 : Honk\n",
      "index: 108 : Wild animals\n",
      "index: 109 : Roaring cats (lions, tigers)\n",
      "index: 110 : Roar\n",
      "index: 111 : Bird\n",
      "index: 112 : Bird vocalization, bird call, bird song\n",
      "index: 113 : Chirp, tweet\n",
      "index: 114 : Squawk\n",
      "index: 115 : Pigeon, dove\n",
      "index: 116 : Coo\n",
      "index: 117 : Crow\n",
      "index: 118 : Caw\n",
      "index: 119 : Owl\n",
      "index: 120 : Hoot\n",
      "index: 121 : Bird flight, flapping wings\n",
      "index: 122 : Canidae, dogs, wolves\n",
      "index: 123 : Rodents, rats, mice\n",
      "index: 124 : Mouse\n",
      "index: 125 : Patter\n",
      "index: 126 : Insect\n",
      "index: 127 : Cricket\n",
      "index: 128 : Mosquito\n",
      "index: 129 : Fly, housefly\n",
      "index: 130 : Buzz\n",
      "index: 131 : Bee, wasp, etc.\n",
      "index: 132 : Frog\n",
      "index: 133 : Croak\n",
      "index: 134 : Snake\n",
      "index: 135 : Rattle\n",
      "index: 136 : Whale vocalization\n",
      "index: 137 : Music\n",
      "index: 138 : Musical instrument\n",
      "index: 139 : Plucked string instrument\n",
      "index: 140 : Guitar\n",
      "index: 141 : Electric guitar\n",
      "index: 142 : Bass guitar\n",
      "index: 143 : Acoustic guitar\n",
      "index: 144 : Steel guitar, slide guitar\n",
      "index: 145 : Tapping (guitar technique)\n",
      "index: 146 : Strum\n",
      "index: 147 : Banjo\n",
      "index: 148 : Sitar\n",
      "index: 149 : Mandolin\n",
      "index: 150 : Zither\n",
      "index: 151 : Ukulele\n",
      "index: 152 : Keyboard (musical)\n",
      "index: 153 : Piano\n",
      "index: 154 : Electric piano\n",
      "index: 155 : Organ\n",
      "index: 156 : Electronic organ\n",
      "index: 157 : Hammond organ\n",
      "index: 158 : Synthesizer\n",
      "index: 159 : Sampler\n",
      "index: 160 : Harpsichord\n",
      "index: 161 : Percussion\n",
      "index: 162 : Drum kit\n",
      "index: 163 : Drum machine\n",
      "index: 164 : Drum\n",
      "index: 165 : Snare drum\n",
      "index: 166 : Rimshot\n",
      "index: 167 : Drum roll\n",
      "index: 168 : Bass drum\n",
      "index: 169 : Timpani\n",
      "index: 170 : Tabla\n",
      "index: 171 : Cymbal\n",
      "index: 172 : Hi-hat\n",
      "index: 173 : Wood block\n",
      "index: 174 : Tambourine\n",
      "index: 175 : Rattle (instrument)\n",
      "index: 176 : Maraca\n",
      "index: 177 : Gong\n",
      "index: 178 : Tubular bells\n",
      "index: 179 : Mallet percussion\n",
      "index: 180 : Marimba, xylophone\n",
      "index: 181 : Glockenspiel\n",
      "index: 182 : Vibraphone\n",
      "index: 183 : Steelpan\n",
      "index: 184 : Orchestra\n",
      "index: 185 : Brass instrument\n",
      "index: 186 : French horn\n",
      "index: 187 : Trumpet\n",
      "index: 188 : Trombone\n",
      "index: 189 : Bowed string instrument\n",
      "index: 190 : String section\n",
      "index: 191 : Violin, fiddle\n",
      "index: 192 : Pizzicato\n",
      "index: 193 : Cello\n",
      "index: 194 : Double bass\n",
      "index: 195 : Wind instrument, woodwind instrument\n",
      "index: 196 : Flute\n",
      "index: 197 : Saxophone\n",
      "index: 198 : Clarinet\n",
      "index: 199 : Harp\n",
      "index: 200 : Bell\n",
      "index: 201 : Church bell\n",
      "index: 202 : Jingle bell\n",
      "index: 203 : Bicycle bell\n",
      "index: 204 : Tuning fork\n",
      "index: 205 : Chime\n",
      "index: 206 : Wind chime\n",
      "index: 207 : Change ringing (campanology)\n",
      "index: 208 : Harmonica\n",
      "index: 209 : Accordion\n",
      "index: 210 : Bagpipes\n",
      "index: 211 : Didgeridoo\n",
      "index: 212 : Shofar\n",
      "index: 213 : Theremin\n",
      "index: 214 : Singing bowl\n",
      "index: 215 : Scratching (performance technique)\n",
      "index: 216 : Pop music\n",
      "index: 217 : Hip hop music\n",
      "index: 218 : Beatboxing\n",
      "index: 219 : Rock music\n",
      "index: 220 : Heavy metal\n",
      "index: 221 : Punk rock\n",
      "index: 222 : Grunge\n",
      "index: 223 : Progressive rock\n",
      "index: 224 : Rock and roll\n",
      "index: 225 : Psychedelic rock\n",
      "index: 226 : Rhythm and blues\n",
      "index: 227 : Soul music\n",
      "index: 228 : Reggae\n",
      "index: 229 : Country\n",
      "index: 230 : Swing music\n",
      "index: 231 : Bluegrass\n",
      "index: 232 : Funk\n",
      "index: 233 : Folk music\n",
      "index: 234 : Middle Eastern music\n",
      "index: 235 : Jazz\n",
      "index: 236 : Disco\n",
      "index: 237 : Classical music\n",
      "index: 238 : Opera\n",
      "index: 239 : Electronic music\n",
      "index: 240 : House music\n",
      "index: 241 : Techno\n",
      "index: 242 : Dubstep\n",
      "index: 243 : Drum and bass\n",
      "index: 244 : Electronica\n",
      "index: 245 : Electronic dance music\n",
      "index: 246 : Ambient music\n",
      "index: 247 : Trance music\n",
      "index: 248 : Music of Latin America\n",
      "index: 249 : Salsa music\n",
      "index: 250 : Flamenco\n",
      "index: 251 : Blues\n",
      "index: 252 : Music for children\n",
      "index: 253 : New-age music\n",
      "index: 254 : Vocal music\n",
      "index: 255 : A capella\n",
      "index: 256 : Music of Africa\n",
      "index: 257 : Afrobeat\n",
      "index: 258 : Christian music\n",
      "index: 259 : Gospel music\n",
      "index: 260 : Music of Asia\n",
      "index: 261 : Carnatic music\n",
      "index: 262 : Music of Bollywood\n",
      "index: 263 : Ska\n",
      "index: 264 : Traditional music\n",
      "index: 265 : Independent music\n",
      "index: 266 : Song\n",
      "index: 267 : Background music\n",
      "index: 268 : Theme music\n",
      "index: 269 : Jingle (music)\n",
      "index: 270 : Soundtrack music\n",
      "index: 271 : Lullaby\n",
      "index: 272 : Video game music\n",
      "index: 273 : Christmas music\n",
      "index: 274 : Dance music\n",
      "index: 275 : Wedding music\n",
      "index: 276 : Happy music\n",
      "index: 277 : Funny music\n",
      "index: 278 : Sad music\n",
      "index: 279 : Tender music\n",
      "index: 280 : Exciting music\n",
      "index: 281 : Angry music\n",
      "index: 282 : Scary music\n",
      "index: 283 : Wind\n",
      "index: 284 : Rustling leaves\n",
      "index: 285 : Wind noise (microphone)\n",
      "index: 286 : Thunderstorm\n",
      "index: 287 : Thunder\n",
      "index: 288 : Water\n",
      "index: 289 : Rain\n",
      "index: 290 : Raindrop\n",
      "index: 291 : Rain on surface\n",
      "index: 292 : Stream\n",
      "index: 293 : Waterfall\n",
      "index: 294 : Ocean\n",
      "index: 295 : Waves, surf\n",
      "index: 296 : Steam\n",
      "index: 297 : Gurgling\n",
      "index: 298 : Fire\n",
      "index: 299 : Crackle\n",
      "index: 300 : Vehicle\n",
      "index: 301 : Boat, Water vehicle\n",
      "index: 302 : Sailboat, sailing ship\n",
      "index: 303 : Rowboat, canoe, kayak\n",
      "index: 304 : Motorboat, speedboat\n",
      "index: 305 : Ship\n",
      "index: 306 : Motor vehicle (road)\n",
      "index: 307 : Car\n",
      "index: 308 : Vehicle horn, car horn, honking\n",
      "index: 309 : Toot\n",
      "index: 310 : Car alarm\n",
      "index: 311 : Power windows, electric windows\n",
      "index: 312 : Skidding\n",
      "index: 313 : Tire squeal\n",
      "index: 314 : Car passing by\n",
      "index: 315 : Race car, auto racing\n",
      "index: 316 : Truck\n",
      "index: 317 : Air brake\n",
      "index: 318 : Air horn, truck horn\n",
      "index: 319 : Reversing beeps\n",
      "index: 320 : Ice cream truck, ice cream van\n",
      "index: 321 : Bus\n",
      "index: 322 : Emergency vehicle\n",
      "index: 323 : Police car (siren)\n",
      "index: 324 : Ambulance (siren)\n",
      "index: 325 : Fire engine, fire truck (siren)\n",
      "index: 326 : Motorcycle\n",
      "index: 327 : Traffic noise, roadway noise\n",
      "index: 328 : Rail transport\n",
      "index: 329 : Train\n",
      "index: 330 : Train whistle\n",
      "index: 331 : Train horn\n",
      "index: 332 : Railroad car, train wagon\n",
      "index: 333 : Train wheels squealing\n",
      "index: 334 : Subway, metro, underground\n",
      "index: 335 : Aircraft\n",
      "index: 336 : Aircraft engine\n",
      "index: 337 : Jet engine\n",
      "index: 338 : Propeller, airscrew\n",
      "index: 339 : Helicopter\n",
      "index: 340 : Fixed-wing aircraft, airplane\n",
      "index: 341 : Bicycle\n",
      "index: 342 : Skateboard\n",
      "index: 343 : Engine\n",
      "index: 344 : Light engine (high frequency)\n",
      "index: 345 : Dental drill, dentist's drill\n",
      "index: 346 : Lawn mower\n",
      "index: 347 : Chainsaw\n",
      "index: 348 : Medium engine (mid frequency)\n",
      "index: 349 : Heavy engine (low frequency)\n",
      "index: 350 : Engine knocking\n",
      "index: 351 : Engine starting\n",
      "index: 352 : Idling\n",
      "index: 353 : Accelerating, revving, vroom\n",
      "index: 354 : Door\n",
      "index: 355 : Doorbell\n",
      "index: 356 : Ding-dong\n",
      "index: 357 : Sliding door\n",
      "index: 358 : Slam\n",
      "index: 359 : Knock\n",
      "index: 360 : Tap\n",
      "index: 361 : Squeak\n",
      "index: 362 : Cupboard open or close\n",
      "index: 363 : Drawer open or close\n",
      "index: 364 : Dishes, pots, and pans\n",
      "index: 365 : Cutlery, silverware\n",
      "index: 366 : Chopping (food)\n",
      "index: 367 : Frying (food)\n",
      "index: 368 : Microwave oven\n",
      "index: 369 : Blender\n",
      "index: 370 : Water tap, faucet\n",
      "index: 371 : Sink (filling or washing)\n",
      "index: 372 : Bathtub (filling or washing)\n",
      "index: 373 : Hair dryer\n",
      "index: 374 : Toilet flush\n",
      "index: 375 : Toothbrush\n",
      "index: 376 : Electric toothbrush\n",
      "index: 377 : Vacuum cleaner\n",
      "index: 378 : Zipper (clothing)\n",
      "index: 379 : Keys jangling\n",
      "index: 380 : Coin (dropping)\n",
      "index: 381 : Scissors\n",
      "index: 382 : Electric shaver, electric razor\n",
      "index: 383 : Shuffling cards\n",
      "index: 384 : Typing\n",
      "index: 385 : Typewriter\n",
      "index: 386 : Computer keyboard\n",
      "index: 387 : Writing\n",
      "index: 388 : Alarm\n",
      "index: 389 : Telephone\n",
      "index: 390 : Telephone bell ringing\n",
      "index: 391 : Ringtone\n",
      "index: 392 : Telephone dialing, DTMF\n",
      "index: 393 : Dial tone\n",
      "index: 394 : Busy signal\n",
      "index: 395 : Alarm clock\n",
      "index: 396 : Siren\n",
      "index: 397 : Civil defense siren\n",
      "index: 398 : Buzzer\n",
      "index: 399 : Smoke detector, smoke alarm\n",
      "index: 400 : Fire alarm\n",
      "index: 401 : Foghorn\n",
      "index: 402 : Whistle\n",
      "index: 403 : Steam whistle\n",
      "index: 404 : Mechanisms\n",
      "index: 405 : Ratchet, pawl\n",
      "index: 406 : Clock\n",
      "index: 407 : Tick\n",
      "index: 408 : Tick-tock\n",
      "index: 409 : Gears\n",
      "index: 410 : Pulleys\n",
      "index: 411 : Sewing machine\n",
      "index: 412 : Mechanical fan\n",
      "index: 413 : Air conditioning\n",
      "index: 414 : Cash register\n",
      "index: 415 : Printer\n",
      "index: 416 : Camera\n",
      "index: 417 : Single-lens reflex camera\n",
      "index: 418 : Tools\n",
      "index: 419 : Hammer\n",
      "index: 420 : Jackhammer\n",
      "index: 421 : Sawing\n",
      "index: 422 : Filing (rasp)\n",
      "index: 423 : Sanding\n",
      "index: 424 : Power tool\n",
      "index: 425 : Drill\n",
      "index: 426 : Explosion\n",
      "index: 427 : Gunshot, gunfire\n",
      "index: 428 : Machine gun\n",
      "index: 429 : Fusillade\n",
      "index: 430 : Artillery fire\n",
      "index: 431 : Cap gun\n",
      "index: 432 : Fireworks\n",
      "index: 433 : Firecracker\n",
      "index: 434 : Burst, pop\n",
      "index: 435 : Eruption\n",
      "index: 436 : Boom\n",
      "index: 437 : Wood\n",
      "index: 438 : Chop\n",
      "index: 439 : Splinter\n",
      "index: 440 : Crack\n",
      "index: 441 : Glass\n",
      "index: 442 : Chink, clink\n",
      "index: 443 : Shatter\n",
      "index: 444 : Liquid\n",
      "index: 445 : Splash, splatter\n",
      "index: 446 : Slosh\n",
      "index: 447 : Squish\n",
      "index: 448 : Drip\n",
      "index: 449 : Pour\n",
      "index: 450 : Trickle, dribble\n",
      "index: 451 : Gush\n",
      "index: 452 : Fill (with liquid)\n",
      "index: 453 : Spray\n",
      "index: 454 : Pump (liquid)\n",
      "index: 455 : Stir\n",
      "index: 456 : Boiling\n",
      "index: 457 : Sonar\n",
      "index: 458 : Arrow\n",
      "index: 459 : Whoosh, swoosh, swish\n",
      "index: 460 : Thump, thud\n",
      "index: 461 : Thunk\n",
      "index: 462 : Electronic tuner\n",
      "index: 463 : Effects unit\n",
      "index: 464 : Chorus effect\n",
      "index: 465 : Basketball bounce\n",
      "index: 466 : Bang\n",
      "index: 467 : Slap, smack\n",
      "index: 468 : Whack, thwack\n",
      "index: 469 : Smash, crash\n",
      "index: 470 : Breaking\n",
      "index: 471 : Bouncing\n",
      "index: 472 : Whip\n",
      "index: 473 : Flap\n",
      "index: 474 : Scratch\n",
      "index: 475 : Scrape\n",
      "index: 476 : Rub\n",
      "index: 477 : Roll\n",
      "index: 478 : Crushing\n",
      "index: 479 : Crumpling, crinkling\n",
      "index: 480 : Tearing\n",
      "index: 481 : Beep, bleep\n",
      "index: 482 : Ping\n",
      "index: 483 : Ding\n",
      "index: 484 : Clang\n",
      "index: 485 : Squeal\n",
      "index: 486 : Creak\n",
      "index: 487 : Rustle\n",
      "index: 488 : Whir\n",
      "index: 489 : Clatter\n",
      "index: 490 : Sizzle\n",
      "index: 491 : Clicking\n",
      "index: 492 : Clickety-clack\n",
      "index: 493 : Rumble\n",
      "index: 494 : Plop\n",
      "index: 495 : Jingle, tinkle\n",
      "index: 496 : Hum\n",
      "index: 497 : Zing\n",
      "index: 498 : Boing\n",
      "index: 499 : Crunch\n",
      "index: 500 : Silence\n",
      "index: 501 : Sine wave\n",
      "index: 502 : Harmonic\n",
      "index: 503 : Chirp tone\n",
      "index: 504 : Sound effect\n",
      "index: 505 : Pulse\n",
      "index: 506 : Inside, small room\n",
      "index: 507 : Inside, large room or hall\n",
      "index: 508 : Inside, public space\n",
      "index: 509 : Outside, urban or manmade\n",
      "index: 510 : Outside, rural or natural\n",
      "index: 511 : Reverberation\n",
      "index: 512 : Echo\n",
      "index: 513 : Noise\n",
      "index: 514 : Environmental noise\n",
      "index: 515 : Static\n",
      "index: 516 : Mains hum\n",
      "index: 517 : Distortion\n",
      "index: 518 : Sidetone\n",
      "index: 519 : Cacophony\n",
      "index: 520 : White noise\n",
      "index: 521 : Pink noise\n",
      "index: 522 : Throbbing\n",
      "index: 523 : Vibration\n",
      "index: 524 : Television\n",
      "index: 525 : Radio\n",
      "index: 526 : Field recording\n"
     ]
    }
   ],
   "source": [
    "whisper.print_label_name(language='en')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OJCbSma4YkaM"
   },
   "source": [
    "Assume we only interested in class 0, 1, 2 (Speech). We can let Whisper-AT only output these classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TK_KI5oeYfQP",
    "outputId": "44aca4ae-8f6c-4b59-a124-00ed8f818cd1"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'time': {'start': 0, 'end': 10}, 'audio tags': [('Speech', 0.932380735874176), ('Male speech, man speaking', -3.2922842502593994)]}\n",
      "{'time': {'start': 10, 'end': 20}, 'audio tags': []}\n",
      "{'time': {'start': 20, 'end': 30}, 'audio tags': []}\n",
      "{'time': {'start': 30, 'end': 40}, 'audio tags': []}\n",
      "{'time': {'start': 40, 'end': 50}, 'audio tags': []}\n",
      "{'time': {'start': 50, 'end': 60}, 'audio tags': []}\n",
      "{'time': {'start': 60, 'end': 70}, 'audio tags': []}\n",
      "{'time': {'start': 70, 'end': 80}, 'audio tags': []}\n",
      "{'time': {'start': 80, 'end': 90}, 'audio tags': [('Speech', -1.0284314155578613)]}\n",
      "{'time': {'start': 90, 'end': 100}, 'audio tags': []}\n"
     ]
    }
   ],
   "source": [
    "audio_tag_result = whisper.parse_at_label(result, language='follow_asr', top_k=10, p_threshold=-5, include_class_list=[0, 1, 2])\n",
    "for segment in audio_tag_result:\n",
    "  print(segment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "366Sq2ozY1jc"
   },
   "source": [
    "Finally, Whisper-AT support multiple languages. The default behavior is to output audio tag label names in the same language as ASR transcripts (i.e., `follow-asr`). But you can specify any supported language. Check supported language by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "h3dR5kjtZGI1",
    "outputId": "96d3066b-8689-4231-9e81-718f40c5f0a1"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "language code: en : english\n",
      "language code: zh : chinese\n",
      "language code: de : german\n",
      "language code: es : spanish\n",
      "language code: ru : russian\n",
      "language code: ko : korean\n",
      "language code: fr : french\n",
      "language code: ja : japanese\n",
      "language code: pt : portuguese\n",
      "language code: tr : turkish\n",
      "language code: pl : polish\n",
      "language code: ca : catalan\n",
      "language code: nl : dutch\n",
      "language code: ar : arabic\n",
      "language code: sv : swedish\n",
      "language code: it : italian\n",
      "language code: id : indonesian\n",
      "language code: hi : hindi\n",
      "language code: fi : finnish\n",
      "language code: vi : vietnamese\n",
      "language code: he : hebrew\n",
      "language code: uk : ukrainian\n",
      "language code: el : greek\n",
      "language code: ms : malay\n",
      "language code: cs : czech\n",
      "language code: ro : romanian\n",
      "language code: da : danish\n",
      "language code: hu : hungarian\n",
      "language code: ta : tamil\n",
      "language code: no : norwegian\n",
      "language code: th : thai\n",
      "language code: ur : urdu\n",
      "language code: hr : croatian\n",
      "language code: bg : bulgarian\n",
      "language code: lt : lithuanian\n",
      "language code: mi : maori\n",
      "language code: ml : malayalam\n",
      "language code: cy : welsh\n",
      "language code: sk : slovak\n",
      "language code: te : telugu\n",
      "language code: fa : persian\n",
      "language code: lv : latvian\n",
      "language code: bn : bengali\n",
      "language code: sr : serbian\n",
      "language code: az : azerbaijani\n",
      "language code: sl : slovenian\n",
      "language code: kn : kannada\n",
      "language code: mk : macedonian\n",
      "language code: eu : basque\n",
      "language code: is : icelandic\n",
      "language code: hy : armenian\n",
      "language code: ne : nepali\n",
      "language code: mn : mongolian\n",
      "language code: bs : bosnian\n",
      "language code: kk : kazakh\n",
      "language code: sq : albanian\n",
      "language code: sw : swahili\n",
      "language code: gl : galician\n",
      "language code: mr : marathi\n",
      "language code: pa : punjabi\n",
      "language code: si : sinhala\n",
      "language code: km : khmer\n",
      "language code: sn : shona\n",
      "language code: yo : yoruba\n",
      "language code: so : somali\n",
      "language code: af : afrikaans\n",
      "language code: ka : georgian\n",
      "language code: sd : sindhi\n",
      "language code: gu : gujarati\n",
      "language code: am : amharic\n",
      "language code: lo : lao\n",
      "language code: uz : uzbek\n",
      "language code: fo : faroese\n",
      "language code: ht : haitian creole\n",
      "language code: ps : pashto\n",
      "language code: tk : turkmen\n",
      "language code: mt : maltese\n",
      "language code: my : myanmar\n",
      "language code: bo : tibetan\n",
      "language code: mg : malagasy\n",
      "language code: as : assamese\n",
      "language code: tt : tatar\n",
      "language code: ln : lingala\n",
      "language code: ha : hausa\n",
      "language code: ba : bashkir\n"
     ]
    }
   ],
   "source": [
    "whisper.print_support_language()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jZjKvvVDZOsQ"
   },
   "source": [
    "Let's say we want the output labels in Chinese (zh):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pAz5DdG_ZWqA",
    "outputId": "15dda663-a6e8-42d0-f771-7e7456cb80cd"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'time': {'start': 0, 'end': 10}, 'audio tags': [('音乐之声', 1.8221180438995361), ('说话的声音', 0.932380735874176)]}\n",
      "{'time': {'start': 10, 'end': 20}, 'audio tags': [('音乐之声', 1.3558011054992676), ('垃圾的声音', -1.350265622138977), ('进步摇滚的声音', -1.424497127532959), ('朋克摇滚之声', -1.5711394548416138)]}\n",
      "{'time': {'start': 20, 'end': 30}, 'audio tags': [('音乐之声', 0.8049014806747437)]}\n",
      "{'time': {'start': 30, 'end': 40}, 'audio tags': [('音乐之声', 1.323777437210083), ('垃圾的声音', -1.3895658254623413), ('进步摇滚的声音', -1.614871859550476), ('朋克摇滚之声', -1.6720657348632812), ('独立音乐之声', -1.8473705053329468)]}\n",
      "{'time': {'start': 40, 'end': 50}, 'audio tags': [('音乐之声', 0.7121016383171082), ('歌声', -1.492790699005127)]}\n",
      "{'time': {'start': 50, 'end': 60}, 'audio tags': [('音乐之声', -1.4463484287261963)]}\n",
      "{'time': {'start': 60, 'end': 70}, 'audio tags': [('音乐之声', 1.2629746198654175), ('垃圾的声音', -0.9662912487983704), ('流行音乐的声音', -1.6864861249923706), ('独立音乐之声', -1.7853952646255493)]}\n",
      "{'time': {'start': 70, 'end': 80}, 'audio tags': [('音乐之声', 0.7769227623939514), ('垃圾的声音', -1.5903935432434082), ('歌声', -1.6665356159210205)]}\n",
      "{'time': {'start': 80, 'end': 90}, 'audio tags': [('说话的声音', -1.0284314155578613), ('寂静的声音', -1.3325611352920532), ('音乐之声', -1.5024266242980957)]}\n",
      "{'time': {'start': 90, 'end': 100}, 'audio tags': [('寂静的声音', 0.9160257577896118)]}\n"
     ]
    }
   ],
   "source": [
    "audio_tag_result = whisper.parse_at_label(result, language='zh', top_k=5, p_threshold=-2, include_class_list=list(range(527)))\n",
    "for segment in audio_tag_result:\n",
    "  print(segment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pu4bjaS2ZgGA"
   },
   "source": [
    "### Step 4. Transcribe a video\n",
    "\n",
    "Let's check the result! The above audio track is actually from a video. You can of course generate .srt. But in this example, we directly put text and audio transcriptions to the video.\n",
    "\n",
    "**Step 4 is independent from the above, replace the URL and play with your own video!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 37
    },
    "id": "MsnGMfxWaQhw",
    "outputId": "35642b2a-b951-4fb2-a189-42bac7b334dd"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'/content/sample_video.mp4'"
      ],
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      }
     },
     "metadata": {},
     "execution_count": 14
    }
   ],
   "source": [
    "from IPython.display import HTML\n",
    "from base64 import b64encode\n",
    "# Replace this URL to play with your own video\n",
    "wget.download('https://www.dropbox.com/s/pzc72c59xtluuc0/case_closed.mp4?dl=1', '/content/sample_video.mp4')\n",
    "# mp4 = open('/content/sample_video.mp4','rb').read()\n",
    "# data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "# HTML(\"\"\"\n",
    "# <video width=800 controls>\n",
    "#       <source src=\"%s\" type=\"video/mp4\">\n",
    "# </video>\n",
    "# \"\"\" % data_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "KKbWD1KAa8e5"
   },
   "outputs": [],
   "source": [
    "!pip install -q ffmpeg-python\n",
    "import os,ffmpeg,cv2\n",
    "\n",
    "def dubbing_video(video_path, out_video_path, text_info, font_size=0.5, font_v_pos=0.95, font_color=(0, 0, 255)):\n",
    "    extract_audio(video_path, './temp_audio.wav')\n",
    "\n",
    "    video = cv2.VideoCapture(video_path)\n",
    "    # Get video properties\n",
    "    frame_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "    frame_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "    fps = video.get(cv2.CAP_PROP_FPS)\n",
    "\n",
    "    # Create output video writer\n",
    "    output_video = cv2.VideoWriter('./temp_video.mp4', cv2.VideoWriter_fourcc(*\"mp4v\"), fps,\n",
    "                                   (frame_width, frame_height))\n",
    "\n",
    "    # Process each frame of the video\n",
    "    current_frame = 0\n",
    "    while video.isOpened():\n",
    "        ret, frame = video.read()\n",
    "        if not ret:\n",
    "            break\n",
    "\n",
    "        # Calculate current time in seconds\n",
    "        current_time = current_frame / fps\n",
    "\n",
    "        # Iterate through text information and add text if within the time interval\n",
    "        for text_start, text_end, text in text_info:\n",
    "            if text_start <= current_time <= text_end:\n",
    "                text_position = (int(frame_width * 0.0), int(frame_height * font_v_pos))\n",
    "                font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "                font_scale = font_size\n",
    "                font_color = font_color\n",
    "                line_type = 1\n",
    "\n",
    "                cv2.putText(frame, text, text_position, font, font_scale, font_color, line_type, cv2.LINE_AA)\n",
    "\n",
    "        # Write the frame to the output video\n",
    "        output_video.write(frame)\n",
    "        current_frame += 1\n",
    "\n",
    "    # Release video resources\n",
    "    video.release()\n",
    "    output_video.release()\n",
    "\n",
    "    combine_audio_video('./temp_video.mp4', './temp_audio.wav', out_video_path)\n",
    "    os.remove('./temp_video.mp4')\n",
    "    os.remove('./temp_audio.wav')\n",
    "\n",
    "def combine_audio_video(video_path, audio_path, output_path):\n",
    "    video = ffmpeg.input(video_path)\n",
    "    audio = ffmpeg.input(audio_path)\n",
    "    output_file = ffmpeg.output(video, audio, output_path)\n",
    "    output_file.overwrite_output().run()\n",
    "\n",
    "def extract_audio(video_path, output_path):\n",
    "    video = ffmpeg.input(video_path)\n",
    "    audio = video.audio\n",
    "    output_file = ffmpeg.output(audio, output_path)\n",
    "    output_file.overwrite_output().run()\n",
    "\n",
    "extract_audio('/content/sample_video.mp4', '/content/sample_audio_from_video.wav')\n",
    "result = model.transcribe(\"/content/sample_audio_from_video.wav\", at_time_res=audio_tagging_time_resolution)\n",
    "\n",
    "# ASR Output\n",
    "text_segments = result['segments']\n",
    "text_annotation = [(x['start'], x['end'], x['text']) for x in text_segments]\n",
    "\n",
    "# Audio Tagging Output\n",
    "audio_tag_result = whisper.parse_at_label(result, language='follow_asr', top_k=5, p_threshold=-2, include_class_list=list(range(527)))\n",
    "\n",
    "all_seg = []\n",
    "for segment in audio_tag_result:\n",
    "    cur_start = segment['time']['start']\n",
    "    cur_end = segment['time']['end']\n",
    "    cur_tags = segment['audio tags']\n",
    "    cur_tags = [x[0] for x in cur_tags]\n",
    "    cur_tags = '; '.join(cur_tags)\n",
    "    all_seg.append((cur_start, cur_end, cur_tags))\n",
    "\n",
    "dubbing_video('/content/sample_video.mp4', '/content/sample_video_at.mp4', all_seg)\n",
    "dubbing_video('/content/sample_video_at.mp4', '/content/sample_video_at_text.mp4', text_annotation, font_color=(0,255,0), font_v_pos=0.85)\n",
    "os.remove('/content/sample_video.mp4')\n",
    "os.remove('/content/sample_video_at.mp4')\n",
    "\n",
    "# mp4 = open('/content/sample_video_at_text.mp4','rb').read()\n",
    "# data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "# HTML(\"\"\"\n",
    "# <video width=800 controls>\n",
    "#       <source src=\"%s\" type=\"video/mp4\">\n",
    "# </video>\n",
    "# \"\"\" % data_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kHA968-ShLzx"
   },
   "source": [
    "That's all. If you like the project, give us a star at https://github.com/YuanGongND/whisper-at."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}