322 lines
11 KiB
Python
322 lines
11 KiB
Python
|
#!/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()
|