preprocess/audio_resample.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Resample audio wavs\n",
        "\n",
        "Refer to: [audio resampling tutorial](https://pytorch.org/audio/0.10.0/tutorials/audio_resampling_tutorial.html)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import torchaudio\n",
        "import torchaudio.transforms as T\n",
        "import concurrent.futures\n",
        "from pathlib import Path\n",
        "import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Example usage:\n",
        "input_directory = \"/path/to/dataset\"\n",
        "output_directory = f\"{input_directory}.cleaned\"\n",
        "orig_sr = 16000\n",
        "new_sr = 22050"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def resample_wav_files(input_dir, output_dir, sr, new_sr):\n",
        "    # Create the output directory if it doesn't exist\n",
        "    os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "    # Create a resampler object\n",
        "    resampler = T.Resample(\n",
        "        sr,\n",
        "        new_sr,\n",
        "        lowpass_filter_width=128,\n",
        "        rolloff=0.99999,\n",
        "        resampling_method=\"sinc_interp_hann\",\n",
        "    )\n",
        "\n",
        "    def resample_file(file_path):\n",
        "        # Load the audio file\n",
        "        waveform, sample_rate = torchaudio.load(file_path)\n",
        "        assert sample_rate == sr\n",
        "\n",
        "        # Resample the audio\n",
        "        resampled_waveform = resampler(waveform)\n",
        "\n",
        "        # Construct the output file path\n",
        "        output_file = Path(output_dir) / Path(file_path).relative_to(input_dir)\n",
        "\n",
        "        # Save the resampled audio\n",
        "        torchaudio.save(output_file, resampled_waveform,\n",
        "                        new_sr, bits_per_sample=16)\n",
        "\n",
        "        return output_file\n",
        "\n",
        "    # Use generator to find .wav files and pre-create output directories\n",
        "    def find_and_prep_wav_files(input_dir, output_dir):\n",
        "        for root, _, files in os.walk(input_dir):\n",
        "            for file in files:\n",
        "                if file.endswith(\".wav\"):\n",
        "                    file_path = Path(root) / file\n",
        "                    output_file = Path(output_dir) / \\\n",
        "                        file_path.relative_to(input_dir)\n",
        "                    os.makedirs(output_file.parent, exist_ok=True)\n",
        "                    yield str(file_path)\n",
        "\n",
        "    # Resample the .wav files using threads for parallel processing\n",
        "    wav_files = find_and_prep_wav_files(input_dir, output_dir)\n",
        "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
        "        for i, output_file in enumerate(executor.map(resample_file, wav_files)):\n",
        "            if i % 1000 == 0:\n",
        "                print(f\"{i}: {output_file}\")\n",
        "\n",
        "\n",
        "resample_wav_files(input_directory, output_directory, orig_sr, new_sr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Test random file to see if it worked\n",
        "out_path = os.path.join(output_directory, os.listdir(output_directory)[random.randint(0, len(os.listdir(output_directory)))])\n",
        "\n",
        "print(torchaudio.info(out_path))\n",
        "resampled_waveform, sample_rate = torchaudio.load(out_path)\n",
        "print(f\"max: {resampled_waveform.max()}, min: {resampled_waveform.min()}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "g2p",
      "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.11.4"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2
}