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,217 @@
+import argparse
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from config.setup import init_yaml_config
+from data.utils import load_audio
+from preprocess.vap.utils import (
+    EVENT_PADDING_SEC,
+    EVENT_SEG_LEN_SEC,
+    SAMPLE_RATE,
+    TEST_GROUPS,
+    VAL_GROUPS,
+    ProcessConfig,
+    event_crop,
+    get_dataset_split,
+    print_dataset_summary,
+)
+
+SEG_LEN_SEC: int = EVENT_SEG_LEN_SEC
+PADDING_SEC: float = EVENT_PADDING_SEC
+
+
+AUDIO_DIRS = {1: "audio", 2: "audio_stereo"}
+
+
+@dataclass(slots=True)
+class TurnEvent:
+    part: str
+    group: str
+    conv_id: str
+    region_start: float
+    region_end: float
+    prev_spk: str
+    next_spk: str
+    timing: str
+    gap_type: str
+    label: str
+    crop_start: float = 0.0
+    crop_end: float = 0.0
+
+
+def parse_line(line: str) -> TurnEvent:
+    tokens = line.strip().split()
+    if len(tokens) != 10:
+        return None
+
+    entry = TurnEvent(
+        part=tokens[0],
+        group=tokens[1],
+        conv_id=tokens[2],
+        region_start=float(tokens[3]),
+        region_end=float(tokens[4]),
+        prev_spk=tokens[5],
+        next_spk=tokens[6],
+        timing=tokens[7],
+        gap_type=tokens[8],
+        label=tokens[9],
+    )
+
+    crop = event_crop(entry.timing, entry.region_start, entry.region_end)
+    if crop is None:
+        return None
+    entry.crop_start, entry.crop_end = crop
+
+    return entry
+
+
+def group_by_conversation(entries: List[TurnEvent]):
+    grouped = defaultdict(list)
+    for entry in entries:
+        grouped[(entry.part, entry.group, entry.conv_id)].append(entry)
+    return dict(grouped)
+
+
+def format_output_line(entry: TurnEvent):
+    location = f"{entry.part} {entry.group} {entry.conv_id} {entry.region_start:.2f} {entry.region_end:.2f}"
+    event = f"{entry.prev_spk} {entry.next_spk} {entry.timing} {entry.gap_type} {entry.label}"
+    return f"{location} {event}"
+
+
+def load_entries_from_disk(cfg: ProcessConfig):
+    splits = {"train": [], "val": [], "test": []}
+    path = os.path.join(cfg.fisher_path, cfg.input_file)
+
+    with open(path, "r") as f:
+        parsed_entries = (parse_line(line) for line in f)
+
+        valid_entries = (e for e in parsed_entries if e is not None)
+
+        for entry in valid_entries:
+            split_name = get_dataset_split(entry.group, cfg)
+            if split_name in cfg.active_splits:
+                splits[split_name].append(entry)
+
+    return splits
+
+
+def save_segment(entry: TurnEvent, audio: np.ndarray, cfg: ProcessConfig, channels: int):
+    seg_samples = int(SEG_LEN_SEC * SAMPLE_RATE)
+    start_sample = int(entry.crop_start * SAMPLE_RATE)
+    end_sample = start_sample + seg_samples
+
+    if end_sample > audio.shape[1]:
+        return False
+
+    seg = audio[:, start_sample:end_sample]
+    if seg.shape[1] != seg_samples:
+        return False
+
+    out_dir = os.path.join(
+        cfg.fisher_path,
+        entry.part,
+        "tune",
+        AUDIO_DIRS[channels],
+        entry.group,
+        entry.conv_id,
+    )
+    os.makedirs(out_dir, exist_ok=True)
+
+    out_path = os.path.join(out_dir, f"{entry.region_start:.2f}.npy")
+    if not os.path.exists(out_path):
+        pcm16 = (seg.clamp(-1.0, 1.0) * 32767).round().to(torch.int16)
+        np.save(out_path, pcm16.numpy())
+
+    return True
+
+
+def process_split(
+    split_name: str, entries: List[TurnEvent], cfg: ProcessConfig, channels: int
+):
+    if not entries:
+        return [], []
+
+    out_lines: List[str] = []
+    successful_events: List[TurnEvent] = []
+    grouped_convos = group_by_conversation(entries)
+
+    for (part, group, cid), convo_entries in tqdm(
+        grouped_convos.items(), desc=f"Processing {split_name}"
+    ):
+        audio_sub_dir = "p1" if part == "p1" else "p2"
+        audio_path = os.path.join(
+            cfg.fisher_path, audio_sub_dir, "audio", group, f"{cid}.wav"
+        )
+
+        if not os.path.exists(audio_path):
+            continue
+
+        audio, _ = load_audio(audio_path, mono=channels == 1)
+
+        for entry in convo_entries:
+            if save_segment(entry, audio, cfg, channels):
+                out_lines.append(format_output_line(entry))
+                successful_events.append(entry)
+
+    return out_lines, successful_events
+
+
+def run_pipeline(cfg: ProcessConfig, channels: int = 1):
+    splits = load_entries_from_disk(cfg)
+    all_saved_events = []
+
+    for split_name in cfg.active_splits:
+        entries = splits.get(split_name, [])
+        if not entries:
+            continue
+
+        out_lines, success_events = process_split(split_name, entries, cfg, channels)
+        all_saved_events.extend(success_events)
+
+        # The manifest lists which events have audio, so it is written once per
+        # channel layout and stays identical between them.
+        if out_lines:
+            out_path = os.path.join(cfg.fisher_path, f"{split_name}_events.txt")
+            with open(out_path, "w") as f:
+                f.write("\n".join(out_lines) + "\n")
+
+    if all_saved_events:
+        print_dataset_summary(all_saved_events)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="config/yaml/vap_role.yaml")
+    parser.add_argument(
+        "--splits",
+        nargs="+",
+        choices=("train", "val", "test"),
+        default=["val", "test"],
+        help="splits to cut clips for; training reads segments, not events, so "
+        "train clips are only needed to enlarge the probe's tuning pool",
+    )
+    parser.add_argument(
+        "--channels",
+        type=int,
+        choices=(1, 2),
+        default=1,
+        help="1 writes tune/audio for the mono setups, 2 writes tune/audio_stereo for original VAP",
+    )
+    args = parser.parse_args()
+    yaml_cfg = init_yaml_config(args.config)
+
+    config = ProcessConfig(
+        fisher_path=yaml_cfg["fisher_path"],
+        input_file="turn_events.txt",
+        val_groups=VAL_GROUPS,
+        test_groups=TEST_GROUPS,
+        active_splits=set(args.splits),
+    )
+
+    run_pipeline(config, channels=args.channels)