{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f72992c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchaudio\n",
    "import librosa\n",
    "import pyrubberband as pyrb\n",
    "from torch import Tensor, no_grad, device\n",
    "from torchaudio.transforms import  MelSpectrogram, AmplitudeToDB, Resample, TimeStretch\n",
    "from pathlib import Path\n",
    "from librosa import effects\n",
    "import torchaudio.functional as F\n",
    "from IPython.display import Audio\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "\n",
    "from tokenizer import ASRTokenizer\n",
    "from model import ASRModel\n",
    "\n",
    "CONFIG = {\n",
    "    # 模型配置\n",
    "    'input_dim': 640,\n",
    "    'num_heads': 8,\n",
    "    'ffn_dim': 2048,\n",
    "    'num_layers': 8,\n",
    "    'dropout': 0.1,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e77899f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ASRInference:\n",
    "    def __init__(self, model_path: Path, vocab_path: Path, device: device, augment: bool = True, augment_prob: float = 0.5) -> None:\n",
    "        self.device = device\n",
    "        self.augment: bool = augment\n",
    "        self.augment_prob: float = augment_prob\n",
    "        self.tokenizer = ASRTokenizer(vocab_path=vocab_path)\n",
    "        self.model = ASRModel(vocab_size=self.tokenizer.vocab_size(), **CONFIG).to(device=device)\n",
    "        self.model.load_state_dict(torch.load(model_path, map_location=device)['model_state_dict'])\n",
    "        self.model.eval()\n",
    "\n",
    "        print(f\"params params: {self.model.get_num_params():,}\",)\n",
    "\n",
    "        self.sample_rate = 16000\n",
    "        self.mel_transform = MelSpectrogram(\n",
    "            sample_rate=self.sample_rate,\n",
    "            n_fft=400,\n",
    "            win_length=400,\n",
    "            hop_length=160 // 1,\n",
    "            n_mels=80 * 4,\n",
    "            f_min=0,\n",
    "            f_max=8000,\n",
    "            power=3.0,\n",
    "        )\n",
    "        self.amplitude_to_db = AmplitudeToDB()\n",
    "        \n",
    "    def _load_audio(self, audio_path: Path) -> Tensor:\n",
    "        waveform, sample_rate = torchaudio.load_with_torchcodec(audio_path)\n",
    "\n",
    "        if sample_rate != self.sample_rate:\n",
    "          waveform = Resample(sample_rate, self.sample_rate)(waveform)\n",
    "\n",
    "        if waveform.shape[0] > 1:\n",
    "            waveform = waveform.mean(dim=0, keepdim=True)\n",
    "\n",
    "        waveform = waveform / (waveform.abs().max() + 1e-8)\n",
    "        return waveform\n",
    "    \n",
    "    def _extract_features(self, waveform: Tensor) -> Tensor:\n",
    "        mel_spec: Tensor = self.mel_transform(waveform)\n",
    "        log_mel_spec: Tensor = self.amplitude_to_db(mel_spec)\n",
    "        return log_mel_spec.squeeze(0) # [n_mels, time]\n",
    "    \n",
    "    def _augment_waveform(self, waveform: Tensor) -> Tensor:\n",
    "        if not self.augment or random.random() > self.augment_prob:\n",
    "            return waveform\n",
    "        \n",
    "        # 1. voice Stretch/Compress \n",
    "        if random.random() < 0.5:\n",
    "            waveform = self._voice_stretch_or_compress(waveform=waveform)\n",
    "\n",
    "        if random.random() < 0.3:\n",
    "            waveform = self._drop_frames(waveform)\n",
    "\n",
    "        if random.random() < 0.4:\n",
    "            waveform = self._add_noise(waveform)\n",
    "        \n",
    "        return waveform\n",
    "    \n",
    "    def _voice_stretch_or_compress1(self, waveform: Tensor) -> Tensor:\n",
    "        speed_factor = random.uniform(0.6, 1.4) # (Speed Change: 0.6x - 1.4x)\n",
    "        waveform_np = waveform.squeeze(0).cpu().numpy()\n",
    "        y_stretch = librosa.effects.time_stretch(waveform_np, rate=speed_factor)\n",
    "        waveform = torch.from_numpy(y_stretch).float().to(waveform.device).unsqueeze(0)\n",
    "        return waveform\n",
    "    \n",
    "    def _voice_stretch_or_compress(self, waveform: Tensor) -> Tensor:\n",
    "        speed_factor = random.uniform(0.6, 1.4) # (Speed Change: 0.6x - 1.4x)\n",
    "        spec = torch.stft(\n",
    "            waveform.squeeze(0),\n",
    "            n_fft=400,\n",
    "            hop_length=160,\n",
    "            window=torch.hann_window(400).to(waveform.device),\n",
    "            return_complex=True\n",
    "        )\n",
    "        \n",
    "        # 时间拉伸（不改变音高）\n",
    "        stretch = TimeStretch(\n",
    "            hop_length=160,\n",
    "            n_freq=201,\n",
    "            fixed_rate=speed_factor\n",
    "        )\n",
    "        stretched_spec = stretch(spec)\n",
    "        \n",
    "        # 转回波形\n",
    "        waveform_stretched = torch.istft(\n",
    "            stretched_spec,\n",
    "            n_fft=400,\n",
    "            hop_length=160,\n",
    "            window=torch.hann_window(400).to(waveform.device)\n",
    "        ).unsqueeze(0)\n",
    "    \n",
    "        return waveform_stretched\n",
    "\n",
    "    def _voice_stretch_or_compress2(self, waveform: Tensor) -> Tensor:\n",
    "        speed_factor = random.uniform(0.6, 1.4) # (Speed Change: 0.6x - 1.4x)\n",
    "        waveform_np = waveform.squeeze(0).cpu().numpy()\n",
    "        \n",
    "        y_stretch = pyrb.time_stretch(waveform_np, self.sample_rate, speed_factor)\n",
    "        \n",
    "        return torch.from_numpy(y_stretch).float().to(waveform.device).unsqueeze(0)\n",
    "    \n",
    "    def _drop_frames(self, waveform: Tensor) -> Tensor:\n",
    "        audio_len = waveform.shape[1]\n",
    "        drop_ratio = random.uniform(0.05, 0.15)\n",
    "        drop_len = int(audio_len * drop_ratio)\n",
    "\n",
    "        if audio_len > drop_len:\n",
    "            start_pos = random.randint(0, audio_len - drop_len)\n",
    "            # clean\n",
    "            waveform = torch.cat([waveform[:, :start_pos], waveform[:, start_pos + drop_len:]], dim=1)\n",
    "\n",
    "        return waveform\n",
    "    \n",
    "    def _add_noise(self, waveform: Tensor) -> Tensor:\n",
    "        snr_db = random.uniform(10, 20)\n",
    "        signal_power = torch.mean(waveform ** 2)\n",
    "        \n",
    "        snr_linear = 10 ** (snr_db / 10)\n",
    "        noise_power = signal_power / snr_linear\n",
    "        \n",
    "        noise = torch.randn_like(waveform) * torch.sqrt(noise_power)\n",
    "        return waveform + noise\n",
    "    \n",
    "    def transcribe(self, audio_path: Path) -> str:\n",
    "        waveform = self._load_audio(audio_path=audio_path)\n",
    "        # waveform = self._augment_waveform(waveform)\n",
    "        mel_spec = self._extract_features(waveform=waveform)\n",
    "\n",
    "        mel_spec = mel_spec.unsqueeze(0).to(device=self.device)\n",
    "        mel_length = torch.tensor([mel_spec.shape[2]], dtype=torch.long, device=self.device)\n",
    "\n",
    "        with no_grad():\n",
    "            log_probs, _ = self.model(mel_specs=mel_spec, mel_lengths=mel_length)\n",
    "        \n",
    "        text = self.tokenizer.ctc_greedy_decode(log_probs=log_probs[0])\n",
    "        return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5af7c836",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load Checkpoint: /home/blacksheep/projekts/study_asr/.checkpoints/best_wer_model.pt\n",
      "params params: 66,766,816\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/blacksheep/projekts/study_asr/.venv/lib/python3.11/site-packages/torchaudio/functional/functional.py:585: UserWarning: At least one mel filterbank has all zero values. The value for `n_mels` (320) may be set too high. Or, the value for `n_freqs` (201) may be set too low.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:1')\n",
    "workspace_dir = Path.cwd().parent\n",
    "\n",
    "# checkpoint = sorted(workspace_dir.glob('.checkpoints/checkpoint_step_*.pt'), key=lambda p: int(p.stem.split('_')[-1]))[-1]\n",
    "checkpoint = workspace_dir / \".checkpoints/best_wer_model.pt\"\n",
    "# checkpoint = workspace_dir / \".checkpoints/best_wer_model.pt\"\n",
    "print(f\"Load Checkpoint: {checkpoint}\")\n",
    "inference = ASRInference(model_path=checkpoint, vocab_path=workspace_dir / 'config/asr_vocab.json' , device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "967e23cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 208051])\n",
      "Orginal audio\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 554864 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Torch audio\n",
      "\n",
      "ئامېرىكا جاتېز رىسدىكى رەپشى تۇچوڭ، چاۋشراتكى دەمتىق بېقىس ئۇرۇشىۋان ئۇيۈم غەلىبىسى دەپگو ھاتقى گەۋنىقال قولدارلارقال كې، يوقياۋان پارقىدە ئىس كۆئىمخەتكى خىدىڭنامىسى.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_27571678.mp3'\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_40794549.mp3'\n",
    "# audio_path = workspace_dir / 'data/test/let_me_see_the_computer.mp3'\n",
    "# audio_path = '/mnt/train/audio/Dilnaz/only_vocals/3272800876_100006098_(vocals)_melband_roformer_big_beta5e.m4a'\n",
    "# audio_path = workspace_dir / 'data/test/how_are_you.mp3'\n",
    "audio_path = workspace_dir / 'data/test/split.m4a'\n",
    "# audio_path = workspace_dir / 'data/test/let_me_see_the_computer.mp3'\n",
    "# audio_path = workspace_dir / 'data/test/F001_001.wav'\n",
    "orginal_waveform = inference._load_audio(audio_path=audio_path)\n",
    "print(orginal_waveform.shape)\n",
    "orginal_waveform = orginal_waveform[:, :]\n",
    "print(\"Orginal audio\")\n",
    "display(Audio(orginal_waveform, rate=inference.sample_rate))\n",
    "print(\"Torch audio\")\n",
    "# display(Audio(torch_audio_waveform, rate=inference.sample_rate))\n",
    "text = inference.transcribe(audio_path=audio_path)\n",
    "print(f\"\\n{text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3671d334",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48.266026 -100.0\n",
      "1.0 0.0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 3397484 characters]",
      "text/plain": [
       "<Figure size 3000x1200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([320, 1301])\n",
      "torch.Size([1, 320, 1301])\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1104380 characters]",
      "text/plain": [
       "<Figure size 1800x1200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mel 频谱图形状: torch.Size([1, 320, 1301])\n"
     ]
    }
   ],
   "source": [
    "mel_spec = inference._extract_features(waveform=orginal_waveform)\n",
    "################################\n",
    "from PIL import Image, ImageFilter\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "img = (mel_spec.numpy())\n",
    "print(img.max(), img.min())\n",
    "img = img + np.abs(img.min())\n",
    "img = img / img.max()\n",
    "print(img.max(), img.min())\n",
    "img = Image.fromarray((img * 255).astype(np.uint8))\n",
    "shapened = img.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))\n",
    "img.save(\"original_hd.jpg\")\n",
    "# img.save(\"sharpen_hd.jpg\")\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(30, 12))\n",
    "\n",
    "axes[0].imshow(\n",
    "    img,\n",
    "    cmap=\"gray\",\n",
    "    origin=\"lower\",\n",
    "    vmin=0,\n",
    "    vmax=255,\n",
    "    aspect=\"auto\"\n",
    ")\n",
    "axes[0].set_title(\"Original\")\n",
    "axes[1].imshow(\n",
    "    shapened,\n",
    "    cmap=\"gray\",\n",
    "    origin=\"lower\",\n",
    "    vmin=0,\n",
    "    vmax=255,\n",
    "    aspect=\"auto\"\n",
    ")\n",
    "axes[1].set_title(\"Sharpened\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "################################\n",
    "print(mel_spec.shape)\n",
    "mel_spec = mel_spec.unsqueeze(0)\n",
    "print(mel_spec.shape)\n",
    "# 创建可视化\n",
    "fig, axes = plt.subplots(2, 1, figsize=(18, 12))\n",
    "\n",
    "# 原始波形\n",
    "axes[0].plot(orginal_waveform[0].numpy())\n",
    "axes[0].set_title(\"Original waveform\")\n",
    "axes[0].set_xlabel(\"sample\")\n",
    "axes[0].set_ylabel(\"Amplitude\")\n",
    "axes[0].grid(True)\n",
    "\n",
    "# Mel 频谱图\n",
    "im = axes[1].imshow(\n",
    "    mel_spec[0].numpy(),\n",
    "    cmap='gray',\n",
    "    aspect='auto',\n",
    "    origin='lower',\n",
    "    # vmin=-150,\n",
    "    # vmax=1,\n",
    "    extent=[0, orginal_waveform.shape[-1] / inference.sample_rate, 0, 80]\n",
    ")\n",
    "axes[1].set_title(\"Log Mel Spectrogram\")\n",
    "axes[1].set_xlabel(\"Time (second)\")\n",
    "axes[1].set_ylabel(\"Mels\")\n",
    "# plt.colorbar(im, ax=axes[1], format='%+2.0f dB')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "\n",
    "print(f\"Mel 频谱图形状: {mel_spec.shape}\")  # (channel, n_mels, time_frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a776864",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "study-asr (3.11.12)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
