SA-colab.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "fJUJLWQ92g6R"
      },
      "outputs": [],
      "source": [
        "!pip install -q gradio soundfile configargparse"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e200MmBU2aLT"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "import numpy as np\n",
        "import torch\n",
        "import torchaudio\n",
        "import sys\n",
        "\n",
        "\n",
        "def convert_to_16_bit_wav(data):\n",
        "    # Based on: https://docs.scipy.org/doc/scipy/reference/generated/scipy.io.wavfile.write.html\n",
        "    #breakpoint()\n",
        "    if data.dtype == np.float32:\n",
        "        print(\n",
        "            \"Audio data is not in 16-bit integer format.\",\n",
        "            \"Trying to convert to 16-bit int format.\",\n",
        "            file=sys.stderr\n",
        "        )\n",
        "        data = data / np.abs(data).max()\n",
        "        data = data * 32767\n",
        "        data = data.astype(np.int16)\n",
        "    elif data.dtype == np.int32:\n",
        "        print(\n",
        "            \"Audio data is not in 16-bit integer format.\",\n",
        "            \"Trying to convert to 16-bit int format.\",\n",
        "            file=sys.stderr\n",
        "        )\n",
        "        data = data / 65538\n",
        "        data = data.astype(np.int16)\n",
        "    elif data.dtype == np.int16:\n",
        "        pass\n",
        "    elif data.dtype == np.uint8:\n",
        "        print(\n",
        "            \"Audio data is not in 16-bit integer format.\",\n",
        "            \"Trying to convert to 16-bit int format.\",\n",
        "            file=sys.stderr\n",
        "        )\n",
        "        data = data * 257 - 32768\n",
        "        data = data.astype(np.int16)\n",
        "    else:\n",
        "        raise ValueError(\"Audio data cannot be converted to \" \"16-bit int format.\")\n",
        "    return data\n",
        "\n",
        "def pcm2float(sig, dtype='float32'):\n",
        "    \"\"\"\n",
        "    https://gist.github.com/HudsonHuang/fbdf8e9af7993fe2a91620d3fb86a182\n",
        "    \"\"\"\n",
        "    sig = np.asarray(sig)\n",
        "    if sig.dtype.kind not in 'iu':\n",
        "        raise TypeError(\"'sig' must be an array of integers\")\n",
        "    dtype = np.dtype(dtype)\n",
        "    if dtype.kind != 'f':\n",
        "        raise TypeError(\"'dtype' must be a floating point type\")\n",
        "\n",
        "    i = np.iinfo(sig.dtype)\n",
        "    abs_max = 2 ** (i.bits - 1)\n",
        "    offset = i.min + abs_max\n",
        "    return (sig.astype(dtype) - offset) / abs_max\n",
        "\n",
        "\n",
        "def float2pcm(sig, dtype='int16'):\n",
        "    \"\"\"\n",
        "    https://gist.github.com/HudsonHuang/fbdf8e9af7993fe2a91620d3fb86a182\n",
        "    \"\"\"\n",
        "    sig = np.asarray(sig)\n",
        "    if sig.dtype.kind != 'f':\n",
        "        raise TypeError(\"'sig' must be a float array\")\n",
        "    dtype = np.dtype(dtype)\n",
        "    if dtype.kind not in 'iu':\n",
        "        raise TypeError(\"'dtype' must be an integer type\")\n",
        "    i = np.iinfo(dtype)\n",
        "    abs_max = 2 ** (i.bits - 1)\n",
        "    offset = i.min + abs_max\n",
        "    return (sig * abs_max + offset).clip(i.min, i.max).astype(dtype)\n",
        "\n",
        "\n",
        "@torch.no_grad()\n",
        "def inference(audio, model_tag):\n",
        "  sr, audio = audio\n",
        "  audio = convert_to_16_bit_wav(audio)\n",
        "  audio = pcm2float(audio)\n",
        "  audio = torch.tensor(audio).unsqueeze(0)\n",
        "  audio = torchaudio.transforms.Resample(orig_freq=sr,\n",
        "                                         new_freq=16000)(audio)\n",
        "  model = torch.hub.load(\"deep-privacy/SA-toolkit\", \"anonymization\", tag_version=model_tag, trust_repo=True)\n",
        "  model.eval()\n",
        "  wav_conv = model.convert(audio, target=\"6081\") # hard coded target\n",
        "  return 16000, float2pcm(wav_conv.squeeze().cpu().numpy())\n",
        "\n",
        "\n",
        "article = \"<p style='text-align: center'><a href='https://arxiv.org/abs/2308.04455' target='_blank'>PhD thesis: Anonymizing Speech: Evaluating and Designing Speaker Anonymization Techniques</a> | <a href='https://github.com/deep-privacy/SA-toolkit' target='_blank'>Github Repo</a></p>\"\n",
        "\n",
        "\n",
        "def toggle_audio_src(choice):\n",
        "    if choice == \"mic\":\n",
        "        return gr.update(source=\"microphone\", value=None, label=\"Microphone (best with a headset)\")\n",
        "    else:\n",
        "        return gr.update(source=\"upload\", value=None, label=\"File\")\n",
        "\n",
        "with gr.Blocks() as interface:\n",
        "  gr.Markdown(\n",
        "            \"\"\"\n",
        "            # SA-toolkit\n",
        "            Demo: Speaker speech anonymization toolkit in python\n",
        "            \"\"\"\n",
        "        )\n",
        "  with gr.Row():\n",
        "    with gr.Column():\n",
        "      radio = gr.Radio([\"file\", \"mic\"], value=\"file\",\n",
        "                       label=\"Input speech (File or Mic)\")\n",
        "      audio_input = gr.Audio(source=\"upload\", type=\"numpy\", label=\"File\",\n",
        "                        interactive=True, elem_id=\"melody-input\")\n",
        "      model_tag = gr.Dropdown(['hifigan_bn_tdnnf_wav2vec2_vq_48_v1',\n",
        "                              'hifigan_bn_tdnnf_wav2vec2_100h_aug_v1',\n",
        "                              'hifigan_bn_tdnnf_600h_aug_v1',\n",
        "                              'hifigan_bn_tdnnf_100h_vq_64_v1',\n",
        "                              'hifigan_bn_tdnnf_100h_vq_256_v1',\n",
        "                              'hifigan_bn_tdnnf_100h_aug_v1'], type='value',\n",
        "                              value='hifigan_bn_tdnnf_wav2vec2_vq_48_v1',\n",
        "                             label='Model')\n",
        "      with gr.Row():\n",
        "        submit = gr.Button(\"Submit\")\n",
        "    with gr.Column():\n",
        "      audio_output = gr.Audio(label=\"Output\")\n",
        "  submit.click(inference, inputs=[audio_input, model_tag],\n",
        "                 outputs=[audio_output], batch=False)\n",
        "  radio.change(toggle_audio_src, radio, [audio_input], queue=False, show_progress=False)\n",
        "  gr.Examples(fn=inference,\n",
        "              examples=[['https://datasets-server.huggingface.co/assets/librispeech_asr/--/all/train.clean.100/14/audio/audio.mp3']],\n",
        "              inputs=[audio_input, \"hifigan_bn_tdnnf_wav2vec2_vq_48_v1\"],\n",
        "              outputs=[audio_output], batch=False)\n",
        "\n",
        "\n",
        "  gr.HTML(article)\n",
        "  interface.queue().launch(debug=True)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}