535 lines
18 KiB
Python
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")
|