GPU Spectrogram Transformation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GPU Backed Spectogarms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchaudio.transforms import MelSpectrogram, PadTrim\n",
    "from fastai.vision import Image, open_image, image2np\n",
    "from fastai.basics import *\n",
    "import time\n",
    "import PIL\n",
    "import numpy as np\n",
    "from torchaudio import load as load_audio\n",
    "from concurrent.futures import ProcessPoolExecutor\n",
    "import multiprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "AUDIO_EXTENSIONS = tuple(str.lower(k) for k,v in mimetypes.types_map.items() \n",
    "                         if v.startswith('audio/'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "class SPEC2DB(object):\n",
    "    \"\"\"Turns a spectrogram from the power/amplitude scale to the decibel scale.\n",
    "\n",
    "    Args:\n",
    "        stype (str): scale of input spectrogram (\"power\" or \"magnitude\").  The\n",
    "            power being the elementwise square of the magnitude. default: \"power\"\n",
    "        top_db (float, optional): minimum negative cut-off in decibels.  A reasonable number\n",
    "            is -80.\n",
    "    \"\"\"\n",
    "    def __init__(self, stype=\"power\", top_db=None):\n",
    "        self.stype = stype\n",
    "        self.top_db = -top_db if top_db > 0 else top_db\n",
    "        self.multiplier = 10. if stype == \"power\" else 20.\n",
    "\n",
    "    def __call__(self, spec):\n",
    "        spec_db = self.multiplier * torch.log10(spec / spec.max())  # power -> dB\n",
    "        if self.top_db is not None:\n",
    "            spec_db = torch.max(spec_db, spec_db.new([self.top_db]))\n",
    "        return spec_db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/h/.fastai/data/ST-AEDS-20180100_1-OS')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## The actual url is http://www.openslr.org/resources/45/ST-AEDS-20180100_1-OS.tgz\n",
    "## but we need to strip off the extension otherwise fastai gets confused.\n",
    "data_url = 'http://www.openslr.org/resources/45/ST-AEDS-20180100_1-OS'\n",
    "## Need this because the source tar file doesn't extract to its own folder\n",
    "p = datapath4file(url2name(data_url))\n",
    "untar_data(data_url, dest=p)\n",
    "p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_spectogram(p, to_db_scale=True, force_cache=False, n_fft=1024,\n",
    "                ws=None, hop=72, f_min=0.0, f_max=200, pad=0, n_mels=224, max_to_pad=16000):\n",
    "    sig, sr = load_audio(p)\n",
    "    image_path = p.with_suffix('.jpg')\n",
    "    sig = sig.clone().cuda()\n",
    "    sig = PadTrim(max_len=max_to_pad)(sig).squeeze()\n",
    "    mel = MelSpectrogram(sr=sr, n_mels=n_mels, n_fft=n_fft, ws=ws, hop=hop,\n",
    "                    f_min=f_min, f_max=f_max, pad=pad)(sig.reshape(1, -1))\n",
    "    mel = mel.permute(0, 2, 1) # swap dimension...\n",
    "    mel = SPEC2DB(stype='magnitude', top_db=f_max)(mel)\n",
    "    x = image2np(mel).astype(np.uint8)\n",
    "    PIL.Image.fromarray(x).save(image_path)\n",
    "\n",
    "def spectogram_files(folder, **kwargs):\n",
    "    files = get_files(folder, extensions=AUDIO_EXTENSIONS)\n",
    "    f = partial(create_spectogram, **kwargs)\n",
    "    [f(x) for x in files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11.725131034851074"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max_length = 4*16000\n",
    "\n",
    "tfm_params = {\n",
    "    'max_to_pad' : max_length,\n",
    "    'to_db_scale':True,\n",
    "    'f_max': 80\n",
    "}\n",
    "\n",
    "t = time.time()\n",
    "spectogram_files(p, **tfm_params)\n",
    "time.time() - t"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}