Mimir/scripts/capture_co_streamer_surfaces.py
2026-05-18 20:12:39 +01:00

316 lines
14 KiB
Python

from __future__ import annotations
import argparse
import json
import math
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path
from tempfile import NamedTemporaryFile
import numpy as np
from scipy import signal
from scipy.io import wavfile
DEFAULT_RUN = Path("calibration/runs/audio-program-live-20260518-180226")
def main() -> int:
parser = argparse.ArgumentParser(
description=(
"Capture the late-arriving co-streamer audio surfaces, align them to "
"local loopback ground truth, and write synced OBS stem inputs."
)
)
parser.add_argument("--run", type=Path, default=DEFAULT_RUN)
parser.add_argument("--seconds", type=float, default=24.0)
parser.add_argument("--sample-rate", type=int, default=48000)
parser.add_argument("--ssh-target", default="madman's lullaby@192.168.1.84")
parser.add_argument("--ffmpeg", default=r"C:\Users\Madman's Lullaby\AppData\Local\Microsoft\WinGet\Links\ffmpeg.exe")
parser.add_argument("--remote-dir", default=r"C:\Meta\LocalCastBridge\calibration\remote-captures")
parser.add_argument("--focusrite-device", default="Analogue 1 + 2 (Focusrite USB Audio)")
parser.add_argument("--loopback-device", default="")
parser.add_argument("--loopback-capture", choices=("wasapi", "dshow"), default="wasapi")
parser.add_argument("--remote-wasapi-script", default=r"C:\Meta\LocalCastBridge\scripts\wasapi-loopback-capture.ps1")
parser.add_argument("--remote-wasapi-role", choices=("Console", "Multimedia", "Communications"), default="Console")
parser.add_argument("--local-loopback-query", default="Scarlett")
parser.add_argument("--max-lag-ms", type=float, default=3000.0)
parser.add_argument("--min-loopback-rms", type=float, default=1e-4)
parser.add_argument("--min-loopback-score", type=float, default=0.01)
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
run = args.run.resolve()
surface_dir = run / "co_streamer_surfaces"
surface_dir.mkdir(parents=True, exist_ok=True)
stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
remote_focusrite = f"{args.remote_dir.rstrip('\\\\/')}\\co_streamer_focusrite-{stamp}.wav"
remote_loopback_ext = "f32" if args.loopback_capture == "wasapi" else "wav"
remote_loopback = f"{args.remote_dir.rstrip('\\\\/')}\\co_streamer_loopback-{stamp}.{remote_loopback_ext}"
local_focusrite = surface_dir / "raw_focusrite.wav"
local_remote_loopback = surface_dir / f"raw_loopback.{remote_loopback_ext}"
local_reference = surface_dir / "local_loopback_reference.wav"
commands = [
remote_ffmpeg_command(args, args.focusrite_device, remote_focusrite, channels=1),
remote_loopback_command(args, remote_loopback),
]
if args.dry_run:
for command in commands:
print(command)
return 0
print("starting remote co-streamer captures")
procs = [subprocess.Popen(["ssh", args.ssh_target, command], text=True) for command in commands]
time.sleep(0.2)
local_loopback = record_local_loopback(
query=args.local_loopback_query,
seconds=float(args.seconds),
sample_rate=int(args.sample_rate),
channels=2,
)
wavfile.write(local_reference, int(args.sample_rate), local_loopback.astype(np.float32))
failures = []
for proc in procs:
rc = proc.wait()
if rc:
failures.append(rc)
if failures:
raise SystemExit(f"remote capture failed with exit codes {failures}")
sftp_get(args.ssh_target, remote_focusrite, local_focusrite)
sftp_get(args.ssh_target, remote_loopback, local_remote_loopback)
sample_rate = int(args.sample_rate)
_, reference = read_wav_float(local_reference, sample_rate)
_, focusrite = read_wav_float(local_focusrite, sample_rate)
remote_loop = read_loopback_capture(local_remote_loopback, sample_rate, args.loopback_capture)
reference_mono = mono(reference)
remote_loop_mono = mono(remote_loop)
max_lag = int(float(args.max_lag_ms) * sample_rate / 1000.0)
remote_loop_rms = rms(remote_loop_mono)
focusrite_mono = mono(focusrite)
focusrite_rms = rms(focusrite_mono)
loopback_delay, loopback_score = normalized_delay(remote_loop_mono, reference_mono, max_lag=max_lag)
focusrite_delay, focusrite_score = normalized_delay(focusrite_mono, reference_mono, max_lag=max_lag)
if remote_loop_rms >= float(args.min_loopback_rms) and loopback_score >= float(args.min_loopback_score):
delay_samples = loopback_delay
score = loopback_score
alignment_witness = "remote-loopback"
else:
delay_samples = focusrite_delay
score = focusrite_score
alignment_witness = "neighbor-focusrite-program-bleed"
aligned_focusrite = shift_by_delay(focusrite, delay_samples, len(reference))
aligned_loopback = shift_by_delay(remote_loop, delay_samples, len(reference))
wavfile.write(surface_dir / "aligned_focusrite.wav", sample_rate, aligned_focusrite.astype(np.float32))
wavfile.write(surface_dir / "aligned_loopback.wav", sample_rate, aligned_loopback.astype(np.float32))
report = {
"schema_version": "gamecult.localcast.co_streamer_surfaces.v1",
"created_utc": datetime.now(timezone.utc).isoformat(),
"sample_rate": sample_rate,
"seconds": float(args.seconds),
"alignment_witness": alignment_witness,
"delay_samples": int(delay_samples),
"delay_ms": 1000.0 * float(delay_samples) / sample_rate,
"correlation_score": float(score),
"remote_loopback": {
"rms": remote_loop_rms,
"delay_samples": int(loopback_delay),
"delay_ms": 1000.0 * float(loopback_delay) / sample_rate,
"correlation_score": float(loopback_score),
},
"neighbor_focusrite": {
"rms": focusrite_rms,
"delay_samples": int(focusrite_delay),
"delay_ms": 1000.0 * float(focusrite_delay) / sample_rate,
"correlation_score": float(focusrite_score),
},
"interpretation": (
"positive delay means the remote loopback arrived later than local loopback "
"and was advanced into the local presentation timeline"
),
"files": {
"local_loopback_reference": str(local_reference),
"raw_focusrite": str(local_focusrite),
"raw_loopback": str(local_remote_loopback),
"aligned_focusrite": str(surface_dir / "aligned_focusrite.wav"),
"aligned_loopback": str(surface_dir / "aligned_loopback.wav"),
},
}
(surface_dir / "alignment-report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print(json.dumps(report, indent=2))
return 0
def remote_ffmpeg_command(args: argparse.Namespace, device: str, output: str, *, channels: int) -> str:
remote_dir = args.remote_dir.rstrip("\\/")
return (
f'cmd /c if not exist "{remote_dir}" mkdir "{remote_dir}" && '
f'"{args.ffmpeg}" -y -hide_banner -nostdin -loglevel warning '
f'-f dshow -t {float(args.seconds):.3f} -i audio="{device}" '
f'-vn -ac {channels} -ar {int(args.sample_rate)} -c:a pcm_f32le "{output}"'
)
def remote_loopback_command(args: argparse.Namespace, output: str) -> str:
if args.loopback_capture == "dshow":
if not args.loopback_device:
raise SystemExit("--loopback-device is required when --loopback-capture dshow")
return remote_ffmpeg_command(args, args.loopback_device, output, channels=2)
log = output.rsplit(".", 1)[0] + ".log"
task = "LocalCastWasapiLoopbackCapture"
return (
"powershell -NoProfile -ExecutionPolicy Bypass -Command "
+ quote_for_cmd(
"$ErrorActionPreference='Stop'; "
f"$out='{escape_ps(output)}'; $log='{escape_ps(log)}'; "
"Remove-Item -LiteralPath $out,$log -ErrorAction SilentlyContinue; "
f"$cmd='powershell.exe -NoProfile -ExecutionPolicy Bypass -File \"{escape_ps(args.remote_wasapi_script)}\" "
f"-Output \"{escape_ps(output)}\" -Seconds {float(args.seconds):.3f} "
f"-SampleRate {int(args.sample_rate)} -Channels 2 -Role {args.remote_wasapi_role} *> \"{escape_ps(log)}\"'; "
"$bat=[IO.Path]::ChangeExtension($out,'.cmd'); "
"Set-Content -LiteralPath $bat -Encoding ASCII -Value ('@echo off' + [Environment]::NewLine + $cmd + [Environment]::NewLine); "
"$action=New-ScheduledTaskAction -Execute $bat; "
"$trigger=New-ScheduledTaskTrigger -Once -At (Get-Date).AddMinutes(10); "
"$principal=New-ScheduledTaskPrincipal -UserId ($env:COMPUTERNAME + '\\' + $env:USERNAME) -LogonType Interactive -RunLevel Limited; "
f"Register-ScheduledTask -TaskName '{task}' -Action $action -Trigger $trigger -Principal $principal -Force | Out-Null; "
f"Start-ScheduledTask -TaskName '{task}'; "
f"$deadline=(Get-Date).AddSeconds({float(args.seconds) + 12.0:.3f}); "
"while((Get-Date) -lt $deadline){ "
" $info=Get-ScheduledTaskInfo -TaskName '" + task + "'; "
" if((Test-Path -LiteralPath $out) -and ((Get-Item -LiteralPath $out).Length -gt 0) -and $info.LastTaskResult -eq 0){ break } "
" Start-Sleep -Milliseconds 250 "
"}; "
"if(-not (Test-Path -LiteralPath $out)){ if(Test-Path -LiteralPath $log){Get-Content $log -Tail 40}; throw 'WASAPI loopback capture did not create output' }; "
"Write-Host ('wasapiLoopbackBytes=' + (Get-Item -LiteralPath $out).Length)"
)
)
def quote_for_cmd(command: str) -> str:
return '"' + command.replace('"', '\\"') + '"'
def escape_ps(value: str) -> str:
return value.replace("'", "''")
def sftp_get(ssh_target: str, remote_file: str, local_file: Path) -> None:
remote_sftp_path = "/C:/" + remote_file.replace("\\", "/").split("C:/", 1)[-1]
batch = f'get "{remote_sftp_path}" "{local_file}"\n'
with NamedTemporaryFile("w", encoding="ascii", delete=False, suffix=".sftp") as fh:
fh.write(batch)
batch_path = fh.name
try:
completed = subprocess.run(["sftp", "-b", batch_path, ssh_target], text=True, capture_output=True)
if completed.returncode:
raise SystemExit(completed.stderr or completed.stdout)
finally:
Path(batch_path).unlink(missing_ok=True)
def record_local_loopback(query: str, seconds: float, sample_rate: int, channels: int) -> np.ndarray:
try:
import soundcard as sc
except Exception as exc:
raise SystemExit(f"soundcard loopback capture unavailable: {exc!r}") from exc
matches = [
mic
for mic in sc.all_microphones(include_loopback=True)
if "loopback" in repr(mic).lower() and query.lower() in mic.name.lower()
]
if not matches:
raise SystemExit(f"No loopback device matching {query!r}")
with matches[0].recorder(samplerate=sample_rate, channels=channels) as recorder:
return recorder.record(numframes=int(seconds * sample_rate)).astype(np.float32)
def read_wav_float(path: Path, target_rate: int) -> tuple[int, np.ndarray]:
rate, samples = wavfile.read(path)
samples = pcm_to_float(samples)
if int(rate) != int(target_rate):
gcd = math.gcd(int(rate), int(target_rate))
samples = signal.resample_poly(samples, int(target_rate) // gcd, int(rate) // gcd, axis=0).astype(np.float32)
rate = target_rate
return int(rate), ensure_2d(samples)
def read_loopback_capture(path: Path, target_rate: int, capture_kind: str) -> np.ndarray:
if capture_kind == "wasapi":
samples = np.fromfile(path, dtype=np.float32)
if samples.size % 2:
samples = samples[:-1]
return samples.reshape(-1, 2).astype(np.float32)
_, samples = read_wav_float(path, target_rate)
return samples
def pcm_to_float(samples: np.ndarray) -> np.ndarray:
array = np.asarray(samples)
if array.ndim == 1:
array = array[:, None]
if np.issubdtype(array.dtype, np.floating):
return array.astype(np.float32, copy=False)
if np.issubdtype(array.dtype, np.signedinteger):
scale = float(max(abs(np.iinfo(array.dtype).min), np.iinfo(array.dtype).max))
return (array.astype(np.float32) / scale).clip(-1.0, 1.0)
raise TypeError(f"unsupported WAV dtype: {array.dtype}")
def ensure_2d(samples: np.ndarray) -> np.ndarray:
return samples[:, None] if samples.ndim == 1 else samples
def mono(samples: np.ndarray) -> np.ndarray:
samples = ensure_2d(samples)
return np.mean(samples, axis=1, dtype=np.float32)
def normalized_delay(signal_a: np.ndarray, signal_b: np.ndarray, max_lag: int) -> tuple[int, float]:
count = min(len(signal_a), len(signal_b))
if count <= 0:
return 0, 0.0
a = signal_a[:count].astype(np.float32)
b = signal_b[:count].astype(np.float32)
if float(np.std(a)) <= 1e-12 or float(np.std(b)) <= 1e-12:
return 0, 0.0
a = (a - float(np.mean(a))) / (float(np.std(a)) + 1e-9)
b = (b - float(np.mean(b))) / (float(np.std(b)) + 1e-9)
corr = signal.correlate(a, b, mode="full", method="fft")
mid = len(b) - 1
segment = corr[mid - max_lag : mid + max_lag + 1]
index = int(np.argmax(np.abs(segment)))
lag = index - max_lag
score = float(abs(segment[index]) / max(1, len(a)))
return lag, score
def rms(samples: np.ndarray) -> float:
samples = np.asarray(samples, dtype=np.float32)
return float(np.sqrt(np.mean(samples * samples))) if samples.size else 0.0
def shift_by_delay(samples: np.ndarray, delay_samples: int, frame_count: int) -> np.ndarray:
samples = ensure_2d(samples)
out = np.zeros((frame_count, samples.shape[1]), dtype=np.float32)
if delay_samples >= 0:
src_start = delay_samples
dst_start = 0
else:
src_start = 0
dst_start = -delay_samples
count = min(frame_count - dst_start, samples.shape[0] - src_start)
if count > 0:
out[dst_start : dst_start + count] = samples[src_start : src_start + count]
return out
if __name__ == "__main__":
raise SystemExit(main())