speaker_diarization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_speakers = 3 #@param {type:\"integer\"}\n",
    "\n",
    "language = 'English' #@param ['any', 'English']\n",
    "\n",
    "model_size = 'medium' #@param ['tiny', 'base', 'small', 'medium', 'large']\n",
    "\n",
    "\n",
    "model_name = model_size\n",
    "if language == 'English' and model_size != 'large':\n",
    "  model_name += '.en'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/saravanan/anaconda3/envs/audio_summary/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/Users/saravanan/anaconda3/envs/audio_summary/lib/python3.10/site-packages/pyannote/audio/core/io.py:43: UserWarning: torchaudio._backend.set_audio_backend has been deprecated. With dispatcher enabled, this function is no-op. You can remove the function call.\n",
      "  torchaudio.set_audio_backend(\"soundfile\")\n",
      "/Users/saravanan/anaconda3/envs/audio_summary/lib/python3.10/site-packages/torch_audiomentations/utils/io.py:27: UserWarning: torchaudio._backend.set_audio_backend has been deprecated. With dispatcher enabled, this function is no-op. You can remove the function call.\n",
      "  torchaudio.set_audio_backend(\"soundfile\")\n"
     ]
    }
   ],
   "source": [
    "import whisper\n",
    "import datetime\n",
    "\n",
    "import subprocess\n",
    "\n",
    "import torch\n",
    "import pyannote.audio\n",
    "from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding\n",
    "embedding_model = PretrainedSpeakerEmbedding( \n",
    "    \"speechbrain/spkrec-ecapa-voxceleb\",\n",
    "    device=torch.device(\"cpu\"))\n",
    "from pyannote.audio import Audio\n",
    "from pyannote.core import Segment\n",
    "\n",
    "import wave\n",
    "import contextlib\n",
    "\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = whisper.load_model(model_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.transcribe('example-mono.wav')\n",
    "segments = result[\"segments\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with contextlib.closing(wave.open('example-mono.wav','r')) as f:\n",
    "  frames = f.getnframes()\n",
    "  rate = f.getframerate()\n",
    "  duration = frames / float(rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio()\n",
    "\n",
    "def segment_embedding(segment):\n",
    "  start = segment[\"start\"]\n",
    "  # Whisper overshoots the end timestamp in the last segment\n",
    "  end = min(duration, segment[\"end\"])\n",
    "  clip = Segment(start, end)\n",
    "  waveform, sample_rate = audio.crop('example-mono.wav', clip)\n",
    "  return embedding_model(waveform[None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = np.zeros(shape=(len(segments), 192))\n",
    "for i, segment in enumerate(segments):\n",
    "  embeddings[i] = segment_embedding(segment)\n",
    "\n",
    "embeddings = np.nan_to_num(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "clustering = AgglomerativeClustering(num_speakers).fit(embeddings)\n",
    "labels = clustering.labels_\n",
    "for i in range(len(segments)):\n",
    "  segments[i][\"speaker\"] = 'SPEAKER ' + str(labels[i] + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def time(secs):\n",
    "      return datetime.timedelta(seconds=round(secs))\n",
    "\n",
    "f = open(\"Example.txt\", \"w\")\n",
    "\n",
    "for (i, segment) in enumerate(segments):\n",
    "  if i == 0 or segments[i - 1][\"speaker\"] != segment[\"speaker\"]:\n",
    "    f.write(\"\\n\" + segment[\"speaker\"] + ' ' + str(time(segment[\"start\"])) + '\\n')\n",
    "  f.write(segment[\"text\"][1:] + ' ')\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "audio_summary",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}