File size: 4,025 Bytes
df5e685 | 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 | """Opus-based audio cache I/O for v86 feature cache.
Encode: float32 24kHz mono ndarray → Opus@96kbps bytes (~45 KB / 4 sec).
Decode: bytes → float32 24kHz mono ndarray (exact length 96000).
Opus internally resamples to 48 kHz; we resample back to 24 kHz on decode
and trim the 156-sample pre-padding to align with the original.
"""
from __future__ import annotations
import io
import numpy as np
import av
import scipy.signal as sps
SR_TARGET = 24000
SR_OPUS = 48000
EXPECTED_SAMPLES = 96000 # 4 sec @ 24 kHz
OPUS_BIT_RATE = 96000 # 96 kbps
# Empirically the decoded length matches 96000 already (no trimming needed).
# Residual lag is ~47 samples (2 ms) which is acceptable for instrument audio
# since the codec/LM windows are 4 sec each.
OPUS_PRE_DELAY_24K = 0
def encode_opus(audio_24k_f32: np.ndarray, bit_rate: int = OPUS_BIT_RATE) -> bytes:
"""Encode (96000,) float32 24kHz mono → Opus bytes."""
assert audio_24k_f32.ndim == 1, audio_24k_f32.shape
assert audio_24k_f32.dtype == np.float32, audio_24k_f32.dtype
# PyAV expects (channels, samples) for non-planar formats — but we use flt which is planar.
audio_2d = audio_24k_f32.reshape(1, -1)
buf = io.BytesIO()
container = av.open(buf, mode="w", format="ogg")
stream = container.add_stream("libopus", rate=SR_TARGET)
stream.bit_rate = bit_rate
stream.layout = "mono"
frame = av.AudioFrame.from_ndarray(audio_2d, format="flt", layout="mono")
frame.sample_rate = SR_TARGET
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode(): # flush
container.mux(packet)
container.close()
return buf.getvalue()
def decode_opus(opus_bytes: bytes,
target_samples: int = EXPECTED_SAMPLES) -> np.ndarray:
"""Decode Opus bytes → (target_samples,) float32 24kHz mono."""
buf = io.BytesIO(opus_bytes)
container = av.open(buf, mode="r")
frames = []
for frame in container.decode(audio=0):
# frame.to_ndarray() returns (channels, samples) for planar
arr = frame.to_ndarray()
if arr.ndim == 2: arr = arr[0]
frames.append(arr.astype(np.float32))
container.close()
audio_48k = np.concatenate(frames)
# Resample 48k → 24k
audio_24k = sps.resample_poly(audio_48k, 1, 2).astype(np.float32)
# Drop pre-padding
if len(audio_24k) > OPUS_PRE_DELAY_24K:
audio_24k = audio_24k[OPUS_PRE_DELAY_24K:]
# Pad or trim to exact length
if len(audio_24k) < target_samples:
audio_24k = np.pad(audio_24k, (0, target_samples - len(audio_24k)))
elif len(audio_24k) > target_samples:
audio_24k = audio_24k[:target_samples]
return audio_24k
if __name__ == "__main__":
# Smoke test with real-ish audio (sine + noise)
import time
sr = SR_TARGET
t = np.arange(EXPECTED_SAMPLES) / sr
audio = (
0.3 * np.sin(2 * np.pi * 440 * t) +
0.2 * np.sin(2 * np.pi * 880 * t) +
0.05 * np.random.randn(EXPECTED_SAMPLES)
).astype(np.float32)
audio = np.clip(audio, -1, 1)
t0 = time.time()
enc = encode_opus(audio)
enc_ms = (time.time() - t0) * 1000
t0 = time.time()
dec = decode_opus(enc)
dec_ms = (time.time() - t0) * 1000
print(f"Input: {audio.shape} {audio.dtype} max={np.abs(audio).max():.3f}")
print(f"Encoded: {len(enc)} bytes (encode took {enc_ms:.1f} ms)")
print(f"Decoded: {dec.shape} {dec.dtype} (decode took {dec_ms:.1f} ms)")
print(f"Compression vs float32: {audio.nbytes / len(enc):.1f}x")
diff = audio - dec
print(f"Time-domain max abs diff: {np.abs(diff).max():.4f}")
print(f" L1 mean diff: {np.abs(diff).mean():.4f}")
# Spectral check
import scipy.signal as sps_
f, Pxx_in = sps_.welch(audio, fs=sr, nperseg=2048)
f, Pxx_out = sps_.welch(dec, fs=sr, nperseg=2048)
spec_db = 10 * np.log10(np.mean((np.log10(Pxx_in + 1e-10) - np.log10(Pxx_out + 1e-10))**2))
print(f" Spectral log-power MSE (dB): {spec_db:.2f}")
|