| |
| |
| """Trace pyannote's own front-end batching (`Inference.slide`, `SpeakerDiarization.get_embeddings`) with both |
| networks replaced by recorders, on CPU, and print the batches as JSON. Used by tests/test_front.py. |
| |
| The waveform carries its own sample index (sample k has value k + 1, exact in float32 below 2^24), so each |
| recorded window row tells which samples it holds; zero padding stays 0. |
| |
| python -s tests/ref_front_trace.py LENGTHS_JSON MASK_SEED OUT_JSON |
| """ |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import refpkg |
|
|
| refpkg.load() |
| import torch |
| from pyannote.core import SlidingWindowFeature |
| from anuvani_diar_ref.frontend import load_bundle |
|
|
| REL = Path(__file__).resolve().parents[2] / "anuvani_diar_v1" |
|
|
|
|
| def describe(rows): |
| """For each (1, WIN) row: first sample index, number of non-padding samples, padding all zero.""" |
| out = [] |
| for r in rows: |
| r = r.reshape(-1) |
| nz = np.flatnonzero(r) |
| valid = int(nz[-1]) + 1 if len(nz) else 0 |
| first = int(r[0]) - 1 if valid else -1 |
| ok = bool(valid == 0 or np.array_equal(r[:valid], np.arange(first + 1, first + 1 + valid, dtype=np.float32))) |
| out.append([first, valid, ok and not r[valid:].any()]) |
| return out |
|
|
|
|
| def main(): |
| lengths = json.loads(sys.argv[1]) |
| rng = np.random.default_rng(int(sys.argv[2])) |
| model = load_bundle(REL / "bundle", "cpu", 8, 2) |
| pipe = model["pipeline"] |
| |
| real_min = pipe._embedding.min_num_samples |
| seg_rec, emb_rec = [], [] |
|
|
| def fake_infer(chunks, soft=False): |
| x = chunks.numpy() |
| seg_rec.append(dict(B=int(x.shape[0]), rows=describe(x))) |
| return np.zeros((x.shape[0], 799, 16), np.float32) |
|
|
| class FakeEmbedding: |
| sample_rate = 16000 |
| min_num_samples = real_min |
| dimension = 256 |
| metric = "cosine" |
|
|
| def __call__(self, waveforms, masks=None): |
| emb_rec.append(dict(B=int(waveforms.shape[0]), rows=describe(waveforms.numpy()), |
| masks=masks.numpy().astype(np.uint8).tolist())) |
| return np.zeros((waveforms.shape[0], 256), np.float32) |
|
|
| pipe._segmentation.infer = fake_infer |
| pipe._embedding = FakeEmbedding() |
| result = dict(min_num_samples=int(real_min), receptive_field=[model["frames"].start, model["frames"].duration, |
| model["frames"].step], |
| mapping=model["mapping"].tolist(), runs=[]) |
| for n in lengths: |
| assert n < 2 ** 24 |
| audio = np.arange(1, n + 1, dtype=np.float32) |
| file = {"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000} |
| seg_rec.clear(); emb_rec.clear() |
| seg = pipe.get_segmentations(file, soft=False) |
| c = seg.data.shape[0] |
| density = rng.choice([0.05, 0.3, 0.6, 0.95]) |
| masks = (rng.random((c, 799, 4)) < density).astype(np.float32) |
| |
| if c > 1: |
| masks[1] = 0 |
| |
| masks[1, :2, 0] = 1 |
| |
| masks[1, 10:13, 1] = 1 |
| |
| masks[1, 13:20, 1:3] = 1 |
| sw = seg.sliding_window |
| emb = pipe.get_embeddings(file, SlidingWindowFeature(masks, sw), exclude_overlap=True) |
| result["runs"].append(dict(n=n, chunks=c, sw=[sw.start, sw.duration, sw.step], seg=list(seg_rec), |
| emb=list(emb_rec), masks=masks.astype(np.uint8).tolist(), |
| emb_shape=list(emb.shape))) |
| Path(sys.argv[3]).write_text(json.dumps(result)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
