| """AMI real-recording dataset — genuine mixtures with genuine ground truth.
|
|
|
| Each AMI meeting is recorded twice, simultaneously:
|
| Headset-N : close-talking mic on speaker N -> near-clean target voice
|
| Array1-01 : mic in the room -> the real mixture of everyone
|
|
|
| So (Array1-01, Headset-N) is a real (mixture, target) pair: same room, same
|
| moment, same acoustics, conversational speech, real mic chain. That is exactly
|
| the physics our synthetic pipeline can't fake.
|
|
|
| Three problems real data brings that synthesis never did, and how we handle them:
|
|
|
| 1. LEVEL/CHANNEL MISMATCH. The headset signal is not "the target as the room mic
|
| heard it" — different mic, distance, and gain. Training SI-SDR against the raw
|
| headset would ask the model to invent a channel transform. Fix: project the
|
| headset signal onto the room-mic signal (least-squares scalar) so the target
|
| is expressed *in the mixture's own channel*. SI-SDR is scale-invariant, but
|
| this keeps target and mixture in the same acoustic space.
|
|
|
| 2. SILENCE. Meetings are mostly one person talking; a random 3 s window often has
|
| the target silent, which is a degenerate SI-SDR target. Fix: index only
|
| windows where the target is genuinely active (energy threshold), while still
|
| allowing partial activity so the model keeps learning "go quiet when I stop".
|
|
|
| 3. TARGET LEAKAGE INTO OTHER HEADSETS. Not our problem here — we only use the
|
| target's own headset — but it's why we don't build interference from the
|
| other headsets.
|
|
|
| Enrollment comes from a *different* window of the same speaker's headset, so it
|
| matches deployment: a clean-ish reference clip of the person you want.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import json
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| import soundfile as sf
|
| import torch
|
| from torch.utils.data import Dataset, get_worker_info
|
|
|
| from vanta.utils.audio import EPS, peak_normalize
|
|
|
| ARRAY_MIC = "Array1-01"
|
|
|
|
|
| def _rms(x: np.ndarray) -> float:
|
| return float(np.sqrt(np.mean(x**2) + EPS))
|
|
|
|
|
| def build_index(
|
| ami_root: Path,
|
| clip_seconds: float = 3.0,
|
| hop_seconds: float = 2.0,
|
| sr: int = 16000,
|
| min_active_rms: float = 0.005,
|
| min_active_frac: float = 0.4,
|
| ) -> list[dict]:
|
| """Scan meetings and index windows where a speaker is genuinely talking.
|
|
|
| Returns entries: {meeting, speaker, start_sample, n_samples}. Reads the audio
|
| once per file with soundfile blocks, so memory stays flat.
|
| """
|
| entries: list[dict] = []
|
| n = int(clip_seconds * sr)
|
| hop = int(hop_seconds * sr)
|
|
|
| for mdir in sorted(p for p in ami_root.iterdir() if p.is_dir()):
|
| meeting = mdir.name
|
| array_path = mdir / f"{meeting}.{ARRAY_MIC}.wav"
|
| if not array_path.exists():
|
| continue
|
| for spk in range(4):
|
| hs_path = mdir / f"{meeting}.Headset-{spk}.wav"
|
| if not hs_path.exists():
|
| continue
|
| try:
|
| info = sf.info(str(hs_path))
|
| hs, file_sr = sf.read(str(hs_path), dtype="float32", always_2d=False)
|
| except Exception:
|
| continue
|
| if hs.ndim > 1:
|
| hs = hs.mean(axis=1)
|
| if file_sr != sr:
|
| import soxr
|
|
|
| hs = soxr.resample(hs, file_sr, sr, quality="QQ").astype(np.float32)
|
|
|
|
|
| frame = int(0.05 * sr)
|
| n_frames = len(hs) // frame
|
| if n_frames == 0:
|
| continue
|
| fr = hs[: n_frames * frame].reshape(n_frames, frame)
|
| active = (np.sqrt((fr**2).mean(axis=1) + EPS) > min_active_rms)
|
|
|
| frames_per_win = max(int(clip_seconds * sr) // frame, 1)
|
| for start_f in range(0, n_frames - frames_per_win, max(hop // frame, 1)):
|
| win = active[start_f : start_f + frames_per_win]
|
| if win.mean() >= min_active_frac:
|
| entries.append(
|
| {
|
| "meeting": meeting,
|
| "speaker": spk,
|
| "start": int(start_f * frame),
|
| "n": n,
|
| }
|
| )
|
| return entries
|
|
|
|
|
| class AmiDataset(Dataset):
|
| """Real (mixture, target, enrollment) triples from AMI meetings."""
|
|
|
| def __init__(
|
| self,
|
| ami_root: str | Path,
|
| index_path: str | Path,
|
| clip_seconds: float = 3.0,
|
| enroll_seconds: float = 5.0,
|
| sr: int = 16000,
|
| seed: int = 0,
|
| ):
|
| self.root = Path(ami_root)
|
| self.sr = sr
|
| self.clip_n = int(clip_seconds * sr)
|
| self.enroll_n = int(enroll_seconds * sr)
|
| self.seed = seed
|
| with open(index_path) as f:
|
| self.entries = [json.loads(line) for line in f]
|
|
|
|
|
| self.by_spk: dict[tuple[str, int], list[int]] = {}
|
| for i, e in enumerate(self.entries):
|
| self.by_spk.setdefault((e["meeting"], e["speaker"]), []).append(i)
|
| self._rng: np.random.Generator | None = None
|
|
|
| def __len__(self) -> int:
|
| return len(self.entries)
|
|
|
| def _get_rng(self) -> np.random.Generator:
|
| if self._rng is None:
|
| info = get_worker_info()
|
| seed = self.seed if info is None else (torch.initial_seed() % (2**31))
|
| self._rng = np.random.default_rng(seed)
|
| return self._rng
|
|
|
| def _read(self, meeting: str, mic: str, start: int, n: int) -> np.ndarray:
|
| path = self.root / meeting / f"{meeting}.{mic}.wav"
|
| with sf.SoundFile(str(path)) as f:
|
| file_sr = f.samplerate
|
| if file_sr != self.sr:
|
| ratio = file_sr / self.sr
|
| f.seek(int(start * ratio))
|
| raw = f.read(int(n * ratio), dtype="float32", always_2d=False)
|
| if raw.ndim > 1:
|
| raw = raw.mean(axis=1)
|
| import soxr
|
|
|
| w = soxr.resample(raw, file_sr, self.sr, quality="HQ").astype(np.float32)
|
| else:
|
| f.seek(start)
|
| w = f.read(n, dtype="float32", always_2d=False)
|
| if w.ndim > 1:
|
| w = w.mean(axis=1)
|
| if len(w) < n:
|
| w = np.pad(w, (0, n - len(w)))
|
| return w[:n].astype(np.float32)
|
|
|
| def __getitem__(self, idx: int) -> dict:
|
| rng = self._get_rng()
|
| e = self.entries[idx]
|
| meeting, spk = e["meeting"], e["speaker"]
|
|
|
| mixture = self._read(meeting, ARRAY_MIC, e["start"], self.clip_n)
|
| headset = self._read(meeting, f"Headset-{spk}", e["start"], self.clip_n)
|
|
|
|
|
|
|
|
|
| denom = float(np.dot(headset, headset)) + EPS
|
| alpha = float(np.dot(mixture, headset)) / denom
|
| target = (alpha * headset).astype(np.float32)
|
|
|
|
|
| pool = self.by_spk[(meeting, spk)]
|
| j = int(rng.integers(0, len(pool)))
|
| if len(pool) > 1 and pool[j] == idx:
|
| j = (j + 1) % len(pool)
|
| enr_entry = self.entries[pool[j]]
|
| enrollment = self._read(
|
| meeting, f"Headset-{spk}", enr_entry["start"], self.enroll_n
|
| )
|
| enrollment = peak_normalize(enrollment, peak=0.95)
|
|
|
|
|
|
|
| peak = float(np.max(np.abs(mixture))) + EPS
|
| if peak > 0.95:
|
| s = 0.95 / peak
|
| mixture = mixture * s
|
| target = target * s
|
|
|
| return {
|
| "mixture": torch.from_numpy(mixture.astype(np.float32)),
|
| "target": torch.from_numpy(target.astype(np.float32)),
|
| "enrollment": torch.from_numpy(enrollment.astype(np.float32)),
|
| "target_speaker": f"{meeting}-{spk}",
|
| }
|
|
|
