Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/preprocess/vap/03_prep_vad_rate.py

@@ -0,0 +1,112 @@
+"""Write the segment VAD again at a different frame rate.
+
+`00_prep_fisher.py` rasterizes the word annotations at `VAD_HZ` and stores one
+`.npy` per segment beside the audio. A projection window running at another rate
+needs its labels on that grid: resampling the 25 Hz array would inherit its
+quantisation, so this rasterizes the annotations again from their original
+continuous times.
+
+Only the VAD is rewritten. The audio is rate-independent and stays where it is,
+so this touches ~1 GB rather than the 1.12 TB of segments, and it writes to
+`seg/vad{rate}/` rather than over `seg/vad/`.
+
+    python preprocess/vap/03_prep_vad_rate.py --frame-hz 12.5 \
+        --fisher-path /path/to/fisher --splits train val test
+"""
+
+import argparse
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+
+from preprocess.vap.utils import SAMPLE_RATE
+
+SEG_SEC = 30
+CONTEXT_SEC = 2.0
+
+
+def rasterize(words_dir, conversation, frames, frame_hz):
+    """`[2, frames]` activity for a whole conversation, from the word times."""
+    vad = np.zeros((2, frames), dtype=np.float32)
+    paths = sorted(words_dir.glob(f"{conversation}_*_words.txt"))
+    if len(paths) != 2:
+        raise FileNotFoundError(
+            f"expected two *_words.txt for {conversation} in {words_dir}, found {len(paths)}"
+        )
+    for speaker, path in enumerate(paths):
+        with open(path) as handle:
+            for line in handle:
+                if not line.strip():
+                    continue
+                begin, end, *_ = line.split()
+                first = max(0, int(round(float(begin) * frame_hz)))
+                last = min(frames, int(round(float(end) * frame_hz)))
+                if last > first:
+                    vad[speaker, first:last] = 1.0
+    return vad
+
+
+def segment_window(vad, start_sec, frame_hz, context_frames):
+    """The stored slice for one segment: the segment plus context, zero-padded.
+
+    Mirrors `extract_and_save_chunks`, so a reader can treat these files exactly
+    as it treats the 25 Hz ones.
+    """
+    seg_frames = int(round(SEG_SEC * frame_hz))
+    start = int(round(start_sec * frame_hz))
+    end = start + seg_frames
+    lo, hi = max(0, start - context_frames), min(vad.shape[-1], end + context_frames)
+    chunk = vad[:, lo:hi]
+    pad_left = max(0, context_frames - start)
+    pad_right = max(0, (end + context_frames) - vad.shape[-1])
+    if pad_left or pad_right:
+        chunk = np.pad(chunk, ((0, 0), (pad_left, pad_right)))
+    return chunk
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--fisher-path", required=True)
+    parser.add_argument("--frame-hz", type=float, required=True)
+    parser.add_argument("--splits", nargs="+", default=["train", "val", "test"])
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    root = Path(args.fisher_path)
+    label = f"{args.frame_hz:g}".replace(".", "_")
+    context_frames = int(round(CONTEXT_SEC * args.frame_hz))
+
+    # One rasterization per conversation, then slice every segment out of it.
+    by_conversation = defaultdict(list)
+    for split in args.splits:
+        for line in open(root / f"{split}.txt"):
+            if not line.strip():
+                continue
+            part, group, conversation, seg_id, start, *_ = line.split()
+            by_conversation[(part, group, conversation)].append((int(seg_id), float(start)))
+
+    written = skipped = 0
+    for index, ((part, group, conversation), segments) in enumerate(sorted(by_conversation.items()), 1):
+        out = root / part / "seg" / f"vad{label}" / group / conversation
+        out.mkdir(parents=True, exist_ok=True)
+        targets = [(s, p) for s, p in ((s, out / f"{s}.npy") for s, _ in segments)
+                   if args.overwrite or not p.exists()]
+        if not targets:
+            skipped += len(segments)
+            continue
+        last_start = max(start for _, start in segments)
+        frames = int(round((last_start + SEG_SEC + CONTEXT_SEC) * args.frame_hz)) + 1
+        vad = rasterize(root / part / "regions_vap" / group, conversation, frames, args.frame_hz)
+        starts = dict(segments)
+        for seg_id, path in targets:
+            np.save(path, segment_window(vad, starts[seg_id], args.frame_hz, context_frames))
+            written += 1
+        if index % 500 == 0:
+            print(f"  {index}/{len(by_conversation)} conversations, {written} written", flush=True)
+    print(f"done: {written} written, {skipped} already present, "
+          f"{len(by_conversation)} conversations -> seg/vad{label}/")
+
+
+if __name__ == "__main__":
+    main()