{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "married-breast", "metadata": {}, "outputs": [], "source": [ "# voicebank, vctk: 44h, 109 spks, british accent\n", "\n", "# spectral map\n", "# 2mins per epoch, ~30 epochs gets 2.65 pesq\n", "\n", "# mimic loss\n", "# 10mins per epoch, ~50 epochs get 3.0 pesq\n", "\n", "# metricgan\n", "# 30mins per epoch, ~600 epochs get 3.15 pesq" ] }, { "cell_type": "code", "execution_count": null, "id": "forced-indie", "metadata": {}, "outputs": [], "source": [ "## Tutorial\n", "# https://colab.research.google.com/drive/18RyiuKupAhwWX7fh3LCatwQGU5eIS3TR\n", "\n", "## SpectralMask\n", "# https://github.com/speechbrain/speechbrain/tree/develop/recipes/Voicebank/enhance/spectral_mask\n", "\n", "## MimicLoss\n", "# https://github.com/speechbrain/speechbrain/tree/develop/recipes/Voicebank/MTL/ASR_enhance\n", "\n", "## MetricGan\n", "# https://github.com/speechbrain/speechbrain/tree/develop/recipes/Voicebank/enhance/MetricGAN\n", "\n", "\n", "\n", "## Sepformer\n", "# https://github.com/speechbrain/speechbrain/issues/499" ] }, { "cell_type": "code", "execution_count": null, "id": "instructional-acquisition", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ['CUDA_VISIBLE_DEVICES'] = ''" ] }, { "cell_type": "code", "execution_count": 7, "id": "dental-drinking", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torchaudio\n", "import numpy as np\n", "from pydub import AudioSegment\n", "from scipy.io import wavfile\n", "from IPython.display import display_html\n", "\n", "def convert_array(arr):\n", " if isinstance(arr, torch.Tensor):\n", " arr = arr.cpu().numpy()[0]\n", " if \"float\" in str(arr.dtype):\n", " arr = np.floor(arr.copy() * 32768).astype(np.int16)\n", " return arr\n", "\n", "def play_array(arr):\n", " arr = convert_array(arr)\n", " display_html(AudioSegment(\n", " arr.tobytes(), \n", " frame_rate=16000,\n", " sample_width=2, \n", " channels=1\n", " ))\n", " " ] }, { "cell_type": "markdown", "id": "chubby-friendly", "metadata": {}, "source": [ "### cloud model" ] }, { "cell_type": "code", "execution_count": 19, "id": "nominated-leisure", "metadata": {}, "outputs": [], "source": [ "from speechbrain.pretrained import SpectralMaskEnhancement" ] }, { "cell_type": "code", "execution_count": 20, "id": "tracked-america", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " " ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "enhance_model = SpectralMaskEnhancement.from_hparams(\n", " source=\"speechbrain/metricgan-plus-voicebank\",\n", " savedir=\"/data-ssd-3/georg/tmp/sb-models/metricgan-plus-voicebank\",\n", ")\n", "\n", "# Load and add fake batch dimension\n", "noisy = enhance_model.load_audio(\"tmp/bg_noisy.wav\").unsqueeze(0)\n", "\n", "# Add relative length tensor\n", "enhanced = enhance_model.enhance_batch(noisy, lengths=torch.tensor([1.]))\n", "# enhanced = enhance_model.enhance_batch(noisy)\n", "\n", "play_array(enhanced)" ] }, { "cell_type": "code", "execution_count": 21, "id": "frequent-democrat", "metadata": {}, "outputs": [], "source": [ "# enhance_model = SpectralMaskEnhancement.from_hparams(\n", "# source=\"speechbrain/mtl-mimic-voicebank\",\n", "# savedir=\"/data-ssd-3/georg/tmp/sb-models/mtl-mimic-voicebank\",\n", "# )\n", "\n", "# noisy = enhance_model.load_audio(\"tmp/bg_noisy.wav\").unsqueeze(0)\n", "\n", "# enhanced = enhance_model.enhance_batch(noisy)\n", "\n", "# play_array(enhanced)" ] }, { "cell_type": "markdown", "id": "fifth-bottle", "metadata": {}, "source": [ "### local model" ] }, { "cell_type": "code", "execution_count": 126, "id": "nervous-palmer", "metadata": {}, "outputs": [], "source": [ "# use generic interface\n", "# NOTE - had to add: `pretrainer: !new:speechbrain.utils.parameter_transfer.Pretraine` to yaml\n", "from speechbrain.pretrained import Pretrained\n", "\n", "enhance_model = Pretrained.from_hparams(\n", " source=\"/data-ssd-3/georg/tmp/sb-models/spectral-mask_4234\",\n", " savedir=\"/data-ssd-3/georg/tmp/sb-models/test\",\n", ")\n", "\n", "noisy = enhance_model.load_audio(\"tmp/bg_noisy.wav\").unsqueeze(0)" ] }, { "cell_type": "code", "execution_count": 127, "id": "european-event", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 127, "metadata": {}, "output_type": "execute_result" } ], "source": [ "cpp = \"/data-ssd-3/georg/tmp/sb-models/spectral-mask_4234/save/CKPT+2021-03-17+18-29-07+00/model.ckpt\"\n", "state_dict = torch.load(cpp, map_location='cpu')\n", "enhance_model.modules.model.load_state_dict(state_dict)" ] }, { "cell_type": "code", "execution_count": 125, "id": "nasty-relative", "metadata": {}, "outputs": [], "source": [ "# implement non-standard wrapper similar to SpectralMaskEnhancement (find info in train.py)\n", "from speechbrain.processing.features import spectral_magnitude\n", "\n", "def _compute_features(model, wavs):\n", " feats = model.hparams.compute_STFT(wavs)\n", " feats = spectral_magnitude(feats, power=0.5)\n", " return torch.log1p(feats)\n", "\n", "def enhance_batch_generic(model, noisy):\n", " noisy = noisy.to(model.device)\n", " noisy_features = _compute_features(model, noisy)\n", "\n", " mask = model.modules.model(noisy_features)\n", " mask = torch.squeeze(mask, 2)\n", " enhanced = torch.mul(mask, noisy_features)\n", "\n", " return model.hparams.resynth(torch.expm1(enhanced), noisy)" ] }, { "cell_type": "code", "execution_count": 122, "id": "medical-marijuana", "metadata": {}, "outputs": [], "source": [ "out = enhance_batch_generic(enhance_model, noisy)" ] }, { "cell_type": "code", "execution_count": 123, "id": "resident-glasgow", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " " ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "play_array(out)" ] }, { "cell_type": "code", "execution_count": null, "id": "outstanding-sleeping", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "signed-scholarship", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "infrared-bradford", "metadata": {}, "source": [ "### data prep" ] }, { "cell_type": "code", "execution_count": null, "id": "moral-courtesy", "metadata": {}, "outputs": [], "source": [ "# read_df = pd.read_csv(\"/data-ssd-3/georg/data/DNS-Challenge/datasets/clean/read_speech_meta.csv\")" ] }, { "cell_type": "code", "execution_count": 183, "id": "conceptual-amplifier", "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "TRAIN_DATA_DIR = \"/data-ssd-3/georg/data/DNS-Challenge/datasets/training_data-09_02_2021/\"\n", "\n", "with open(TRAIN_DATA_DIR + \"meta_tr/clean.json\", 'r') as f:\n", " clean_data_tr = json.load(f)\n", "with open(TRAIN_DATA_DIR + \"meta_tr/noisy.json\", 'r') as f:\n", " noisy_data_tr = json.load(f)\n", "with open(TRAIN_DATA_DIR + \"meta_val/clean.json\", 'r') as f:\n", " clean_data_val = json.load(f)\n", "with open(TRAIN_DATA_DIR + \"meta_val/noisy.json\", 'r') as f:\n", " noisy_data_val = json.load(f)" ] }, { "cell_type": "code", "execution_count": 159, "id": "rapid-excess", "metadata": {}, "outputs": [], "source": [ "tr_dict = {\n", " \"tr_\" + str(n): {\n", " \"clean_wav\": clean_path,\n", " \"noisy_wav\": noisy_path,\n", " \"length\": n_frames / 16000,\n", " } \n", " for n, ((clean_path, n_frames), (noisy_path, _)) in enumerate(zip(clean_data_tr, noisy_data_tr))\n", "}\n", "\n", "val_dict = {\n", " \"val_\" + str(n): {\n", " \"clean_wav\": clean_path,\n", " \"noisy_wav\": noisy_path,\n", " \"length\": n_frames / 16000,\n", " } \n", " for n, ((clean_path, n_frames), (noisy_path, _)) in enumerate(zip(clean_data_val, noisy_data_val))\n", "}" ] }, { "cell_type": "code", "execution_count": 160, "id": "opened-internet", "metadata": {}, "outputs": [], "source": [ "with open(TRAIN_DATA_DIR + \"sb_meta/train.json\", 'w') as f:\n", " f.write(json.dumps(tr_dict, indent=4))\n", "with open(TRAIN_DATA_DIR + \"sb_meta/val.json\", 'w') as f:\n", " f.write(json.dumps(val_dict, indent=4))" ] }, { "cell_type": "code", "execution_count": null, "id": "collected-wallpaper", "metadata": {}, "outputs": [], "source": [ "# /data-ssd-3/georg/data/DNS-Challenge/datasets/training_data-09_02_2021/sb_meta/train.json\n", "# /data-ssd-3/georg/data/DNS-Challenge/datasets/training_data-09_02_2021/sb_meta/val.json" ] }, { "cell_type": "code", "execution_count": null, "id": "invalid-badge", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "cardiovascular-juice", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "brilliant-express", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "fluid-topic", "metadata": {}, "source": [ "### Train model" ] }, { "cell_type": "code", "execution_count": null, "id": "automated-grace", "metadata": {}, "outputs": [], "source": [ "# 500h, ~20mins per epoch" ] }, { "cell_type": "code", "execution_count": null, "id": "future-silver", "metadata": {}, "outputs": [], "source": [ "# cd /data-ssd-3/georg/code/speechbrain/recipes/Voicebank/enhance/spectral_mask\n", "python train.py hparams/train.yaml --data_parallel_backend --name=lr_0005" ] }, { "cell_type": "code", "execution_count": null, "id": "general-enterprise", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "wooden-veteran", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "burning-virgin", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "juvenile-philosophy", "metadata": {}, "source": [ "## Playground" ] }, { "cell_type": "code", "execution_count": null, "id": "armed-married", "metadata": {}, "outputs": [], "source": [ "cd recipes///\n", "python experiment.py params.yaml --data_parallel_backend" ] }, { "cell_type": "code", "execution_count": null, "id": "naughty-sequence", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "liquid-premium", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "external-denial", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "acceptable-touch", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "addressed-submission", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "buried-gilbert", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.8.10" } }, "nbformat": 4, "nbformat_minor": 5 }