icassp2027-experiment/easycom_manifest.py

#!/usr/bin/env python3
"""Build and validate an EgoSpeak-compatible EasyCom raw-media manifest.

EasyCom stores each session as approximately one-minute MP4 segments named by
their session-relative start time (``mm-ss-SSS.mp4``).  This script preserves
that global time origin and optionally checks EgoSpeak ``target_perframe``
arrays at 5 Hz.

Example:
    python easycom_manifest.py \
      --dataset-root /data1/yfzx/wangx/Datasets/easycom/extracted_easycom \
      --target-root /path/to/EasyCom/target_perframe \
      --output easycom_manifest.json
"""

from __future__ import annotations

import argparse
import json
import re
import subprocess
from fractions import Fraction
from pathlib import Path
from typing import Any


TASK_FPS = 5
TRAIN_SESSIONS = (
    "Session_7",
    "Session_11",
    "Session_10",
    "Session_6",
    "Session_9",
    "Session_4",
    "Session_8",
    "Session_5",
    "Session_12",
)
TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
ALL_SESSIONS = TRAIN_SESSIONS + TEST_SESSIONS
TIMECODE_RE = re.compile(r"^(\d{2})-(\d{2})-(\d{3})$")


def parse_timecode(stem: str) -> float:
    match = TIMECODE_RE.fullmatch(stem)
    if match is None:
        raise ValueError(f"Expected mm-ss-SSS filename, got {stem!r}")
    minutes, seconds, milliseconds = (int(value) for value in match.groups())
    if seconds >= 60:
        raise ValueError(f"Invalid seconds in timecode {stem!r}")
    return minutes * 60.0 + seconds + milliseconds / 1000.0


def probe_video(path: Path) -> dict[str, Any]:
    completed = subprocess.run(
        [
            "ffprobe",
            "-v",
            "error",
            "-select_streams",
            "v:0",
            "-show_entries",
            "stream=width,height,avg_frame_rate,nb_frames,duration:format=duration",
            "-of",
            "json",
            str(path),
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    payload = json.loads(completed.stdout)
    streams = payload.get("streams", [])
    if not streams:
        raise RuntimeError(f"No video stream found in {path}")
    stream = streams[0]
    rate_text = stream.get("avg_frame_rate", "0/0")
    if rate_text in ("0/0", "N/A"):
        raise RuntimeError(f"No usable video frame rate in {path}")
    fps = float(Fraction(rate_text))

    frame_text = stream.get("nb_frames")
    if frame_text not in (None, "N/A"):
        num_frames = int(frame_text)
    else:
        duration_text = stream.get("duration") or payload.get("format", {}).get("duration")
        if duration_text is None:
            raise RuntimeError(f"No frame count or duration in {path}")
        num_frames = round(float(duration_text) * fps)

    # EgoSpeak processes EasyCom at 20 video frames/s and groups four frames
    # per 5 Hz task step.  The released MP4s are nominally already 20 FPS.
    duration = num_frames / fps
    frames_at_20hz = round(duration * 20.0)
    width = int(stream["width"])
    height = int(stream["height"])
    return {
        "width": width,
        "height": height,
        "fps": fps,
        "num_frames": num_frames,
        "video_duration_seconds": duration,
        "frames_at_20hz": frames_at_20hz,
        # EasyCom videos may contain an ID banner to the left of the 1920px
        # video. EgoSpeak crops it and retains the rightmost video region.
        "egospeak_crop": {
            "x": max(0, width - 1920),
            "y": 0,
            "width": min(1920, width),
            "height": height,
        },
    }


def find_target(target_root: Path, session: str) -> Path | None:
    direct = target_root / f"{session}.npy"
    if direct.is_file():
        return direct
    matches = sorted(target_root.rglob(f"{session}.npy"))
    if len(matches) > 1:
        raise RuntimeError(
            f"Multiple targets found for {session}: "
            + ", ".join(str(path) for path in matches)
        )
    return matches[0] if matches else None


def target_metadata(path: Path) -> dict[str, Any]:
    try:
        import numpy as np
    except ImportError as exc:
        raise RuntimeError("Install numpy to inspect target_perframe arrays") from exc

    target = np.load(path, mmap_mode="r")
    if target.ndim not in (1, 2):
        raise ValueError(f"Unexpected target shape in {path}: {target.shape}")
    if target.ndim == 2 and target.shape[1] != 3:
        raise ValueError(f"Expected three classes in {path}, got {target.shape}")
    return {
        "path": str(path.resolve()),
        "shape": list(target.shape),
        "num_steps": int(target.shape[0]),
        "duration_seconds_at_5hz": float(target.shape[0] / TASK_FPS),
    }


def build_manifest(
    dataset_root: Path,
    *,
    target_root: Path | None,
    probe_media: bool,
) -> dict[str, Any]:
    dataset_root = dataset_root.resolve()
    video_root = dataset_root / "Main" / "Video_Compressed"
    if not video_root.is_dir():
        raise FileNotFoundError(f"EasyCom video directory not found: {video_root}")

    sessions: list[dict[str, Any]] = []
    all_issues: list[str] = []
    all_discontinuities: list[str] = []
    for session in ALL_SESSIONS:
        session_dir = video_root / session
        if not session_dir.is_dir():
            all_issues.append(f"missing session directory: {session_dir}")
            continue

        clip_paths = sorted(
            session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem)
        )
        if not clip_paths:
            all_issues.append(f"no MP4 clips in {session_dir}")
            continue

        clips: list[dict[str, Any]] = []
        session_issues: list[str] = []
        session_discontinuities: list[str] = []
        previous_source_end: float | None = None
        concatenated_20hz_frames = 0
        for path in clip_paths:
            source_start = parse_timecode(path.stem)
            media = probe_video(path) if probe_media else None
            duration = media["video_duration_seconds"] if media is not None else None
            source_end = source_start + duration if duration is not None else None

            if previous_source_end is not None:
                delta = source_start - previous_source_end
                if abs(delta) > 0.05:
                    kind = "gap" if delta > 0 else "overlap"
                    session_discontinuities.append(
                        f"{kind} of {abs(delta):.3f}s before {path.name}"
                    )
            if source_end is not None:
                previous_source_end = source_end

            record: dict[str, Any] = {
                "path": str(path.resolve()),
                "timecode": path.stem,
                "source_session_start_seconds": source_start,
            }
            if media is not None and source_end is not None:
                concat_start_frame = concatenated_20hz_frames
                concatenated_20hz_frames += int(media["frames_at_20hz"])
                record.update(
                    {
                        **media,
                        "source_session_end_seconds": source_end,
                        "egospeak_concat_start_seconds": concat_start_frame / 20.0,
                        "egospeak_concat_end_seconds": concatenated_20hz_frames / 20.0,
                        "target_start_index": concat_start_frame // 4,
                        "target_end_index_exclusive": concatenated_20hz_frames // 4,
                    }
                )
            clips.append(record)

        source_session_end = (
            max(float(clip["source_session_end_seconds"]) for clip in clips)
            if probe_media
            else None
        )
        split = "train" if session in TRAIN_SESSIONS else "test"
        session_record: dict[str, Any] = {
            "session": session,
            "split": split,
            "clips": clips,
            "issues": session_issues,
            "source_timeline_discontinuities": session_discontinuities,
        }
        if source_session_end is not None:
            session_record["source_session_end_seconds"] = source_session_end
            session_record["egospeak_concat_duration_seconds"] = (
                concatenated_20hz_frames / 20.0
            )
            session_record["expected_5hz_steps"] = concatenated_20hz_frames // 4

        if target_root is not None:
            target_path = find_target(target_root.resolve(), session)
            if target_path is None:
                session_issues.append(f"target_perframe missing for {session}")
            else:
                target = target_metadata(target_path)
                session_record["target"] = target
                if source_session_end is not None:
                    difference = target["num_steps"] - (concatenated_20hz_frames // 4)
                    session_record["target_minus_media_steps"] = difference
                    # Small encoder/container rounding differences are normal;
                    # anything beyond one task step needs manual inspection.
                    if abs(difference) > 1:
                        session_issues.append(
                            f"target/media length mismatch: {difference:+d} steps"
                        )

        all_issues.extend(f"{session}: {issue}" for issue in session_issues)
        all_discontinuities.extend(
            f"{session}: {item}" for item in session_discontinuities
        )
        sessions.append(session_record)

    return {
        "dataset": "EasyCom",
        "dataset_root": str(dataset_root),
        "video_root": str(video_root),
        "task_fps": TASK_FPS,
        "prediction_step_seconds": 1.0 / TASK_FPS,
        "class_names": [
            "background",
            "camera wearer speaking",
            "other people speaking",
        ],
        "future_target_offsets": list(range(1, 11)),
        "train_sessions": list(TRAIN_SESSIONS),
        "test_sessions": list(TEST_SESSIONS),
        "sessions": sessions,
        "source_timeline_discontinuities": all_discontinuities,
        "issues": all_issues,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-root", type=Path, required=True)
    parser.add_argument("--target-root", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--no-probe",
        action="store_true",
        help="Skip ffprobe duration and continuity checks.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    manifest = build_manifest(
        args.dataset_root,
        target_root=args.target_root,
        probe_media=not args.no_probe,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(manifest, indent=2) + "\n")

    num_clips = sum(len(session["clips"]) for session in manifest["sessions"])
    print(
        f"Saved {len(manifest['sessions'])} sessions / {num_clips} clips "
        f"to {args.output}"
    )
    if manifest["issues"]:
        print(f"Found {len(manifest['issues'])} issue(s):")
        for issue in manifest["issues"]:
            print(f"  - {issue}")
    else:
        print("No missing-session or target-length issues found.")
    discontinuities = manifest["source_timeline_discontinuities"]
    if discontinuities:
        print(
            f"Observed {len(discontinuities)} source-timeline discontinuity(ies); "
            "EgoSpeak removes these when concatenating retained clips."
        )


if __name__ == "__main__":
    main()