{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":70203,"databundleVersionId":8068726,"sourceType":"competition"}],"dockerImageVersionId":30684,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Audio Denoising Model\nThis model was trained for speech, but seems to do ok with birds. \n\n[GitHub](https://github.com/will-rice/denoisers)\n\nDisclaimer: This is my model","metadata":{}},{"cell_type":"markdown","source":"## Install Dependencies","metadata":{}},{"cell_type":"code","source":"!pip install denoisers pedalboard pesq","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-04-11T18:11:15.903043Z","iopub.execute_input":"2024-04-11T18:11:15.903421Z","iopub.status.idle":"2024-04-11T18:11:40.765287Z","shell.execute_reply.started":"2024-04-11T18:11:15.903387Z","shell.execute_reply":"2024-04-11T18:11:40.763886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Denoiser Model","metadata":{}},{"cell_type":"code","source":"from pathlib import Path\nimport random\n\nimport torch\nimport torchaudio\nfrom denoisers import UNet1DModel\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport IPython.display as ipd\nimport librosa\n\n\nclass Denoiser:\n    def __init__(self, model_name=\"wrice/unet1d-vctk-48khz\"):\n        self.model = UNet1DModel.from_pretrained(model_name)\n        \n        if torch.cuda.is_available():\n            self.model.cuda()\n        \n    \n    def __call__(self, audio, sr):\n        if sr != self.model.config.sample_rate:\n            audio = torchaudio.functional.resample(audio, sr, self.model.config.sample_rate)\n\n        if audio.size(0) > 1:\n            audio = audio.mean(0, keepdim=True)\n\n        chunk_size = self.model.config.max_length\n\n        padding = abs(audio.size(-1) % chunk_size - chunk_size)\n        padded = torch.nn.functional.pad(audio, (0, padding))\n\n        clean = []\n        for i in tqdm(range(0, padded.shape[-1], chunk_size)):\n            audio_chunk = padded[:, i:i + chunk_size]\n            with torch.no_grad():\n                if torch.cuda.is_available():\n                    audio_chunk = audio_chunk.cuda()\n                clean_chunk = self.model(audio_chunk[None]).audio.cpu()\n            clean.append(clean_chunk.squeeze(0))\n\n        denoised = torch.concat(clean, 1)[:, :audio.shape[-1]]\n        \n        if sr != self.model.config.sample_rate:\n            denoised = torchaudio.functional.resample(denoised, self.model.config.sample_rate, sr)\n            \n        return denoised\n    \n","metadata":{"execution":{"iopub.status.busy":"2024-04-11T18:11:40.767710Z","iopub.execute_input":"2024-04-11T18:11:40.768168Z","iopub.status.idle":"2024-04-11T18:11:51.957654Z","shell.execute_reply.started":"2024-04-11T18:11:40.768049Z","shell.execute_reply":"2024-04-11T18:11:51.956647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Denoise Random Train Sample","metadata":{}},{"cell_type":"code","source":"root = Path(\"/kaggle/input/birdclef-2024/train_audio/\")\npaths = list(root.glob(\"**/*.ogg\"))\n\ndenoiser = Denoiser()\n\npath = random.choice(paths)\naudio, sr = torchaudio.load(path)\ndenoised = denoiser(audio, sr)\n\nipd.display(ipd.Audio(audio, rate=sr))\nipd.display(ipd.Audio(denoised, rate=sr))","metadata":{"execution":{"iopub.status.busy":"2024-04-11T18:11:51.958923Z","iopub.execute_input":"2024-04-11T18:11:51.959314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize","metadata":{}},{"cell_type":"code","source":"spec_fn = torchaudio.transforms.MelSpectrogram(\n    sample_rate=32000,\n    n_fft=2048,\n    win_length=2048,\n    hop_length=256,\n    power=2.0,\n    normalized=True\n)\namp_to_db = torchaudio.transforms.AmplitudeToDB(top_db=80.0)\n\noriginal_spec = spec_fn(audio[:, :32000*5])\nclean_spec = spec_fn(denoised[:, :32000*5])\n\nfig, ax = plt.subplots(1, 2, figsize=(50, 10), facecolor=\"white\")\nax[0].imshow(amp_to_db(original_spec.squeeze()), origin=\"lower\", aspect=\"auto\")\nax[0].title.set_text(\"original\")\nax[0].axis(\"off\")\n\nax[1].imshow(amp_to_db(clean_spec.squeeze()), origin=\"lower\", aspect=\"auto\")\nax[1].title.set_text(\"denoised\")\nax[1].axis(\"off\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}