helo-ayush/Diarization_VoiceFingerprinted

1# ==============================================================================
2# AUDIO PROCESSOR
3# Handles noise reduction, VAD (Voice Activity Detection), and audio conversion.
4# This prevents garbage audio from being sent to transcription APIs.
5#
6# FIX: Previously used collect_chunks() which hard-concatenated speech segments
7# with zero gap between them, causing acoustic discontinuity that confused
8# Sarvam's language model context window. Now replaces removed silence with a
9# short 300ms synthetic pad + 10ms fade in/out at each boundary.
10# ==============================================================================
11import io
12import subprocess
13import tempfile
14import os
15import numpy as np
16import soundfile as sf
17import noisereduce as nr
18
19import torch
20import warnings
21
22warnings.filterwarnings("ignore")
23
24# ==============================================================================
25# VAD CONFIGURATION
26# ==============================================================================
27# How long of a silence gap to insert between joined speech chunks.
28# 300ms sounds like a natural short pause — long enough for Sarvam's context
29# window to register a speaker breath, short enough to still reduce file size.
30SILENCE_PAD_BETWEEN_CHUNKS_S = 0.3
31
32# How many samples to use for fade in/out at each chunk boundary.
33# 10ms prevents amplitude clicking at hard splice points.
34FADE_DURATION_S = 0.01
35
36print("⏳ Loading Silero VAD model (downloading if first time)...")
37vad_model, vad_utils = torch.hub.load(
38    repo_or_dir="snakers4/silero-vad",
39    model="silero_vad",
40    force_reload=False,
41    onnx=False
42)
43(get_speech_timestamps, save_audio, read_audio, VADIterator, collect_chunks) = vad_utils
44print("✅ Silero VAD loaded.")
45
46
47# ==============================================================================
48# INTERNAL HELPERS
49# ==============================================================================
50
51def _apply_fade(chunk: np.ndarray, fade_samples: int) -> np.ndarray:
52    """
53    Apply a short linear fade-in at the start and fade-out at the end of a
54    chunk. Removes the audible click/pop that occurs when two audio segments
55    are hard-spliced together at different amplitude levels.
56    """
57    if len(chunk) < fade_samples * 2:
58        # Chunk is too short to safely apply fade without overlap — skip it
59        return chunk
60
61    c = chunk.copy()
62    c[:fade_samples]  *= np.linspace(0.0, 1.0, fade_samples)   # fade in
63    c[-fade_samples:] *= np.linspace(1.0, 0.0, fade_samples)   # fade out
64    return c
65
66
67def _join_speech_chunks(
68    speech_timestamps: list,
69    tensor_samples: torch.Tensor,
70    sample_rate: int,
71) -> np.ndarray:
72    """
73    Joins speech segments from VAD with a short silence pad between each one,
74    instead of hard-concatenating them with zero gap (old collect_chunks behavior).
75
76    This preserves natural acoustic pacing so Sarvam's language model doesn't
77    see abrupt teleporting between unrelated audio frames.
78    """
79    fade_samples = int(FADE_DURATION_S * sample_rate)
80    silence_pad  = np.zeros(
81        int(SILENCE_PAD_BETWEEN_CHUNKS_S * sample_rate),
82        dtype=np.float32
83    )
84
85    chunks = []
86    for i, ts in enumerate(speech_timestamps):
87        # Extract this speech segment as a numpy array
88        chunk = tensor_samples[ts["start"] : ts["end"]].numpy()
89
90        # Apply fade in/out to smooth the amplitude at splice boundaries
91        chunk = _apply_fade(chunk, fade_samples)
92
93        chunks.append(chunk)
94
95        # Insert silence gap between chunks (not after the last one)
96        if i < len(speech_timestamps) - 1:
97            chunks.append(silence_pad)
98
99    return np.concatenate(chunks)
100
101
102# ==============================================================================
103# MAIN PUBLIC FUNCTION
104# ==============================================================================
105
106def clean_audio_for_diarization(audio_bytes: bytes, original_filename: str = "input.wav") -> bytes:
107    """
108    Clean audio using ML-based VAD (Silero) + Noise Reduction (noisereduce).
109
110    Pipeline:
111    1. ffmpeg converts any input format → WAV (16kHz, mono) in memory
112    2. Silero VAD detects speech segments
113    3. Speech segments are joined with 300ms silence pads + 10ms fades
114       (replaces the old hard-concat collect_chunks approach)
115    4. noisereduce applies ML spectral gating noise reduction
116    5. Highpass (80Hz) + Lowpass (8kHz) filtering via FFT
117    6. ffmpeg compresses result → OGG/Opus for fast upload
118    """
119    print(f"🔧 Cleaning audio '{original_filename}' with ML noise reduction & VAD...")
120
121    # ── Step 1: Convert input to WAV 16kHz mono via ffmpeg ──────────────────
122    ext = os.path.splitext(original_filename)[1] or ".input"
123    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp_in:
124        tmp_in.write(audio_bytes)
125        tmp_in_path = tmp_in.name
126
127    tmp_wav_path   = tmp_in_path + ".wav"
128    tmp_clean_path = tmp_in_path + ".clean.wav"
129    tmp_ogg_path   = tmp_in_path + ".ogg"
130
131    try:
132        subprocess.run(
133            [
134                "ffmpeg", "-y", "-i", tmp_in_path,
135                "-ac", "1",            # mono
136                "-ar", "16000",        # 16 kHz
137                "-sample_fmt", "s16",  # 16-bit PCM
138                tmp_wav_path,
139            ],
140            capture_output=True,
141            check=True,
142        )
143
144        # ── Step 2: Load WAV into numpy ──────────────────────────────────────
145        samples, sample_rate = sf.read(tmp_wav_path, dtype="float32")
146        original_dur = len(samples) / sample_rate
147
148        # ── Step 3: Silero VAD — detect speech timestamps ───────────────────
149        tensor_samples = torch.from_numpy(samples)
150
151        speech_timestamps = get_speech_timestamps(
152            tensor_samples,
153            vad_model,
154            sampling_rate=sample_rate,
155            threshold=0.3,                # Lower = catch quieter speech
156            min_speech_duration_ms=100,   # Catch short words like "haan", "ok"
157            min_silence_duration_ms=700,  # Don't split on pauses < 700ms
158            speech_pad_ms=250,            # 250ms padding around each chunk
159        )
160
161        if not speech_timestamps:
162            print("⚠️  No speech detected by Silero VAD — using original audio.")
163            speech_samples = samples
164        else:
165            # ── Step 3b: Join chunks with silence pads (THE KEY FIX) ────────
166            speech_samples = _join_speech_chunks(
167                speech_timestamps,
168                tensor_samples,
169                sample_rate,
170            )
171            new_dur = len(speech_samples) / sample_rate
172            saved   = original_dur - new_dur
173            print(
174                f"✂️  VAD trimming: {original_dur:.1f}s → {new_dur:.1f}s "
175                f"({saved:.1f}s of silence removed, "
176                f"{len(speech_timestamps)} chunks joined with {SILENCE_PAD_BETWEEN_CHUNKS_S*1000:.0f}ms pads)"
177            )
178
179        # ── Step 4: ML Noise Reduction ───────────────────────────────────────
180        cleaned = nr.reduce_noise(
181            y=speech_samples,
182            sr=sample_rate,
183            prop_decrease=0.85,      # Remove 85% of noise
184            stationary=False,        # Handle non-stationary noise (traffic, fans)
185            n_fft=2048,
186            freq_mask_smooth_hz=500,
187        )
188
189        # ── Step 5: Highpass (80Hz) + Lowpass (8kHz) via FFT ─────────────────
190        fft   = np.fft.rfft(cleaned)
191        freqs = np.fft.rfftfreq(len(cleaned), d=1 / sample_rate)
192        fft[freqs < 80]   = 0   # cut sub-bass rumble
193        fft[freqs > 8000] = 0   # cut high-frequency hiss
194        cleaned = np.fft.irfft(fft, n=len(cleaned))
195
196        # Normalize to prevent clipping
197        peak = np.max(np.abs(cleaned))
198        if peak > 0:
199            cleaned = cleaned / peak
200
201        # ── Step 6: Write cleaned WAV ─────────────────────────────────────────
202        sf.write(tmp_clean_path, cleaned, sample_rate, subtype="PCM_16")
203
204        # ── Step 7: Compress to OGG/Opus ─────────────────────────────────────
205        subprocess.run(
206            [
207                "ffmpeg", "-y", "-i", tmp_clean_path,
208                "-c:a", "libopus",
209                "-b:a", "32k",
210                tmp_ogg_path,
211            ],
212            capture_output=True,
213            check=True,
214        )
215
216        with open(tmp_ogg_path, "rb") as f:
217            ogg_bytes = f.read()
218
219        original_kb = len(audio_bytes) / 1024
220        cleaned_kb  = len(ogg_bytes)   / 1024
221        print(
222            f"✅ Audio cleaned: {original_kb:.0f} KB → {cleaned_kb:.0f} KB "
223            f"(ML denoised + VAD-trimmed + compressed)"
224        )
225
226        return ogg_bytes
227
228    finally:
229        # Always clean up temp files even if an exception occurred
230        for path in [tmp_in_path, tmp_wav_path, tmp_clean_path, tmp_ogg_path]:
231            try:
232                os.unlink(path)
233            except (OSError, UnboundLocalError):
234                pass