271 lines
10 KiB
Python
271 lines
10 KiB
Python
"""Crop raw MSDWild videos into aligned 30-second speaker-track samples."""
|
|
|
|
import argparse
|
|
import csv
|
|
import shutil
|
|
import sys
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from data.utils import load_audio # noqa: E402
|
|
from preprocess.asd.video import crop_face, load_visual # noqa: E402
|
|
from preprocess.asd.writer import ShardWriter # noqa: E402
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Detection:
|
|
bbox: tuple[float, float, float, float]
|
|
|
|
|
|
def read_annotations(path: Path) -> dict[int, dict[str, Detection]]:
|
|
annotations: dict[int, dict[str, Detection]] = defaultdict(dict)
|
|
with path.open(newline="") as handle:
|
|
for row_number, row in enumerate(csv.reader(handle), start=1):
|
|
if len(row) < 8:
|
|
raise ValueError(f"{path}:{row_number}: expected 8 columns")
|
|
try:
|
|
frame = int(row[0])
|
|
speaker_id = row[2]
|
|
bbox = tuple(float(value) for value in row[3:7])
|
|
int(row[7])
|
|
except ValueError as error:
|
|
raise ValueError(f"{path}:{row_number}: invalid annotation") from error
|
|
annotations[frame][speaker_id] = Detection(bbox)
|
|
return dict(annotations)
|
|
|
|
|
|
def read_rttm(path: Path) -> dict[str, dict[str, list[tuple[float, float]]]]:
|
|
activity: dict[str, dict[str, list[tuple[float, float]]]] = defaultdict(
|
|
lambda: defaultdict(list)
|
|
)
|
|
with path.open() as handle:
|
|
for line_number, line in enumerate(handle, start=1):
|
|
parts = line.split()
|
|
if not parts:
|
|
continue
|
|
if len(parts) < 8 or parts[0] != "SPEAKER":
|
|
raise ValueError(f"{path}:{line_number}: invalid RTTM row")
|
|
try:
|
|
start = float(parts[3])
|
|
duration = float(parts[4])
|
|
except ValueError as error:
|
|
raise ValueError(f"{path}:{line_number}: invalid time") from error
|
|
activity[parts[1]][parts[7]].append((start, start + duration))
|
|
return {
|
|
video_id: {speaker: sorted(intervals) for speaker, intervals in speakers.items()}
|
|
for video_id, speakers in activity.items()
|
|
}
|
|
|
|
|
|
def retained_speakers(
|
|
annotations: dict[int, dict[str, Detection]],
|
|
start_frame: int,
|
|
frame_count: int,
|
|
min_face_ratio: float,
|
|
) -> list[str]:
|
|
counts: dict[str, int] = defaultdict(int)
|
|
for frame in range(start_frame, start_frame + frame_count):
|
|
for speaker_id in annotations.get(frame, {}):
|
|
counts[speaker_id] += 1
|
|
minimum = min_face_ratio * frame_count
|
|
return sorted(speaker_id for speaker_id, count in counts.items() if count >= minimum)
|
|
|
|
|
|
def clip_starts(first_face_frame: int, last_face_frame: int, stride_frames: int):
|
|
return range(first_face_frame, last_face_frame + 1, stride_frames)
|
|
|
|
|
|
def video_metadata(path: Path) -> tuple[float, int]:
|
|
cap = cv2.VideoCapture(str(path))
|
|
if not cap.isOpened():
|
|
raise FileNotFoundError(f"could not open video: {path}")
|
|
fps = float(cap.get(cv2.CAP_PROP_FPS))
|
|
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
cap.release()
|
|
if fps <= 0 or frame_count <= 0:
|
|
raise ValueError(f"invalid video metadata: {path}")
|
|
return fps, frame_count
|
|
|
|
|
|
def process_video(
|
|
video_path: Path,
|
|
annotation_path: Path,
|
|
activity: dict[str, list[tuple[float, float]]],
|
|
writer: ShardWriter,
|
|
clip_seconds: float,
|
|
overlap_seconds: float,
|
|
min_face_ratio: float,
|
|
min_clip_seconds: float,
|
|
) -> int:
|
|
annotations = read_annotations(annotation_path)
|
|
if not annotations:
|
|
return 0
|
|
|
|
fps, video_frames = video_metadata(video_path)
|
|
clip_frames = round(clip_seconds * fps)
|
|
stride_frames = round((clip_seconds - overlap_seconds) * fps)
|
|
min_clip_frames = max(1, round(min_clip_seconds * fps))
|
|
if clip_frames <= 0 or stride_frames <= 0:
|
|
raise ValueError("clip and stride must each contain at least one frame")
|
|
waveform, _ = load_audio(video_path)
|
|
first_face = min(annotations)
|
|
last_face = min(max(annotations), video_frames - 1)
|
|
written = 0
|
|
|
|
for segment_id, start_frame in enumerate(
|
|
clip_starts(first_face, last_face, stride_frames)
|
|
):
|
|
# Take whatever is left at the tail. Windows used to be dropped unless
|
|
# they were complete, which discarded every video shorter than one clip
|
|
# - about a third of MSDWild. The loader pads and chunks, so a short
|
|
# sample costs nothing.
|
|
window_frames = min(clip_frames, video_frames - start_frame)
|
|
if window_frames < min_clip_frames:
|
|
break
|
|
speakers = retained_speakers(
|
|
annotations, start_frame, window_frames, min_face_ratio
|
|
)
|
|
if not speakers:
|
|
continue
|
|
|
|
frames, decoded_fps = load_visual(video_path, start_frame, window_frames)
|
|
if not np.isclose(decoded_fps, fps):
|
|
raise ValueError(f"frame-rate changed while decoding {video_path}")
|
|
if len(frames) < min_clip_frames:
|
|
continue
|
|
window_frames = len(frames)
|
|
|
|
start_time = start_frame / fps
|
|
start_sample = round(start_time * 16_000)
|
|
target_samples = round(window_frames / fps * 16_000)
|
|
audio = waveform[:, start_sample : start_sample + target_samples]
|
|
|
|
for speaker_id in speakers:
|
|
faces = []
|
|
labels = np.zeros(window_frames, dtype=np.uint8)
|
|
for offset, frame in enumerate(frames):
|
|
detection = annotations.get(start_frame + offset, {}).get(speaker_id)
|
|
faces.append(crop_face(frame, detection.bbox if detection else None))
|
|
intervals = activity.get(speaker_id, [])
|
|
for offset in range(window_frames):
|
|
timestamp = start_time + (offset + 0.5) / fps
|
|
labels[offset] = any(start <= timestamp < end for start, end in intervals)
|
|
faces = np.stack(faces).astype(np.uint8, copy=False)
|
|
sample_id = f"{video_path.stem}_{segment_id:05d}_{speaker_id}"
|
|
writer.add_arrays(
|
|
dataset="msdwild",
|
|
sample_id=sample_id,
|
|
faces=faces,
|
|
waveform=audio,
|
|
labels=labels,
|
|
source_fps=fps,
|
|
start_sec=start_time,
|
|
)
|
|
written += 1
|
|
return written
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--input", type=Path, required=True)
|
|
parser.add_argument("--rttm", type=Path)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=Path,
|
|
help="defaults to <dataset root>/packed/<split>, beside the source media",
|
|
)
|
|
parser.add_argument(
|
|
"--split",
|
|
default="all",
|
|
help="official split to pack: all, few.train, few.val, many.val",
|
|
)
|
|
parser.add_argument("--clip-seconds", type=float, default=30.0)
|
|
# Windows no longer overlap: the loader computes projection targets over a
|
|
# whole track before chunking, so duplicated frames buy nothing and would
|
|
# both inflate the pack and over-weight this corpus in an epoch.
|
|
parser.add_argument("--overlap-seconds", type=float, default=0.0)
|
|
parser.add_argument("--min-clip-seconds", type=float, default=2.0)
|
|
parser.add_argument("--min-face-ratio", type=float, default=0.9)
|
|
parser.add_argument("--max-shard-gb", type=float, default=1.0)
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
if args.clip_seconds <= 0:
|
|
raise ValueError("clip-seconds must be positive")
|
|
if not 0 <= args.overlap_seconds < args.clip_seconds:
|
|
raise ValueError("overlap-seconds must be in [0, clip-seconds)")
|
|
if not 0 < args.min_face_ratio <= 1:
|
|
raise ValueError("min-face-ratio must be in (0, 1]")
|
|
if args.max_shard_gb <= 0:
|
|
raise ValueError("max-shard-gb must be positive")
|
|
if not 0 < args.min_clip_seconds <= args.clip_seconds:
|
|
raise ValueError("min-clip-seconds must be in (0, clip-seconds]")
|
|
if args.output is None:
|
|
args.output = args.input.parent / "packed" / args.split
|
|
print(f"packing into {args.output}")
|
|
if args.output.exists():
|
|
if not args.overwrite and any(args.output.iterdir()):
|
|
raise FileExistsError(f"output directory is not empty: {args.output}")
|
|
if args.overwrite:
|
|
shutil.rmtree(args.output)
|
|
|
|
args.output.mkdir(parents=True)
|
|
videos = sorted(args.input.glob("*.mp4"))
|
|
if not videos:
|
|
raise FileNotFoundError(f"no MP4 files found in {args.input}")
|
|
|
|
rttm_path = args.rttm or args.input.parent / "all.rttm"
|
|
if not rttm_path.exists():
|
|
raise FileNotFoundError(f"missing RTTM labels: {rttm_path}")
|
|
activity = read_rttm(rttm_path)
|
|
|
|
if args.split != "all":
|
|
# The official split lists are themselves RTTMs; their video ids are
|
|
# the split membership.
|
|
split_path = args.input.parent / f"{args.split}.rttm"
|
|
if not split_path.exists():
|
|
raise FileNotFoundError(f"missing split list: {split_path}")
|
|
wanted = set(read_rttm(split_path))
|
|
videos = [video for video in videos if video.stem in wanted]
|
|
if not videos:
|
|
raise ValueError(f"no videos matched split {args.split}")
|
|
print(f"{args.split}: {len(videos)} videos")
|
|
total = 0
|
|
writer = ShardWriter(
|
|
args.output, max_shard_bytes=round(args.max_shard_gb * 1_000_000_000)
|
|
)
|
|
try:
|
|
for index, video_path in enumerate(videos, start=1):
|
|
annotation_path = video_path.with_suffix(".csv")
|
|
if not annotation_path.exists():
|
|
raise FileNotFoundError(f"missing annotations: {annotation_path}")
|
|
total += process_video(
|
|
video_path,
|
|
annotation_path,
|
|
activity.get(video_path.stem, {}),
|
|
writer,
|
|
args.clip_seconds,
|
|
args.overlap_seconds,
|
|
args.min_face_ratio,
|
|
args.min_clip_seconds,
|
|
)
|
|
print(f"[{index}/{len(videos)}] {video_path.stem}: {total} samples")
|
|
finally:
|
|
writer.close()
|
|
print(f"Wrote {total} samples to {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|