Files
beefy-player/somafm_tui/spectrum.py
T

535 lines
18 KiB
Python

"""Spectrum analyser.
The audio you hear is played by mpv. The *synced* mode captures the real
PulseAudio/PipeWire monitor of the sink mpv is playing to (``parec``), so the
bars follow exactly what leaves the speakers (status ``SYNC``).
If that is unavailable we fall back to a muted second ffmpeg decode of the
stream (status ``FFT``; it can drift relative to playback), and finally to a
synthesized show (``SIM``) so the UI stays alive.
"""
from __future__ import annotations
import contextlib
import math
import os
import select
import shutil
import subprocess
import threading
import time
from dataclasses import dataclass
import numpy as np
from .api import USER_AGENT
SAMPLE_RATE = 44100
FFT_SIZE = 2048
READ_BYTES = 4096 # 2048 mono s16le samples
BAND_COUNT = 72
WAVE_POINTS = 192
MIN_FREQ = 35.0
MAX_FREQ = 16000.0
SOURCE_MODES = ("auto", "monitor", "stream", "sim")
@dataclass(frozen=True)
class Frame:
bands: tuple[float, ...]
peaks: tuple[float, ...]
wave: tuple[float, ...]
level: float
status: str # "monitor" | "live" | "sim" | "idle" | "off"
@property
def is_live(self) -> bool:
return self.status in ("monitor", "live")
@property
def label(self) -> str:
return {
"monitor": "SYNC",
"live": "FFT",
"sim": "SIM",
"idle": "IDLE",
"off": "OFF",
}.get(self.status, self.status.upper())
class SpectrumEngine:
def __init__(
self,
bands: int = BAND_COUNT,
sample_rate: int = SAMPLE_RATE,
fft_size: int = FFT_SIZE,
wave_points: int = WAVE_POINTS,
ffmpeg_path: str | None = None,
source_mode: str = "auto",
) -> None:
self.band_count = bands
self.sample_rate = sample_rate
self.fft_size = fft_size
self.wave_points = wave_points
self.source_mode = source_mode if source_mode in SOURCE_MODES else "auto"
self._ffmpeg = ffmpeg_path if ffmpeg_path is not None else shutil.which("ffmpeg")
self._parec = shutil.which("parec")
self._pactl = shutil.which("pactl")
self._window = np.hanning(fft_size).astype(np.float32)
self._groups = self._make_groups()
self._history = np.zeros(fft_size, dtype=np.float32)
self._levels = np.zeros(bands, dtype=np.float32)
self._peaks = np.zeros(bands, dtype=np.float32)
self._wave = np.zeros(wave_points, dtype=np.float32)
self._level = 0.0
self._gain = -24.0
self._status = "idle"
self._url: str | None = None
self._proc: subprocess.Popen[bytes] | None = None
self._capture_kind: str | None = None
self._monitor_failed = False
self._monitor_device_cache: str | None = None
self._lock = threading.Lock()
self._stop = threading.Event()
self._thread: threading.Thread | None = None
self._rng = np.random.default_rng()
self._sim_phase = 0.0
# ------------------------------------------------------------------ #
# public API
# ------------------------------------------------------------------ #
@property
def ffmpeg_available(self) -> bool:
return bool(self._ffmpeg)
@property
def monitor_available(self) -> bool:
return bool(self._parec) and not self._monitor_failed
def start(self) -> None:
if self._thread is not None:
return
self._thread = threading.Thread(
target=self._run, name="somafm-spectrum", daemon=True
)
self._thread.start()
def set_source(self, url: str | None) -> None:
"""Start analysing playback, or idle when ``url`` is ``None``.
The URL is only a fallback source for the stream tap. A monitor
capture is tied to the actual audio output, so it deliberately keeps
running when the station changes underneath it.
"""
url = url or None
with self._lock:
changed = url != self._url
self._url = url
kind = self._capture_kind
if url is None:
self._status = "idle"
if changed and (url is None or kind != "monitor"):
self._kill_proc()
def snapshot(self) -> Frame:
with self._lock:
return Frame(
bands=tuple(float(v) for v in self._levels),
peaks=tuple(float(v) for v in self._peaks),
wave=tuple(float(v) for v in self._wave),
level=float(self._level),
status=self._status,
)
def stop(self) -> None:
self._stop.set()
self._kill_proc()
thread = self._thread
if thread is not None and thread.is_alive():
thread.join(timeout=2.0)
self._thread = None
with self._lock:
self._status = "off"
# ------------------------------------------------------------------ #
# internals
# ------------------------------------------------------------------ #
def _make_groups(self) -> list[tuple[int, int]]:
freqs = np.fft.rfftfreq(self.fft_size, 1.0 / self.sample_rate)
top = min(MAX_FREQ, self.sample_rate / 2 - 100)
edges = np.geomspace(MIN_FREQ, top, self.band_count + 1)
groups: list[tuple[int, int]] = []
last = len(freqs) - 1
for index in range(self.band_count):
lo = int(np.searchsorted(freqs, edges[index], side="left"))
hi = int(np.searchsorted(freqs, edges[index + 1], side="left"))
lo = min(max(lo, 0), last)
hi = min(max(hi, lo + 1), len(freqs))
groups.append((lo, hi))
return groups
def _read_url(self) -> str | None:
with self._lock:
return self._url
def _kill_proc(self) -> None:
with self._lock:
proc = self._proc
self._proc = None
self._capture_kind = None
if proc is not None:
with contextlib.suppress(OSError):
proc.kill()
with contextlib.suppress(Exception):
proc.wait(timeout=1)
def _use_monitor(self) -> bool:
if self._monitor_failed:
return False
if self.source_mode == "sim":
return False
if self.source_mode == "stream":
return False
return bool(self._parec)
def _monitor_device(self) -> str:
override = os.environ.get("SOMAFM_SPECTRUM_DEVICE", "").strip()
if override:
return override
if self._monitor_device_cache is not None:
return self._monitor_device_cache
if self._pactl:
try:
result = subprocess.run(
[self._pactl, "get-default-sink"],
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
text=True,
timeout=1.0,
check=False,
)
sink = result.stdout.strip()
if result.returncode == 0 and sink:
self._monitor_device_cache = f"{sink}.monitor"
return self._monitor_device_cache
except (OSError, subprocess.SubprocessError):
pass
self._monitor_device_cache = "@DEFAULT_MONITOR@"
return self._monitor_device_cache
def _monitor_command(self) -> list[str] | None:
if not self._parec:
return None
return [
self._parec,
f"--device={self._monitor_device()}",
"--format=s16le",
"--rate=44100",
"--channels=1",
"--raw",
"--latency-msec=30",
"--client-name=somafm-tui",
]
def _spawn(
self, command: list[str], kind: str
) -> subprocess.Popen[bytes] | None:
try:
proc = subprocess.Popen(
command,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
bufsize=0,
start_new_session=True,
)
except (OSError, ValueError):
return None
with self._lock:
self._proc = proc
self._capture_kind = kind
self._status = "monitor" if kind == "monitor" else "live"
return proc
def _stop_proc(self, proc: subprocess.Popen[bytes]) -> None:
with self._lock:
if self._proc is proc:
self._proc = None
self._capture_kind = None
with contextlib.suppress(OSError):
proc.kill()
with contextlib.suppress(Exception):
proc.wait(timeout=1)
def _run(self) -> None:
while not self._stop.is_set():
url = self._read_url()
if not url:
self._idle_frame()
self._stop.wait(0.05)
continue
if self.source_mode == "sim":
self._simulate_for(url, 0.05)
continue
if self._use_monitor():
command = self._monitor_command()
proc = self._spawn(command, "monitor") if command else None
if proc is not None:
self._capture(
proc, url, stop_on_url_change=False, status="monitor"
)
self._stop_proc(proc)
if self._stop.is_set():
break
if self._read_url() is None:
continue
# The monitor device failed (or vanished). Use the stream
# tap for the rest of the session.
self._monitor_failed = True
with self._lock:
self._status = "sim"
continue
self._monitor_failed = True
if not self._ffmpeg:
self._simulate_for(url, 0.05)
continue
proc = self._spawn(self._command(url), "stream")
if proc is None:
with self._lock:
self._status = "sim"
self._simulate_for(url, 0.06, duration=5.0)
continue
self._capture(proc, url, stop_on_url_change=True, status="live")
self._stop_proc(proc)
if self._stop.is_set():
break
if self._read_url() == url:
with self._lock:
self._status = "sim"
self._simulate_for(url, 0.06, duration=12.0)
def _command(self, url: str) -> list[str]:
return [
self._ffmpeg or "ffmpeg",
"-hide_banner",
"-loglevel",
"error",
"-nostdin",
"-reconnect",
"1",
"-reconnect_streamed",
"1",
"-reconnect_at_eof",
"1",
"-reconnect_delay_max",
"5",
"-fflags",
"nobuffer",
"-flags",
"low_delay",
"-probesize",
"32768",
"-analyzeduration",
"0",
"-user_agent",
USER_AGENT,
"-i",
url,
"-map",
"0:a:0",
"-vn",
"-ac",
"1",
"-ar",
str(self.sample_rate),
"-f",
"s16le",
"-acodec",
"pcm_s16le",
"-",
]
def _capture(
self,
proc: subprocess.Popen[bytes],
url: str,
*,
stop_on_url_change: bool,
status: str,
) -> bool:
"""Read PCM until playback stops, the source changes, or the process dies."""
stdout = proc.stdout
if stdout is None:
return False
fd = stdout.fileno()
got_data = False
buffer = bytearray()
while not self._stop.is_set():
current = self._read_url()
if current is None:
break
if stop_on_url_change and current != url:
break
# A monitor can go quiet if the device disappears; don't block the
# analyser thread forever. os.read returns whatever is available
# (instead of waiting for a full 4 KiB), which keeps latency low.
if status == "monitor" and os.name == "posix":
try:
ready, _, _ = select.select([fd], [], [], 2.0)
except (OSError, ValueError):
ready = [fd]
if not ready:
break
try:
data = os.read(fd, READ_BYTES)
except OSError:
break
if not data:
break
buffer.extend(data)
# Process fixed-size blocks so attack/decay smoothing is tied to
# time (~46 ms per block), regardless of how the pipe fragments.
while len(buffer) >= READ_BYTES:
block = bytes(buffer[:READ_BYTES])
del buffer[:READ_BYTES]
got_data = True
self.feed_pcm(block, status)
return got_data
# ------------------------------------------------------------------ #
# signal processing
# ------------------------------------------------------------------ #
def feed_pcm(self, data: bytes, status: str = "live") -> None:
"""Feed raw little-endian mono s16 PCM (also used by the tests)."""
if not data or len(data) < 2:
return
usable = len(data) - (len(data) % 2)
try:
samples = np.frombuffer(data[:usable], dtype="<i2")
except ValueError:
return
if samples.size == 0:
return
self._process(samples.astype(np.float32) / 32768.0, status)
def _process(self, samples: np.ndarray, status: str = "live") -> None:
size = self.fft_size
if samples.size >= size:
self._history[:] = samples[-size:]
else:
count = samples.size
self._history[:-count] = self._history[count:]
self._history[-count:] = samples
chunk = self._history
spectrum = np.abs(np.fft.rfft(chunk * self._window))
spectrum /= size * 0.5
db = 20.0 * np.log10(spectrum + 1e-9)
values = np.array(
[db[lo:hi].max() for lo, hi in self._groups], dtype=np.float32
)
peak_db = float(values.max())
self._gain = max(peak_db, self._gain - 0.35)
floor = self._gain - 42.0
norm = np.clip((values - floor) / 42.0, 0.0, 1.0) ** 1.4
rms = float(np.sqrt(np.mean(np.square(chunk)) + 1e-12))
level = float(np.clip((20.0 * math.log10(rms + 1e-9) + 55.0) / 45.0, 0.0, 1.0))
wave = self._downsample(chunk)
self._apply(norm.astype(np.float32), wave, level, status)
def _downsample(self, samples: np.ndarray) -> np.ndarray:
target = self.wave_points
if samples.size < target:
padded = np.zeros(target, dtype=np.float32)
padded[-samples.size :] = samples
samples = padded
step = max(1, samples.size // target)
trimmed = samples[samples.size - step * target :]
try:
folded = trimmed.reshape(target, step).mean(axis=1)
except ValueError:
return np.zeros(target, dtype=np.float32)
peak = float(np.max(np.abs(folded)))
if peak > 1e-4:
folded = folded / peak
return np.clip(folded, -1.0, 1.0).astype(np.float32)
def _apply(
self, values: np.ndarray, wave: np.ndarray, level: float, status: str
) -> None:
with self._lock:
self._levels = np.maximum(values, self._levels * 0.65).astype(np.float32)
self._peaks = np.maximum(values, self._peaks - 0.04).astype(np.float32)
self._wave = (self._wave * 0.55 + wave * 0.45).astype(np.float32)
self._level = float(self._level * 0.5 + level * 0.5)
self._status = status
def _idle_frame(self) -> None:
with self._lock:
self._levels = (self._levels * 0.86).astype(np.float32)
self._peaks = np.maximum(self._peaks - 0.02, 0.0).astype(np.float32)
self._wave = (self._wave * 0.8).astype(np.float32)
self._level *= 0.9
self._status = "idle"
# ------------------------------------------------------------------ #
# simulation
# ------------------------------------------------------------------ #
def _simulate_for(
self, url: str, delay: float, duration: float | None = None
) -> None:
deadline = None if duration is None else time.monotonic() + duration
while not self._stop.is_set():
if self._read_url() != url:
return
self._simulate_step()
if self._stop.wait(delay):
return
if deadline is not None and time.monotonic() >= deadline:
return
def _simulate_step(self) -> None:
self._sim_phase += 0.16
t = self._sim_phase
index = np.linspace(0.0, 1.0, self.band_count, dtype=np.float32)
humps = np.exp(-((index - 0.16) ** 2) / 0.020) * 0.58
humps += np.exp(-((index - 0.46) ** 2) / 0.055) * 0.34
humps += np.exp(-((index - 0.80) ** 2) / 0.130) * 0.20
motion = 0.55 + 0.45 * np.sin(t * 1.9 + index * 11.0)
noise = self._rng.random(self.band_count, dtype=np.float32) * 0.22
beat_phase = (t * 0.55) % 1.0
beat = math.exp(-(beat_phase / 0.16)) * 0.28 * np.exp(-(index**2) / 0.03)
values = np.clip((humps * motion + noise + beat), 0.02, 1.0)
values = np.clip(values * (1.15 - 0.25 * index), 0.0, 1.0).astype(np.float32)
points = np.linspace(0.0, 2.0 * math.pi * 7.0, self.wave_points)
wave = 0.42 * np.sin(points + t * 3.0) * (0.6 + 0.4 * math.sin(t * 0.7))
wave += 0.20 * np.sin(points * 1.9 + t * 4.3)
wave += self._rng.normal(0.0, 0.05, self.wave_points)
wave = np.clip(wave, -1.0, 1.0).astype(np.float32)
level = float(
min(1.0, 0.32 + 0.14 * math.sin(t * 1.7) + 0.25 * math.exp(-(beat_phase / 0.2)))
)
self._apply(values, wave, level, "sim")