demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from audio_embedding import extract_embeddings\n",
    "from model_engine import get_model, get_processor\n",
    "from utils import concat_and_rescale, save_embeddings\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import uuid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AUDIO_PATH = r\"D:\\RAS\\DVORAK_UMAP\\dvorak_01.wav\"\n",
    "OUTPUT_PATH = f\"./outputs/embedding_{uuid.uuid4()}\"\n",
    "\n",
    "BLOCK_LENGTH = 1280\n",
    "TARGET_SR = 16000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-25 14:25:18 | INFO | fairseq.models.wav2vec.wav2vec | Wav2VecModel(\n",
      "  (feature_extractor): ConvFeatureExtractionModel(\n",
      "    (conv_layers): ModuleList(\n",
      "      (0): Sequential(\n",
      "        (0): Conv1d(1, 512, kernel_size=(10,), stride=(5,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): Conv1d(512, 512, kernel_size=(8,), stride=(4,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "      (2): Sequential(\n",
      "        (0): Conv1d(512, 512, kernel_size=(4,), stride=(2,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "      (3): Sequential(\n",
      "        (0): Conv1d(512, 512, kernel_size=(4,), stride=(2,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "      (4): Sequential(\n",
      "        (0): Conv1d(512, 512, kernel_size=(4,), stride=(2,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "      (5): Sequential(\n",
      "        (0): Conv1d(512, 512, kernel_size=(1,), stride=(1,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "      (6): Sequential(\n",
      "        (0): Conv1d(512, 512, kernel_size=(1,), stride=(1,), bias=False)\n",
      "        (1): Dropout(p=0.0, inplace=False)\n",
      "        (2): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (3): ReLU()\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (feature_aggregator): ConvAggegator(\n",
      "    (conv_layers): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): ReplicationPad1d((1, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(2,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): ReplicationPad1d((2, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(3,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (2): Sequential(\n",
      "        (0): ReplicationPad1d((3, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(4,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (3): Sequential(\n",
      "        (0): ReplicationPad1d((4, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(5,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (4): Sequential(\n",
      "        (0): ReplicationPad1d((5, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(6,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (5): Sequential(\n",
      "        (0): ReplicationPad1d((6, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(7,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (6): Sequential(\n",
      "        (0): ReplicationPad1d((7, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(8,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (7): Sequential(\n",
      "        (0): ReplicationPad1d((8, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(9,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (8): Sequential(\n",
      "        (0): ReplicationPad1d((9, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(10,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (9): Sequential(\n",
      "        (0): ReplicationPad1d((10, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(11,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (10): Sequential(\n",
      "        (0): ReplicationPad1d((11, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(12,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "      (11): Sequential(\n",
      "        (0): ReplicationPad1d((12, 0))\n",
      "        (1): Conv1d(512, 512, kernel_size=(13,), stride=(1,))\n",
      "        (2): Dropout(p=0.0, inplace=False)\n",
      "        (3): Fp32GroupNorm(1, 512, eps=1e-05, affine=True)\n",
      "        (4): ReLU()\n",
      "      )\n",
      "    )\n",
      "    (residual_proj): ModuleList(\n",
      "      (0): None\n",
      "      (1): None\n",
      "      (2): None\n",
      "      (3): None\n",
      "      (4): None\n",
      "      (5): None\n",
      "      (6): None\n",
      "      (7): None\n",
      "      (8): None\n",
      "      (9): None\n",
      "      (10): None\n",
      "      (11): None\n",
      "    )\n",
      "  )\n",
      "  (wav2vec_predictions): Wav2VecPredictionsModel(\n",
      "    (project_to_steps): ConvTranspose2d(512, 512, kernel_size=(1, 12), stride=(1, 1))\n",
      "    (dropout): Dropout(p=0.0, inplace=False)\n",
      "  )\n",
      "  (dropout_feats): Dropout(p=0.0, inplace=False)\n",
      "  (dropout_agg): Dropout(p=0.0, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = get_model()\n",
    "processor = get_processor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 40/40 [00:00<?, ?it/s]\n"
     ]
    }
   ],
   "source": [
    "# Extract Embeddings\n",
    "raw_embeddings = extract_embeddings(\n",
    "    audio_path=AUDIO_PATH,\n",
    "    model=model,\n",
    "    processor=processor,\n",
    "    block_length=BLOCK_LENGTH,\n",
    "    target_sr=TARGET_SR,\n",
    ")\n",
    "embeddings = concat_and_rescale(raw_embeddings)\n",
    "print(embeddings.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embeddings\n",
    "pd.DataFrame(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save Embeddings\n",
    "save_embeddings(OUTPUT_PATH, embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Postprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import os\n",
    "from sklearn.preprocessing import MinMaxScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding = np.load(\"outputs\\embedding_dvorak_01.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(274292, 512)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 40/40 [00:13<00:00,  2.87it/s]\n"
     ]
    }
   ],
   "source": [
    "root_path = r\"outputs\"\n",
    "\n",
    "embedding_list = []\n",
    "for filename in tqdm(os.listdir(root_path)):\n",
    "    # File\n",
    "    ABS_PATH = os.path.join(root_path, filename)\n",
    "    embedding_list.append(np.load(ABS_PATH))\n",
    "\n",
    "# OUTPUT_PATH = f\"./outputs/embedding_{os.path.splitext(filename)[0]}.npy\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings_np = np.vstack(embedding_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Min Max Scaling\n",
      "11940 rows deleted.\n",
      "(14964695, 512)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Min Max Scaling\n",
    "print(\"\\nMin Max Scaling\")\n",
    "scaler = MinMaxScaler()\n",
    "scaled_np = scaler.fit_transform(embeddings_np)\n",
    "\n",
    "# Drop duplicate rows\n",
    "unique_rows = np.unique(scaled_np, axis=0)\n",
    "\n",
    "print(f\"{scaled_np.shape[0] - unique_rows.shape[0]} rows deleted.\")\n",
    "print(unique_rows.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Saving process started.\n",
      "(14964695, 512)\n",
      "Saved!\n"
     ]
    }
   ],
   "source": [
    "save_embeddings(\"outputs/all_dvorak15mill.npy\", unique_rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.DataFrame(unique_rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(14964695, 512)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sadtalker",
   "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.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}