{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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 FrequencyMasking, MelSpectrogram, AmplitudeToDB, Resample, TimeMasking, TimeStretch\n",
    "from torch_audiomentations import Gain, PitchShift, LowPassFilter, HighPassFilter\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",
    "from typing import TypedDict\n",
    "\n",
    "from tokenizer import ASRTokenizer\n",
    "from model import ASRModel\n",
    "\n",
    "CONFIG = {\n",
    "    # 模型配置\n",
    "    'input_dim': 256,\n",
    "    'num_heads': 8,\n",
    "    'ffn_dim': 2048,\n",
    "    'num_layers': 8,\n",
    "    'dropout': 0.1,\n",
    "}\n",
    "workspace_dir = Path.cwd().parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15a9a926",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AdvancedAugment:\n",
    "    def __init__(self, noise_folder=None):\n",
    "        # 如果你有噪音文件夹，可以预加载噪音路径列表\n",
    "        # 场景建议：1.汽车 2.餐厅 3.街道 4.办公室 5.下雨 6.风声 7.键盘声 8.婴儿哭声 9.音乐 10.走廊回声 11.白噪音 12.粉红噪音\n",
    "        self.noise_files = [] # 这里存放你的 .wav 噪音文件路径\n",
    "\n",
    "    def _add_noise(self, waveform: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        更丰富的加噪：随机选择 白噪音、粉红噪音 或 真实环境音\n",
    "        \"\"\"\n",
    "        # 1. 随机决定噪音类型\n",
    "        noise_type = random.choice(['white', 'pink', 'environmental'])\n",
    "        \n",
    "        # 2. 随机设定信噪比 (SNR) - 模拟各种清晰度\n",
    "        snr_db = random.uniform(5, 25) \n",
    "        \n",
    "        if noise_type == 'white':\n",
    "            noise = torch.randn_like(waveform)\n",
    "        elif noise_type == 'pink':\n",
    "            \n",
    "            noise = self._generate_pink_noise(waveform)\n",
    "        else:\n",
    "            # 模拟环境音 (如果你有噪音库)\n",
    "            # 这里演示如果没有噪音库，就用多种不同频率的随机噪音组合\n",
    "            noise = self._generate_simulated_env_noise(waveform)\n",
    "\n",
    "        # 3. 计算能量并叠加\n",
    "        return self._mix_signal_noise(waveform, noise, snr_db)\n",
    "\n",
    "    def _generate_pink_noise(self, waveform):\n",
    "        \"\"\"粉红噪音（比白噪音更像自然界的声音，低频能量更高）\"\"\"\n",
    "        # 简易实现：对白噪音做低通滤波\n",
    "        white = torch.randn_like(waveform)\n",
    "        return torch.cumsum(white, dim=-1) / 10.0 # 简单的积分近似\n",
    "\n",
    "    def _generate_simulated_env_noise(self, waveform):\n",
    "        \"\"\"模拟环境噪音（通过叠加不同频率的波形）\"\"\"\n",
    "        # 模拟 12 种以上变化：随机叠加正弦波或窄带噪音\n",
    "        noise = torch.zeros_like(waveform)\n",
    "        for _ in range(random.randint(3, 8)):\n",
    "            freq = random.uniform(50, 4000)\n",
    "            t = torch.arange(waveform.shape[-1]).to(waveform.device)\n",
    "            noise += torch.sin(2 * 3.14159 * freq * t / 16000)\n",
    "        return noise + torch.randn_like(waveform) * 0.5\n",
    "\n",
    "    def _mix_signal_noise(self, signal, noise, snr_db):\n",
    "        \"\"\"核心：根据 SNR 混合信号和噪音\"\"\"\n",
    "        s_p = signal.pow(2).mean()\n",
    "        n_p = noise.pow(2).mean()\n",
    "        \n",
    "        # 防止除以 0\n",
    "        if n_p == 0: return signal\n",
    "        \n",
    "        snr_linear = 10**(snr_db/10)\n",
    "        scale = torch.sqrt(s_p / (n_p * snr_linear))\n",
    "        \n",
    "        noisy = signal + scale * noise\n",
    "        return noisy / (noisy.abs().max() + 1e-7) # 归一化防止爆音\n",
    "\n",
    "    def _simulate_reverb(self, waveform):\n",
    "        \"\"\"\n",
    "        模拟走廊/房间回声 (Simple Reverb)\n",
    "        \"\"\"\n",
    "        # if random.random() < 0.3: # 30% 概率添加回声\n",
    "            # 模拟简单的延迟反馈（模拟大走廊）\n",
    "        delay_samples = random.randint(500, 2000) # 30ms - 125ms 延迟\n",
    "        decay = random.uniform(0.3, 0.6)\n",
    "        \n",
    "        reverb_signal = torch.zeros_like(waveform)\n",
    "        reverb_signal[:, delay_samples:] = waveform[:, :-delay_samples] * decay\n",
    "        return (waveform + reverb_signal) / (1 + decay)\n",
    "        return waveform\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0635d68",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NoiseAugmentor:\n",
    "    def __init__(self, noise_root: Path, sample_rate=16000):\n",
    "        self.sample_rate = sample_rate\n",
    "        # 递归找到目录下所有的 wav 文件\n",
    "        self.noise_files = list(Path(noise_root).rglob(\"*.wav\"))\n",
    "        print(f\"成功加载 {len(self.noise_files)} 个噪音文件\")\n",
    "\n",
    "    def apply_real_noise(self, waveform: Tensor):\n",
    "        # 1. 随机选一个噪音文件\n",
    "        noise_path = random.choice(self.noise_files)\n",
    "        noise_waveform, sr = torchaudio.load_with_torchcodec(noise_path)\n",
    "        \n",
    "        # 2. 统一采样率\n",
    "        if sr != self.sample_rate:\n",
    "            resampler = Resample(sr, self.sample_rate)\n",
    "            noise_waveform = resampler(noise_waveform)\n",
    "        \n",
    "        # 3. 截取或填充，使其长度与语音一致\n",
    "        sig_len = waveform.shape[1]\n",
    "        noise_len = noise_waveform.shape[1]\n",
    "        \n",
    "        if noise_len >= sig_len:\n",
    "            # 随机截取一段\n",
    "            start = random.randint(0, noise_len - sig_len)\n",
    "            noise_waveform = noise_waveform[:, start:start + sig_len]\n",
    "        else:\n",
    "            full_noise = torch.zeros_like(waveform)\n",
    "            start = random.randint(0, sig_len - noise_len)\n",
    "            full_noise[:, start : start + noise_len] = noise_waveform\n",
    "            noise_waveform = full_noise\n",
    "            # 如果噪音太短，循环填充\n",
    "            # repeats = (sig_len // noise_len) + 1\n",
    "            # noise_waveform = noise_waveform.repeat(1, repeats)[:, :sig_len]\n",
    "            \n",
    "        # 4. 设定随机信噪比 SNR (5dB 到 20dB)\n",
    "        snr_db = random.uniform(5, 20)\n",
    "        \n",
    "        # 5. 混合\n",
    "        return self._mix_at_snr(waveform, noise_waveform, snr_db)\n",
    "\n",
    "    def _mix_at_snr(self, signal: Tensor, noise: Tensor, snr_db: float):\n",
    "        s_p = signal.pow(2).mean()\n",
    "        n_p = noise.pow(2).mean()\n",
    "        snr_linear = 10**(snr_db / 10)\n",
    "        scale = torch.sqrt(s_p / (n_p * snr_linear + 1e-8))\n",
    "        \n",
    "        noisy = signal + scale * noise\n",
    "        # 归一化，防止溢出\n",
    "        return noisy / (noisy.abs().max() + 1e-8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e77899f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ASRInference:\n",
    "    def __init__(self, model_path: Path, vocab_path: Path, noise_dir: 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.noise_augmentor = NoiseAugmentor(noise_root=noise_dir)\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",
    "\n",
    "        # self.gain_up = Gain(min_gain_in_db=5, max_gain_in_db=10, p=1.0, output_type='tensor')\n",
    "        # self.gain_down = Gain(min_gain_in_db=-20, max_gain_in_db=-10, p=1.0, output_type='tensor')\n",
    "        # self.pitch_up = PitchShift(min_transpose_semitones=3, max_transpose_semitones=5, p=1.0, sample_rate=self.sample_rate, output_type='tensor')\n",
    "        # self.pitch_down = PitchShift(min_transpose_semitones=-5, max_transpose_semitones=-3, p=1.0, sample_rate=self.sample_rate, output_type='tensor')\n",
    "        # self.lowpass = LowPassFilter(min_cutoff_freq=400, max_cutoff_freq=2400, p=1.0, output_type='tensor')\n",
    "        # self.highpass = HighPassFilter(min_cutoff_freq=1400, max_cutoff_freq=3400, p=1.0, output_type='tensor')\n",
    "\n",
    "        self.gain_up = Gain(min_gain_in_db=4, max_gain_in_db=8, p=1.0, output_type='tensor')\n",
    "        self.gain_down = Gain(min_gain_in_db=-15, max_gain_in_db=-8, p=1.0, output_type='tensor')\n",
    "        self.pitch_up = PitchShift(min_transpose_semitones=1, max_transpose_semitones=4, p=1.0, sample_rate=self.sample_rate, output_type='tensor')\n",
    "        self.pitch_down = PitchShift(min_transpose_semitones=-4, max_transpose_semitones=-1, p=1.0, sample_rate=self.sample_rate, output_type='tensor')\n",
    "        self.lowpass = LowPassFilter(min_cutoff_freq=600, max_cutoff_freq=2000, p=1.0, output_type='tensor')\n",
    "        self.highpass = HighPassFilter(min_cutoff_freq=800, max_cutoff_freq=2000, p=1.0, output_type='tensor')\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 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.4:\n",
    "            waveform = self.noise_augmentor.apply_real_noise(waveform)\n",
    "\n",
    "        if random.random() < 0.3:\n",
    "            waveform = self._time_mask_waveform(waveform=waveform)\n",
    "        \n",
    "        # torch_audiomentations: [1, time] -> [1, 1, time]\n",
    "        if waveform.dim() == 2:\n",
    "            waveform_3d = waveform.unsqueeze(0)\n",
    "        # 随机选择一种频谱增强\n",
    "        choice = random.random()\n",
    "        if choice < 0.15:\n",
    "            # 增益变化（上或下）\n",
    "            if random.random() < 0.5:\n",
    "                waveform_3d = self.gain_up(waveform_3d, sample_rate=self.sample_rate)\n",
    "            else:\n",
    "                waveform_3d = self.gain_down(waveform_3d, sample_rate=self.sample_rate)\n",
    "        elif choice < 0.25:\n",
    "            # 音高变化（上或下）\n",
    "            if random.random() < 0.5:\n",
    "                waveform_3d = self.pitch_up(waveform_3d, sample_rate=self.sample_rate)\n",
    "            else:\n",
    "                waveform_3d = self.pitch_down(waveform_3d, sample_rate=self.sample_rate)\n",
    "        elif choice < 0.30:\n",
    "            # 低通滤波（声音发闷）\n",
    "            waveform_3d = self.lowpass(waveform_3d, sample_rate=self.sample_rate)\n",
    "        elif choice < 0.35:\n",
    "            # 高通滤波（电话效果）\n",
    "            waveform_3d = self.highpass(waveform_3d, sample_rate=self.sample_rate)\n",
    "        \n",
    "        #  [1, 1, time] -> [1, time]\n",
    "        waveform = waveform_3d.squeeze(0)\n",
    "        \n",
    "        # 防止多次 augment 后振幅溢出，最后归一化\n",
    "        max_amp = waveform.abs().max()\n",
    "        if max_amp > 1.0:\n",
    "            print(max_amp)\n",
    "            waveform = waveform / max_amp\n",
    "        \n",
    "        return waveform\n",
    "    \n",
    "    def _voice_stretch_or_compress1(self, waveform: Tensor) -> Tensor:\n",
    "        speed_factor = random.uniform(0.8, 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.8, 1.25) # (Speed Change: 0.6x - 1.2x)\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(hop_length=160, n_freq=spec.shape[-2], fixed_rate=1.4)\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",
    "        y_stretch = pyrb.time_stretch(waveform_np, self.sample_rate, speed_factor)\n",
    "        return torch.from_numpy(y_stretch).float().to(waveform.device).unsqueeze(0)\n",
    "    \n",
    "    def _time_mask_waveform(self, waveform: Tensor) -> Tensor:\n",
    "        audio_len = waveform.shape[1]\n",
    "        sr = self.sample_rate # 16000\n",
    "        \n",
    "        # 设置参数：单次遮盖最长 0.4 秒 (6400个点)\n",
    "        max_mask_time = 0.4\n",
    "        max_mask_samples = int(sr * max_mask_time)\n",
    "        \n",
    "        # 根据音频长度决定遮盖次数：\n",
    "        # 比如每 3 秒钟允许遮盖 1 次\n",
    "        num_masks = max(1, audio_len // (sr * 3)) \n",
    "        \n",
    "        for _ in range(num_masks):\n",
    "            # 每次随机遮盖 0.1s 到 0.4s\n",
    "            current_mask_len = random.randint(int(sr * 0.1), max_mask_samples)\n",
    "            \n",
    "            if audio_len > current_mask_len:\n",
    "                start_pos = random.randint(0, audio_len - current_mask_len)\n",
    "                \n",
    "                # 填充微小噪音（模拟环境底噪）\n",
    "                noise = torch.randn(1, current_mask_len).to(waveform.device) * 0.002\n",
    "                waveform[:, start_pos : start_pos + current_mask_len] = noise\n",
    "                \n",
    "        return waveform\n",
    "\n",
    "    def transcribe(self, waveform: Tensor) -> str:\n",
    "        waveform = waveform.to(device=self.device)\n",
    "        waveform_length = torch.tensor([waveform.shape[1]], dtype=torch.long, device=self.device)\n",
    "\n",
    "        with no_grad():\n",
    "            log_probs, _ = self.model(waveforms=waveform, waveform_lengths=waveform_length)\n",
    "        \n",
    "        text = self.tokenizer.ctc_greedy_decode(log_probs=log_probs[0])\n",
    "        return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5af7c836",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda:1')\n",
    "\n",
    "# checkpoint = sorted(workspace_dir.glob('.checkpoints/checkpoint_epoch_*.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_cer_model.pt\"\n",
    "checkpoint = workspace_dir / \".checkpoints/prodect/checkpoint_epoch_24 copy.pt\"\n",
    "print(f\"Load Checkpoint: {checkpoint}\")\n",
    "inference = ASRInference(model_path=checkpoint, vocab_path=workspace_dir / 'config/asr_vocab.json', noise_dir='/mnt/dataset/dataset/audio/noise', device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "967e23cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_42681424.mp3'\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_26188445.mp3'\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_43349581.mp3'\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_42640878.mp3'\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_26245615.mp3'\n",
    "# audio_path = workspace_dir / '.data/ug/clips/common_voice_ug_40794549.mp3'\n",
    "# audio_path = workspace_dir / 'data/ug/clips/common_voice_ug_26245614.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_1.m4a'\n",
    "# audio_path = workspace_dir / 'data/test/let_me_see_the_computer.mp3'\n",
    "# audio_path = workspace_dir / 'data/test/introduce_myself.mp3'\n",
    "# audio_path = workspace_dir / 'data/test/F001_001.wav'\n",
    "# audio_path = workspace_dir / 'data/test/download.wav'\n",
    "# audio_path = workspace_dir / 'data/test/test_voise_1.m4a'\n",
    "orginal_waveform = inference._load_audio(audio_path=audio_path)\n",
    "print('load audio:', orginal_waveform.shape)\n",
    "# orginal_waveform = orginal_waveform[:, :]\n",
    "\n",
    "def simple_corridor_echo(waveform, sample_rate=16000, delay_ms=80, decay=0.3):\n",
    "    \"\"\"\n",
    "    模拟一次明显的走廊回声：原始声音 + 一个衰减延迟的拷贝\n",
    "    delay_ms: 延迟毫秒数，80ms 模拟中等走廊\n",
    "    decay: 回声的衰减系数，0.2~0.4 比较自然\n",
    "    \"\"\"\n",
    "    delay_samples = int(sample_rate * delay_ms / 1000)\n",
    "    echo = torch.zeros_like(waveform)\n",
    "    echo[..., delay_samples:] = waveform[..., :-delay_samples] * decay\n",
    "    return waveform + echo\n",
    "\n",
    "print(\"Orginal audio\")\n",
    "display(Audio(orginal_waveform, rate=inference.sample_rate))\n",
    "\n",
    "# augment_waveform = inference.augment_waveform(waveform=orginal_waveform)\n",
    "# augment_waveform = simple_corridor_echo(waveform=orginal_waveform)\n",
    "# augment_waveform = inference._voice_stretch_or_compress(waveform=orginal_waveform)\n",
    "# noise_augmentor_waveform = inference.noise_augmentor.apply_real_noise(waveform=orginal_waveform)\n",
    "# display(Audio(augment_waveform, rate=inference.sample_rate))\n",
    "# text = inference.transcribe(waveform=orginal_waveform)\n",
    "# print(f\"\\n{text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16195280",
   "metadata": {},
   "outputs": [],
   "source": [
    "from silero_vad import load_silero_vad, read_audio, get_speech_timestamps\n",
    "class TimestampsType(TypedDict):\n",
    "    start: int\n",
    "    end: int\n",
    "\n",
    "THRESHOLD_MS = 2000\n",
    "SAMPLE_RATE = inference.sample_rate\n",
    "\n",
    "def merge_timestamps(timestamps: int, threshold_ms: int, sample_rate: int):\n",
    "    if not timestamps:\n",
    "        return []\n",
    "    threshold_samples = (threshold_ms / 1000) * sample_rate\n",
    "    merged = []\n",
    "    curr_start = timestamps[0]['start']\n",
    "    curr_end = timestamps[0]['end']\n",
    "    \n",
    "    for i in range(1, len(timestamps)):\n",
    "        # 如果当前积攒的长度不到 2 秒，就一直合并到当前的 end 上\n",
    "        if (curr_end - curr_start) < threshold_samples:\n",
    "            curr_end = timestamps[i]['end']\n",
    "        else:\n",
    "            merged.append({'start': curr_start, 'end': curr_end})\n",
    "            curr_start = timestamps[i]['start']\n",
    "            curr_end = timestamps[i]['end']\n",
    "    merged.append({'start': curr_start, 'end': curr_end})\n",
    "    return merged\n",
    "\n",
    "vad_model = load_silero_vad(onnx=False)\n",
    "waveform_vad = inference._load_audio(audio_path=audio_path)\n",
    "speech_timestamps: list[TimestampsType] = get_speech_timestamps(\n",
    "    waveform_vad, \n",
    "    vad_model, \n",
    "    sampling_rate=inference.sample_rate,\n",
    "    threshold=0.4,  # 可以根据需要调整阈值\n",
    "    min_speech_duration_ms=100, # 最小语音持续时间，防止短噪音被误判\n",
    "    min_silence_duration_ms=200, # 最小静音间隔，用于分割语音块\n",
    "    speech_pad_ms=100, # 在语音块前后添加的填充时间\n",
    ")\n",
    "print('speech_timestamps:', speech_timestamps)\n",
    "print('speech_timestamps:', len(speech_timestamps))\n",
    "final_timestamps = merge_timestamps(speech_timestamps, THRESHOLD_MS, SAMPLE_RATE)\n",
    "print(f'VAD原始片段: {len(speech_timestamps)} -> 合并后片段: {len(final_timestamps)}')\n",
    "for ts in final_timestamps:\n",
    "    ts: TimestampsType\n",
    "    start_sample = int(ts['start'])\n",
    "    end_sample = int(ts['end'])\n",
    "\n",
    "    segment_waveform = waveform_vad[:, start_sample:end_sample]\n",
    "    display(Audio(segment_waveform, rate=inference.sample_rate))\n",
    "\n",
    "    segment_text = inference.transcribe(waveform=segment_waveform)\n",
    "    print(segment_text, end=\" \", flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3671d334",
   "metadata": {},
   "outputs": [],
   "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",
    "# mel_spec = inference._augment_spec(mel_spec=mel_spec)\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)"
   ]
  }
 ],
 "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
}
