Add causal person fusion and MuVAP stage-3 baselines

2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions
@@ -0,0 +1,196 @@
+"""Cut Fisher into fixed segments of stereo audio beside their raw VAD.
+
+Class labels are *not* baked in here. Each segment stores the two-speaker VAD
+padded with `--context-sec` of extra frames on both sides, which is everything
+`ProjectionWindow` needs to produce role-based or speaker-based labels in the
+dataloader. Switching projection setups therefore never requires re-running
+this script.
+"""
+
+import argparse
+import os
+from dataclasses import dataclass
+from glob import glob
+
+import numpy as np
+from tqdm import tqdm
+
+from config.setup import init_yaml_config
+from data.utils import load_audio, load_fisher_vad
+from preprocess.vap.utils import (
+    SAMPLE_RATE,
+    TEST_GROUPS,
+    VAL_GROUPS,
+    ProcessConfig,
+    get_dataset_split,
+)
+
+SEG_SEC: int = 30
+OVERLAP: int = 10
+VAD_HZ: int = 25
+CONTEXT_SEC: float = 2.0
+
+STRIDE_SEC: int = SEG_SEC - OVERLAP
+SEG_SAMPLES: int = int(SEG_SEC * SAMPLE_RATE)
+STRIDE_SAMPLES: int = int(STRIDE_SEC * SAMPLE_RATE)
+SEG_VAD_FRAMES: int = int(SEG_SEC * VAD_HZ)
+
+
+@dataclass()
+class SegmentMeta:
+    part: str
+    group: str
+    cid: str
+    seg_id: int
+    start_t: float
+    end_t: float
+    nframes: int
+
+
+def format_meta(s: SegmentMeta) -> str:
+    return f"{s.part} {s.group} {s.cid} {s.seg_id} {s.start_t:.2f} {s.end_t:.2f} {s.nframes}\n"
+
+
+def extract_and_save_chunks(
+    audio: np.ndarray,
+    vad: np.ndarray,
+    part: str,
+    group: str,
+    cid: str,
+    cfg: ProcessConfig,
+    context_frames: int,
+):
+    total_samples = audio.shape[-1]
+    total_vad_frames = vad.shape[-1]
+
+    audio_out = os.path.join(cfg.fisher_path, part, "seg", "audio", group, cid)
+    vad_out = os.path.join(cfg.fisher_path, part, "seg", "vad", group, cid)
+    os.makedirs(audio_out, exist_ok=True)
+    os.makedirs(vad_out, exist_ok=True)
+
+    segments = []
+    start_sample = 0
+    seg_idx = 0
+
+    while start_sample + SEG_SAMPLES <= total_samples:
+        end_sample = start_sample + SEG_SAMPLES
+
+        chunk_audio = audio[:, start_sample:end_sample]
+        pcm16 = np.round(np.clip(chunk_audio, -1.0, 1.0) * 32767).astype(np.int16)
+        audio_file = os.path.join(audio_out, f"{seg_idx}.npy")
+        if not os.path.exists(audio_file):
+            np.save(audio_file, pcm16)
+
+        curr_vad_start = int((start_sample / SAMPLE_RATE) * VAD_HZ)
+        curr_vad_end = curr_vad_start + SEG_VAD_FRAMES
+
+        ext_start = max(0, curr_vad_start - context_frames)
+        ext_end = min(total_vad_frames, curr_vad_end + context_frames)
+
+        chunk_vad = vad[:, ext_start:ext_end]
+
+        pad_left = max(0, context_frames - curr_vad_start)
+        pad_right = max(0, (curr_vad_end + context_frames) - total_vad_frames)
+
+        if pad_left > 0 or pad_right > 0:
+            chunk_vad = np.pad(
+                chunk_vad, ((0, 0), (pad_left, pad_right)), mode="constant"
+            )
+
+        vad_file = os.path.join(vad_out, f"{seg_idx}.npy")
+        if not os.path.exists(vad_file):
+            np.save(vad_file, chunk_vad.astype(np.float32))
+
+        segments.append(
+            SegmentMeta(
+                part=part,
+                group=group,
+                cid=cid,
+                seg_id=seg_idx,
+                start_t=start_sample / SAMPLE_RATE,
+                end_t=end_sample / SAMPLE_RATE,
+                nframes=total_vad_frames,
+            )
+        )
+
+        start_sample += STRIDE_SAMPLES
+        seg_idx += 1
+
+    return segments
+
+
+def process_entry(
+    part: str,
+    group: str,
+    cid: str,
+    cfg: ProcessConfig,
+    context_frames: int,
+):
+    audio_path = os.path.join(cfg.fisher_path, part, "audio", group, f"{cid}.wav")
+    vad_dir = os.path.join(cfg.fisher_path, part, "regions_vap", group)
+
+    if not os.path.exists(audio_path):
+        return []
+
+    # Both channels are kept: the speaker-based setup needs one per speaker and
+    # the mono setups downmix in the dataloader.
+    audio_tensor, duration = load_audio(audio_path, mono=False)
+    audio = audio_tensor.numpy()
+
+    nframes = int(round(duration * VAD_HZ))
+    vad_tensor = load_fisher_vad(vad_dir, cid, nframes)
+    vad = vad_tensor.numpy()
+
+    return extract_and_save_chunks(audio, vad, part, group, cid, cfg, context_frames)
+
+
+def run_pipeline(cfg: ProcessConfig, context_frames: int):
+    pattern = os.path.join(cfg.fisher_path, "*", "regions_vap", "*", "*_words.txt")
+    all_files = glob(pattern)
+
+    entries = {"train": [], "val": [], "test": []}
+    seen = set()
+
+    for vad_file in tqdm(all_files, desc="Building Dataset"):
+        parts = vad_file.split(os.sep)
+        part, group = parts[-4], parts[-2]
+        cid = "_".join(os.path.basename(vad_file).split("_")[:3])
+
+        if (part, group, cid) in seen:
+            continue
+        seen.add((part, group, cid))
+
+        split = get_dataset_split(group, cfg)
+        if split not in cfg.active_splits:
+            continue
+
+        segments = process_entry(part, group, cid, cfg, context_frames)
+        entries[split].extend(format_meta(s) for s in segments)
+
+    for split in cfg.active_splits:
+        lines = entries.get(split, [])
+        if lines:
+            out_file = os.path.join(cfg.fisher_path, f"{split}.txt")
+            with open(out_file, "w") as f:
+                f.writelines(lines)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="config/yaml/vap_role.yaml")
+    parser.add_argument(
+        "--context-sec",
+        type=float,
+        default=CONTEXT_SEC,
+        help="VAD kept on each side of a segment; must cover the widest projection window",
+    )
+    args = parser.parse_args()
+    yaml_cfg = init_yaml_config(args.config)
+
+    config = ProcessConfig(
+        fisher_path=yaml_cfg["fisher_path"],
+        val_groups=VAL_GROUPS,
+        test_groups=TEST_GROUPS,
+        active_splits={"train", "val", "test"},
+    )
+    run_pipeline(config, int(round(args.context_sec * VAD_HZ)))