File size: 4,224 Bytes
2a8d9cd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
"""Audio post-processing: silence trimming, normalization, quality checks."""

from __future__ import annotations

import io
import os
import tempfile

import numpy as np


def trim_silence(
    audio: np.ndarray,
    sample_rate: int,
    threshold_db: float = -35.0,
    min_silence_sec: float = 0.12,
    max_silence_sec: float = 0.20,
    keep_start: float = 0.03,
    keep_end: float = 0.05,
) -> np.ndarray:
    """Trim leading/trailing silence and compress long pauses between speech segments.

    Args:
        audio: 1-D float32 waveform.
        sample_rate: Samples per second.
        threshold_db: Silence threshold in dB (below this = silence).
        min_silence_sec: Minimum pause to keep (natural breathing room).
        max_silence_sec: Maximum pause allowed — anything longer gets capped here.
        keep_start: Seconds of silence to keep at the start.
        keep_end: Seconds of silence to keep at the end.

    Returns:
        Cleaned waveform (float32).
    """
    if audio.size == 0:
        return audio

    threshold = 10 ** (threshold_db / 20.0)
    sr = sample_rate

    # --- find voiced regions ---
    is_voiced = np.abs(audio) > threshold
    voiced_indices = np.where(is_voiced)[0]

    if voiced_indices.size == 0:
        return audio[: int(sr * 0.1)]

    # keep a small buffer before first voice and after last voice
    start = max(0, int(voiced_indices[0] - sr * keep_start))
    end = min(len(audio), int(voiced_indices[-1] + sr * keep_end))
    trimmed = audio[start:end].copy()

    # --- compress long inter-voice pauses ---
    target_pause = int(min_silence_sec * sr)
    max_pause = int(max_silence_sec * sr)
    result_parts: list[np.ndarray] = []
    fade_len = int(0.005 * sr)  # 5ms crossfade to avoid clicks
    i = 0
    n = len(trimmed)

    while i < n:
        if np.abs(trimmed[i]) <= threshold:
            silence_end = i
            while silence_end < n and np.abs(trimmed[silence_end]) <= threshold:
                silence_end += 1
            silence_len = silence_end - i
            if silence_len > max_pause:
                result_parts.append(trimmed[i: i + target_pause])
            elif silence_len > target_pause:
                result_parts.append(trimmed[i: i + target_pause])
            else:
                result_parts.append(trimmed[i: silence_end])
            i = silence_end
        else:
            voice_end = i
            while voice_end < n and np.abs(trimmed[voice_end]) > threshold:
                voice_end += 1
            result_parts.append(trimmed[i: voice_end])
            i = voice_end

    if not result_parts:
        return trimmed

    out = np.concatenate(result_parts).astype(np.float32)

    # fade-in / fade-out to avoid clicks at trim points
    if len(out) > fade_len * 2:
        fade_in = np.linspace(0, 1, fade_len, dtype=np.float32)
        fade_out = np.linspace(1, 0, fade_len, dtype=np.float32)
        out[:fade_len] *= fade_in
        out[-fade_len:] *= fade_out

    return out


def normalize_audio(
    audio: np.ndarray,
    target_peak_db: float = -1.0,
) -> np.ndarray:
    """Peak-normalize audio to a target level."""
    if audio.size == 0:
        return audio
    peak = float(np.max(np.abs(audio)))
    if peak < 1e-6:
        return audio
    target = 10 ** (target_peak_db / 20.0)
    return (audio * (target / peak)).astype(np.float32)


def postprocess(
    audio: np.ndarray,
    sample_rate: int,
    trim: bool = True,
    normalize: bool = True,
) -> np.ndarray:
    """Full post-processing pipeline: trim silence + normalize."""
    if audio.size == 0:
        return audio
    if trim:
        audio = trim_silence(audio, sample_rate)
    if normalize:
        audio = normalize_audio(audio)
    return audio


def postprocess_wav_bytes(
    wav_bytes: bytes,
    sample_rate: int,
    trim: bool = True,
    normalize: bool = True,
) -> bytes:
    """Post-process WAV bytes (used on Modal GPU before sending back)."""
    import soundfile as sf

    audio, sr = sf.read(io.BytesIO(wav_bytes), dtype="float32")
    audio = postprocess(audio, sr, trim=trim, normalize=normalize)
    buf = io.BytesIO()
    sf.write(buf, audio, sr, format="WAV")
    return buf.getvalue()