314 lines
12 KiB
Python
314 lines
12 KiB
Python
"""Streaming capture conditioning and pause-aware chunking for Blurt.
|
|
|
|
The class deliberately has no UI or sounddevice dependencies. ``feed`` may be
|
|
called with whatever block size the input device supplies; all public sample
|
|
positions are in the 16 kHz output timebase.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from math import gcd
|
|
import logging
|
|
|
|
import numpy as np
|
|
from scipy import signal
|
|
|
|
|
|
OUTPUT_RATE = 16_000
|
|
FRAME = OUTPUT_RATE // 50 # 20 ms
|
|
MIN_SECONDS = 8.0
|
|
TARGET_SECONDS = 20.0
|
|
MAX_SECONDS = 45.0
|
|
EARLY_PAUSE_SECONDS = 0.8
|
|
LATE_PAUSE_SECONDS = 0.5
|
|
HIGHPASS_HZ = 70.0
|
|
TARGET_RMS = 10 ** (-20.0 / 20.0)
|
|
PEAK_CEILING = 10 ** (-1.0 / 20.0)
|
|
MAX_GAIN = 10 ** (30.0 / 20.0)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AudioChunk:
|
|
"""A contiguous section of conditioned audio in the session timebase."""
|
|
|
|
samples: np.ndarray
|
|
start: int
|
|
end: int
|
|
speech: bool
|
|
|
|
|
|
def _mono(frames: np.ndarray, channels: int) -> np.ndarray:
|
|
data = np.asarray(frames, dtype=np.float32)
|
|
if data.ndim == 1:
|
|
if channels != 1:
|
|
raise ValueError("one-dimensional frames are valid only for mono capture")
|
|
return data
|
|
if data.ndim != 2 or data.shape[1] != channels:
|
|
raise ValueError(f"expected frames shaped (n, {channels})")
|
|
return data.mean(axis=1, dtype=np.float32)
|
|
|
|
|
|
def _levels(samples: np.ndarray) -> np.ndarray:
|
|
count = len(samples) // FRAME
|
|
if not count:
|
|
return np.empty(0, dtype=np.float32)
|
|
framed = samples[: count * FRAME].reshape(count, FRAME)
|
|
framed = framed - framed.mean(axis=1, keepdims=True)
|
|
rms = np.sqrt(np.mean(framed * framed, axis=1))
|
|
return (20.0 * np.log10(np.maximum(rms, 1e-7))).astype(np.float32)
|
|
|
|
|
|
def _normalize(samples: np.ndarray) -> np.ndarray:
|
|
"""Normalize one sealed, already high-pass-filtered region."""
|
|
if not len(samples):
|
|
return np.empty(0, dtype=np.int16)
|
|
filtered = samples
|
|
rms = float(np.sqrt(np.mean(filtered * filtered)))
|
|
peak = float(np.percentile(np.abs(filtered), 99.9))
|
|
gain = min(TARGET_RMS / rms, PEAK_CEILING / peak, MAX_GAIN) if rms and peak else 1.0
|
|
scaled = np.clip(filtered * gain, -1.0, 1.0)
|
|
return np.rint(scaled * 32767.0).astype(np.int16)
|
|
|
|
|
|
class StreamingAudio:
|
|
"""Convert native-rate blocks into pause-aware 16 kHz mono chunks.
|
|
|
|
Native input is resampled in one-second, phase-aligned batches with 50 ms
|
|
of context on both sides. Thus callback block size does not affect output
|
|
length, while filter-edge errors stay outside the retained middle. A final
|
|
partial batch is length-corrected to the rounded session duration. Startup
|
|
audio is deliberately retained: without evidence of a device click, an
|
|
automatic trim risks deleting the first quiet consonant.
|
|
"""
|
|
|
|
_GUARD = FRAME * 2
|
|
_CONTEXT_SECONDS = 0.05
|
|
|
|
def __init__(self, rate: int, channels: int):
|
|
if rate <= 0 or channels <= 0:
|
|
raise ValueError("rate and channels must be positive")
|
|
self.rate = int(rate)
|
|
self.channels = int(channels)
|
|
self._pending: list[np.ndarray] = []
|
|
self._pending_len = 0
|
|
self._left_native = np.empty(0, dtype=np.float32)
|
|
self._resampled = np.empty(0, dtype=np.float32)
|
|
self._base = 0
|
|
self._start = 0 # absolute 16 kHz sample position
|
|
self._output_total = 0
|
|
self._sos = signal.butter(
|
|
2, HIGHPASS_HZ, btype="highpass", fs=OUTPUT_RATE, output="sos"
|
|
)
|
|
self._filter_zi = None
|
|
self._finished = False
|
|
|
|
def feed(self, frames: np.ndarray) -> list[AudioChunk]:
|
|
if self._finished:
|
|
raise RuntimeError("cannot feed a finished stream")
|
|
mono = _mono(frames, self.channels)
|
|
previous_available = self._base + len(self._resampled)
|
|
if len(mono):
|
|
self._pending.append(mono.copy())
|
|
self._pending_len += len(mono)
|
|
self._refresh(final=False)
|
|
# Completed native batches already used right-side resampler context.
|
|
stable = self._base + len(self._resampled)
|
|
if stable == previous_available:
|
|
return []
|
|
return self._emit(stable, finishing=False)
|
|
|
|
def finish(self) -> list[AudioChunk]:
|
|
if self._finished:
|
|
return []
|
|
self._finished = True
|
|
self._refresh(final=True)
|
|
chunks = self._emit(self._base + len(self._resampled), finishing=True)
|
|
end = self._base + len(self._resampled)
|
|
if self._start < end:
|
|
chunks.append(self._make_chunk(self._start, end))
|
|
self._start = end
|
|
self._discard_sealed()
|
|
return chunks
|
|
|
|
def checkpoint(self) -> np.ndarray:
|
|
"""Return a conditioned copy of the currently unsealed audio."""
|
|
preview = self._preview()
|
|
local = self._start - self._base
|
|
return _normalize(preview[local:]).copy()
|
|
|
|
def _resample(self, native: np.ndarray) -> np.ndarray:
|
|
if self.rate == OUTPUT_RATE:
|
|
return native.astype(np.float32, copy=True)
|
|
divisor = gcd(self.rate, OUTPUT_RATE)
|
|
return signal.resample_poly(
|
|
native, OUTPUT_RATE // divisor, self.rate // divisor
|
|
).astype(np.float32)
|
|
|
|
def _refresh(self, final: bool) -> None:
|
|
batch = self.rate
|
|
context = round(self.rate * self._CONTEXT_SECONDS)
|
|
while self._pending_len >= batch + (0 if final else context):
|
|
pending = np.concatenate(self._pending)
|
|
take = min(batch, self._pending_len)
|
|
right = min(context, self._pending_len - take)
|
|
left = len(self._left_native)
|
|
source = np.concatenate((self._left_native, pending[: take + right]))
|
|
converted = self._resample(source)
|
|
offset = round(left * OUTPUT_RATE / self.rate)
|
|
wanted = round(take * OUTPUT_RATE / self.rate)
|
|
self._append_filtered(converted[offset : offset + wanted])
|
|
consumed = pending[:take]
|
|
self._left_native = np.concatenate((self._left_native, consumed))[-context:]
|
|
remainder = pending[take:]
|
|
self._pending = [remainder] if len(remainder) else []
|
|
self._pending_len = len(remainder)
|
|
if final and self._pending_len:
|
|
pending = np.concatenate(self._pending)
|
|
left = len(self._left_native)
|
|
converted = self._resample(np.concatenate((self._left_native, pending)))
|
|
offset = round(left * OUTPUT_RATE / self.rate)
|
|
wanted = round(len(pending) * OUTPUT_RATE / self.rate)
|
|
tail = converted[offset : offset + wanted]
|
|
self._append_filtered(np.pad(tail, (0, max(0, wanted - len(tail)))))
|
|
self._pending = []
|
|
self._pending_len = 0
|
|
|
|
def _append_filtered(self, samples: np.ndarray) -> None:
|
|
if not len(samples):
|
|
return
|
|
if self._filter_zi is None:
|
|
self._filter_zi = signal.sosfilt_zi(self._sos) * float(samples[0])
|
|
filtered, self._filter_zi = signal.sosfilt(
|
|
self._sos, samples, zi=self._filter_zi
|
|
)
|
|
self._resampled = np.concatenate((self._resampled, filtered.astype(np.float32)))
|
|
self._output_total += len(samples)
|
|
|
|
def _preview(self) -> np.ndarray:
|
|
"""Resample the outstanding tail without advancing stream state."""
|
|
if not self._pending_len:
|
|
return self._resampled
|
|
context = round(self.rate * self._CONTEXT_SECONDS)
|
|
pending = np.concatenate(self._pending)
|
|
left = len(self._left_native)
|
|
source = np.concatenate((self._left_native, pending))
|
|
converted = self._resample(source)
|
|
offset = round(left * OUTPUT_RATE / self.rate)
|
|
wanted = round(len(pending) * OUTPUT_RATE / self.rate)
|
|
tail = converted[offset : offset + wanted]
|
|
zi = self._filter_zi
|
|
if zi is None and len(tail):
|
|
zi = signal.sosfilt_zi(self._sos) * float(tail[0])
|
|
if zi is not None and len(tail):
|
|
tail = signal.sosfilt(self._sos, tail, zi=zi)[0]
|
|
return np.concatenate((self._resampled, tail.astype(np.float32)))
|
|
|
|
def _emit(self, available: int, finishing: bool) -> list[AudioChunk]:
|
|
chunks: list[AudioChunk] = []
|
|
while True:
|
|
cut = self._next_cut(available)
|
|
if cut is None:
|
|
break
|
|
chunks.append(self._make_chunk(self._start, cut))
|
|
self._start = cut
|
|
self._discard_sealed()
|
|
# finish() leaves the sub-minimum tail to its caller so it is emitted
|
|
# once, after every pause/forced boundary has been handled.
|
|
return chunks
|
|
|
|
def _next_cut(self, available: int) -> int | None:
|
|
length = available - self._start
|
|
minimum = int(MIN_SECONDS * OUTPUT_RATE)
|
|
maximum = int(MAX_SECONDS * OUTPUT_RATE)
|
|
if length < minimum:
|
|
return None
|
|
if length >= maximum:
|
|
search_end = self._start + maximum
|
|
else:
|
|
search_end = available
|
|
|
|
segment = self._resampled[self._start - self._base : search_end - self._base]
|
|
levels = _levels(segment)
|
|
if not len(levels):
|
|
return None
|
|
# The lower quartile follows different microphones/noise floors. The
|
|
# ceiling prevents ordinary quiet speech from being labelled a pause.
|
|
floor = float(np.percentile(levels, 25))
|
|
quiet_below = min(-45.0, floor + 6.0)
|
|
quiet = levels <= quiet_below
|
|
run = 0
|
|
first_eligible = int(MIN_SECONDS * 50)
|
|
for frame_index, is_quiet in enumerate(quiet):
|
|
run = run + 1 if is_quiet else 0
|
|
age = (frame_index + 1) / 50.0
|
|
fraction = np.clip(
|
|
(age - MIN_SECONDS) / (TARGET_SECONDS - MIN_SECONDS), 0.0, 1.0
|
|
)
|
|
required = EARLY_PAUSE_SECONDS + fraction * (
|
|
LATE_PAUSE_SECONDS - EARLY_PAUSE_SECONDS
|
|
)
|
|
required_frames = int(np.ceil(required * 50))
|
|
if frame_index + 1 >= first_eligible and run >= required_frames:
|
|
# Cut in the middle of the accepted pause. Both adjacent
|
|
# ranges still retain their half, so coverage remains exact.
|
|
cut_frame = frame_index + 1 - run // 2
|
|
cut = self._start + max(minimum, cut_frame * FRAME)
|
|
logging.info(
|
|
"Streaming chunk pause cut at %.2fs (age %.2fs, pause %.2fs, "
|
|
"threshold %.1f dBFS)",
|
|
cut / OUTPUT_RATE,
|
|
age,
|
|
run / 50.0,
|
|
quiet_below,
|
|
)
|
|
return cut
|
|
if length >= maximum:
|
|
# Prefer the quietest 20 ms frame in the final second when speech
|
|
# has been uninterrupted, while strictly respecting the limit.
|
|
end_frame = maximum // FRAME
|
|
begin_frame = max(first_eligible, end_frame - 50)
|
|
chosen = begin_frame + int(np.argmin(levels[begin_frame:end_frame]))
|
|
cut = self._start + min((chosen + 1) * FRAME, maximum)
|
|
logging.info(
|
|
"Streaming chunk forced cut at %.2fs after %.2fs without a pause",
|
|
cut / OUTPUT_RATE,
|
|
length / OUTPUT_RATE,
|
|
)
|
|
return cut
|
|
return None
|
|
|
|
def _make_chunk(self, start: int, end: int) -> AudioChunk:
|
|
raw = self._resampled[start - self._base : end - self._base]
|
|
levels = _levels(raw)
|
|
# -80 dBFS distinguishes actual/near digital silence while preserving
|
|
# quiet but valid speech. Consumers spool false chunks but skip upload.
|
|
if len(levels):
|
|
speech_level = float(np.max(levels))
|
|
elif len(raw):
|
|
centered = raw - float(np.mean(raw))
|
|
speech_level = float(20 * np.log10(max(np.sqrt(np.mean(centered * centered)), 1e-7)))
|
|
else:
|
|
speech_level = -140.0
|
|
speech = speech_level > -80.0
|
|
floor = float(np.percentile(levels, 10)) if len(levels) else -140.0
|
|
longest = run = 0
|
|
for level in levels:
|
|
run = run + 1 if level < -75.0 else 0
|
|
longest = max(longest, run)
|
|
logging.info(
|
|
"Streaming chunk %.2f-%.2fs speech=%s floor=%.1f dBFS dead=%.2fs",
|
|
start / OUTPUT_RATE,
|
|
end / OUTPUT_RATE,
|
|
speech,
|
|
floor,
|
|
longest / 50.0,
|
|
)
|
|
return AudioChunk(_normalize(raw), start, end, speech)
|
|
|
|
def _discard_sealed(self) -> None:
|
|
count = self._start - self._base
|
|
if count > 0:
|
|
self._resampled = self._resampled[count:]
|
|
self._base = self._start
|