Add causal person fusion and MuVAP stage-3 baselines
This commit is contained in:
Xin Wang
@@ -0,0 +1 @@
|
||||
"""Canonical active-speaker preprocessing for MuVAP+."""
|
||||
+112
@@ -0,0 +1,112 @@
|
||||
"""Dataset-specific metadata adapters.
|
||||
|
||||
Adapters do not normalize signals. They only convert source layouts into
|
||||
`SourceRecord`; the shared writer owns every media transformation.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import csv
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
from preprocess.asd.schema import SAMPLE_RATE, SourceRecord
|
||||
|
||||
|
||||
def _labels(text: str) -> list[int]:
|
||||
values = ast.literal_eval(text)
|
||||
if not isinstance(values, (list, tuple)):
|
||||
raise ValueError("labels must be a list")
|
||||
result = [int(value) for value in values]
|
||||
if any(value not in (0, 1) for value in result):
|
||||
raise ValueError("labels must contain only 0 and 1")
|
||||
return result
|
||||
|
||||
|
||||
def _existing(path: Path, sample_id: str) -> Path:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"{sample_id}: missing {path}")
|
||||
return path
|
||||
|
||||
|
||||
def clip_directory_records(
|
||||
loader_path: Path,
|
||||
visual_root: Path,
|
||||
audio_root: Path,
|
||||
split: str,
|
||||
dataset: str = "ava",
|
||||
skip_missing: bool = False,
|
||||
) -> Iterable[SourceRecord]:
|
||||
"""Read the AVA-style loader TSV and nested JPEG/WAV layout.
|
||||
|
||||
AVA and WASD ship the identical layout - `clips_videos/<split>/<clip>/
|
||||
<entity>/<timestamp>.jpg` beside `clips_audios/<split>/<clip>/<entity>.wav`
|
||||
- and the identical five-column loader TSV, so one reader serves both.
|
||||
"""
|
||||
face_by_id = {
|
||||
path.name: path for path in (visual_root / split).glob("*/*") if path.is_dir()
|
||||
}
|
||||
audio_by_id = {path.stem: path for path in (audio_root / split).glob("*/*.wav")}
|
||||
with loader_path.open(newline="") as handle:
|
||||
for row_number, row in enumerate(csv.reader(handle, delimiter="\t"), start=1):
|
||||
if len(row) < 4:
|
||||
raise ValueError(
|
||||
f"{loader_path}:{row_number}: expected at least 4 columns"
|
||||
)
|
||||
sample_id, _, fps, label_text = row[:4]
|
||||
face_path = face_by_id.get(sample_id)
|
||||
audio_path = audio_by_id.get(sample_id)
|
||||
if face_path is None or audio_path is None:
|
||||
if skip_missing:
|
||||
continue
|
||||
raise FileNotFoundError(
|
||||
f"{sample_id}: missing face directory or clip-local WAV"
|
||||
)
|
||||
yield SourceRecord(
|
||||
dataset=dataset,
|
||||
sample_id=sample_id,
|
||||
visual_path=face_path,
|
||||
audio_path=audio_path,
|
||||
source_fps=float(fps),
|
||||
labels=_labels(label_text),
|
||||
visual_kind="jpeg_directory",
|
||||
)
|
||||
|
||||
|
||||
def manifest_records(
|
||||
manifest_path: Path,
|
||||
dataset: str,
|
||||
root: Path | None = None,
|
||||
default_npy_audio_sample_rate: int = SAMPLE_RATE,
|
||||
) -> Iterable[SourceRecord]:
|
||||
"""Read the canonical TSV adapter used for AVCC and future datasets.
|
||||
|
||||
Required header fields: sample_id, visual_path, audio_path, source_fps,
|
||||
labels. Optional: start_sec, visual_kind, audio_sample_rate.
|
||||
"""
|
||||
root = root or manifest_path.parent
|
||||
with manifest_path.open(newline="") as handle:
|
||||
reader = csv.DictReader(handle, delimiter="\t")
|
||||
required = {"sample_id", "visual_path", "audio_path", "source_fps", "labels"}
|
||||
missing = required.difference(reader.fieldnames or [])
|
||||
if missing:
|
||||
raise ValueError(f"{manifest_path}: missing columns {sorted(missing)}")
|
||||
for row_number, row in enumerate(reader, start=2):
|
||||
try:
|
||||
yield SourceRecord(
|
||||
dataset=dataset,
|
||||
sample_id=row["sample_id"],
|
||||
visual_path=_existing(root / row["visual_path"], row["sample_id"]),
|
||||
audio_path=_existing(root / row["audio_path"], row["sample_id"]),
|
||||
source_fps=float(row["source_fps"]),
|
||||
labels=json.loads(row["labels"]),
|
||||
start_sec=float(row.get("start_sec") or 0.0),
|
||||
npy_audio_sample_rate=int(
|
||||
row.get("audio_sample_rate") or default_npy_audio_sample_rate
|
||||
),
|
||||
visual_kind=row.get("visual_kind") or "npy",
|
||||
)
|
||||
except (TypeError, ValueError, json.JSONDecodeError) as error:
|
||||
raise ValueError(
|
||||
f"{manifest_path}:{row_number}: invalid row"
|
||||
) from error
|
||||
+270
@@ -0,0 +1,270 @@
|
||||
"""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()
|
||||
+93
@@ -0,0 +1,93 @@
|
||||
"""Build the canonical fast-training format for AVA, WASD, and AVCC."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from preprocess.asd.adapters import clip_directory_records, manifest_records
|
||||
from preprocess.asd.writer import write_records
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("dataset", choices=["ava", "wasd", "msdwild", "avcc"])
|
||||
parser.add_argument(
|
||||
"--loader", type=Path, required=True, help="source TSV/manifest"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
help="defaults to <dataset root>/packed/<split>, beside the source media",
|
||||
)
|
||||
parser.add_argument("--visual-root", type=Path)
|
||||
parser.add_argument("--audio-root", type=Path)
|
||||
parser.add_argument("--root", type=Path, help="base for AVCC manifest paths")
|
||||
parser.add_argument("--split", default="train")
|
||||
parser.add_argument("--audio-sample-rate", type=int, default=16_000)
|
||||
parser.add_argument("--max-shard-gb", type=float, default=1.0)
|
||||
parser.add_argument(
|
||||
"--skip-missing",
|
||||
action="store_true",
|
||||
help="development subsets only; production preprocessing should fail on missing files",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def require_roots(args: argparse.Namespace) -> None:
|
||||
if args.visual_root is None or args.audio_root is None:
|
||||
raise ValueError(f"{args.dataset} requires --visual-root and --audio-root")
|
||||
|
||||
|
||||
def default_output(args: argparse.Namespace) -> Path:
|
||||
"""Keep each corpus's pack beside that corpus.
|
||||
|
||||
`<root>/clips_videos/<split>` becomes `<root>/packed/<split>`, so a dataset
|
||||
stays self-contained and can be copied to fast local storage on its own.
|
||||
"""
|
||||
if args.visual_root is not None:
|
||||
return args.visual_root.parent / "packed" / args.split
|
||||
if args.root is not None:
|
||||
return args.root / "packed" / args.split
|
||||
raise ValueError("--output is required when no dataset root is given")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = build_parser().parse_args()
|
||||
if args.output is None:
|
||||
args.output = default_output(args)
|
||||
print(f"packing into {args.output}")
|
||||
if args.output.exists() and any(args.output.iterdir()):
|
||||
raise FileExistsError(f"output directory is not empty: {args.output}")
|
||||
|
||||
if args.dataset in ("ava", "wasd"):
|
||||
require_roots(args)
|
||||
records = clip_directory_records(
|
||||
args.loader,
|
||||
args.visual_root,
|
||||
args.audio_root,
|
||||
args.split,
|
||||
dataset=args.dataset,
|
||||
skip_missing=args.skip_missing,
|
||||
)
|
||||
elif args.dataset == "msdwild":
|
||||
raise SystemExit(
|
||||
"MSDWild is packed from raw media by preprocess/asd/prep_msdwild.py, "
|
||||
"which needs all.rttm for labels; there is no loader TSV to read here"
|
||||
)
|
||||
else:
|
||||
records = manifest_records(
|
||||
args.loader,
|
||||
dataset="avcc",
|
||||
root=args.root,
|
||||
default_npy_audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
|
||||
count = write_records(
|
||||
records,
|
||||
args.output,
|
||||
max_shard_bytes=round(args.max_shard_gb * 1_000_000_000),
|
||||
)
|
||||
print(f"Wrote {count} {args.dataset} samples to {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
"""Shared source record and packed-shard schema.
|
||||
|
||||
The packed format is a *decoded mirror* of the source corpus, not a
|
||||
model-ready tensor set. Faces, labels, and audio are stored on the native
|
||||
timeline with no gain applied, so any codebase can consume a pack without
|
||||
inheriting MuVAP+'s 25 Hz model timeline or its loudness convention. Every
|
||||
model-specific transform lives in `data.dataloaders.packed` instead.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
FORMAT_NAME = "muvap-asd-mmap"
|
||||
FORMAT_VERSION = 4
|
||||
SUPPORTED_VERSIONS = (4,)
|
||||
|
||||
# The model timeline. It belongs to the loader, not to the stored data, and is
|
||||
# repeated here only because preprocessing reports frame counts against it.
|
||||
MODEL_FPS = 25.0
|
||||
|
||||
SAMPLE_RATE = 16_000
|
||||
FACE_SIZE = 112
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SourceRecord:
|
||||
dataset: str
|
||||
sample_id: str
|
||||
visual_path: Path
|
||||
audio_path: Path
|
||||
source_fps: float
|
||||
labels: Sequence[int]
|
||||
start_sec: float = 0.0
|
||||
npy_audio_sample_rate: int | None = None
|
||||
visual_kind: str = "npy"
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.source_fps <= 0:
|
||||
raise ValueError(f"{self.sample_id}: source_fps must be positive")
|
||||
if self.start_sec < 0:
|
||||
raise ValueError(f"{self.sample_id}: start_sec cannot be negative")
|
||||
if self.visual_kind not in {"npy", "jpeg_directory"}:
|
||||
raise ValueError(f"unsupported visual_kind: {self.visual_kind}")
|
||||
+93
@@ -0,0 +1,93 @@
|
||||
"""Validate packed samples and report per-dataset signal statistics."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
from data.alignment import model_frame_count
|
||||
from data.dataloaders.packed import PackedASDDataset
|
||||
from data.media import WaveformStats, normalization_gain
|
||||
from preprocess.asd.schema import MODEL_FPS
|
||||
from preprocess.asd.writer import native_audio_samples
|
||||
|
||||
|
||||
def percentile(values: list[float], points=(0, 1, 50, 95, 99, 100)) -> dict:
|
||||
return {str(point): float(np.percentile(values, point)) for point in points}
|
||||
|
||||
|
||||
def validate(roots: Path | Sequence[Path]) -> None:
|
||||
dataset = PackedASDDataset(roots)
|
||||
grouped = defaultdict(lambda: defaultdict(list))
|
||||
|
||||
for index in range(len(dataset)):
|
||||
audio, visual, vad, _, metadata = dataset[index]
|
||||
sample_id = metadata["sample_id"]
|
||||
expected = model_frame_count(
|
||||
metadata["source_frames"], metadata["source_fps"], MODEL_FPS
|
||||
)
|
||||
frames = visual.shape[0]
|
||||
if frames == 0:
|
||||
raise ValueError(f"{sample_id}: empty sample")
|
||||
if frames != expected:
|
||||
raise ValueError(
|
||||
f"{sample_id}: {metadata['source_frames']} source frames at "
|
||||
f"{metadata['source_fps']} fps imply {expected} model frames, "
|
||||
f"loader produced {frames}"
|
||||
)
|
||||
if vad.shape[0] != frames or audio.shape[-1] != frames * 640:
|
||||
raise ValueError(f"{sample_id}: stream-length mismatch")
|
||||
if visual.shape[1:] != (112, 112) or visual.min() < 0 or visual.max() > 255:
|
||||
raise ValueError(f"{sample_id}: invalid visual tensor")
|
||||
if not np.isfinite(metadata["audio_peak"]):
|
||||
raise ValueError(f"{sample_id}: invalid audio statistics")
|
||||
if metadata["audio_samples"] != native_audio_samples(
|
||||
metadata["source_frames"], metadata["source_fps"]
|
||||
):
|
||||
raise ValueError(f"{sample_id}: stored audio does not span the faces")
|
||||
|
||||
stats = grouped[metadata["dataset"]]
|
||||
stats["source_audio_peak"].append(metadata["audio_peak"])
|
||||
stats["source_audio_rms"].append(metadata["audio_rms"])
|
||||
stats["source_audio_mean"].append(metadata["audio_mean"])
|
||||
stats["speaking_ratio"].append(metadata["speaking_ratio"])
|
||||
stats["normalized_rms"].append(
|
||||
audio.square().mean().sqrt().item()
|
||||
)
|
||||
stats["normalized_peak"].append(audio.abs().max().item())
|
||||
gain = normalization_gain(
|
||||
WaveformStats(
|
||||
mean=metadata["audio_mean"],
|
||||
rms=metadata["audio_rms"],
|
||||
peak=metadata["audio_peak"],
|
||||
)
|
||||
)
|
||||
stats["gain_db"].append(20 * math.log10(max(gain, 1e-8)))
|
||||
|
||||
report = {"samples": len(dataset), "datasets": {}}
|
||||
for name, stats in grouped.items():
|
||||
report["datasets"][name] = {
|
||||
"samples": len(stats["source_audio_peak"]),
|
||||
"source_audio_peak_percentiles": percentile(stats["source_audio_peak"]),
|
||||
"source_audio_rms_percentiles": percentile(stats["source_audio_rms"]),
|
||||
"normalized_rms_percentiles": percentile(stats["normalized_rms"]),
|
||||
"normalized_peak_percentiles": percentile(stats["normalized_peak"]),
|
||||
"gain_db_percentiles": percentile(stats["gain_db"]),
|
||||
"max_abs_source_mean": max(abs(v) for v in stats["source_audio_mean"]),
|
||||
"speaking_ratio_percentiles": percentile(stats["speaking_ratio"]),
|
||||
}
|
||||
print(json.dumps(report, indent=2))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("roots", type=Path, nargs="+")
|
||||
validate(parser.parse_args().roots)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
"""Reading face crops off a source video.
|
||||
|
||||
Shared by the dataset adapters that pack from video rather than from
|
||||
pre-extracted frames.
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def crop_face(frame, bbox, gray=True, dim=(112, 112)):
|
||||
if bbox is None:
|
||||
return np.full(dim if gray else (dim[1], dim[0], 3), 128, dtype=np.uint8)
|
||||
|
||||
x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
|
||||
|
||||
if x1 >= x2 or y1 >= y2:
|
||||
return np.full(dim if gray else (dim[1], dim[0], 3), 128, dtype=np.uint8)
|
||||
|
||||
h, w = frame.shape[:2]
|
||||
cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
|
||||
size = max(x2 - x1, y2 - y1)
|
||||
half = size // 2
|
||||
|
||||
x1s, y1s = cx - half, cy - half
|
||||
x2s, y2s = x1s + size, y1s + size
|
||||
|
||||
x1c, y1c = max(0, x1s), max(0, y1s)
|
||||
x2c, y2c = min(w, x2s), min(h, y2s)
|
||||
|
||||
if x1c >= x2c or y1c >= y2c:
|
||||
return np.full(dim if gray else (dim[1], dim[0], 3), 128, dtype=np.uint8)
|
||||
|
||||
cropped = frame[y1c:y2c, x1c:x2c]
|
||||
|
||||
if gray:
|
||||
cropped = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)
|
||||
border_val = 128
|
||||
else:
|
||||
cropped = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB)
|
||||
border_val = (128, 128, 128)
|
||||
|
||||
pad_top, pad_left = y1c - y1s, x1c - x1s
|
||||
pad_bottom, pad_right = y2s - y2c, x2s - x2c
|
||||
|
||||
if pad_top > 0 or pad_bottom > 0 or pad_left > 0 or pad_right > 0:
|
||||
cropped = cv2.copyMakeBorder(
|
||||
cropped,
|
||||
pad_top,
|
||||
pad_bottom,
|
||||
pad_left,
|
||||
pad_right,
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=border_val,
|
||||
)
|
||||
|
||||
return cv2.resize(cropped, dim)
|
||||
|
||||
|
||||
def load_visual(video_path, start_frame=0, target_frames=None):
|
||||
"""Decode a native-FPS video interval without changing its timeline."""
|
||||
if start_frame < 0:
|
||||
raise ValueError("start_frame cannot be negative")
|
||||
if target_frames is not None and target_frames < 0:
|
||||
raise ValueError("target_frames cannot be negative")
|
||||
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
if not cap.isOpened():
|
||||
raise FileNotFoundError(f"could not open video: {video_path}")
|
||||
|
||||
fps = float(cap.get(cv2.CAP_PROP_FPS))
|
||||
if fps <= 0:
|
||||
cap.release()
|
||||
raise ValueError(f"video has an invalid frame rate: {video_path}")
|
||||
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
|
||||
frames = []
|
||||
while target_frames is None or len(frames) < target_frames:
|
||||
ok, frame = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
frames.append(frame)
|
||||
cap.release()
|
||||
return frames, fps
|
||||
+204
@@ -0,0 +1,204 @@
|
||||
"""Pack decoded native-rate samples into memory-mapped NumPy shards.
|
||||
|
||||
Nothing here is specific to the model that will read the shards. Faces and
|
||||
labels keep the source frame rate, audio keeps its full duration and its
|
||||
original gain, and the per-recording level statistics needed to reproduce a
|
||||
normalization travel in the index. Resampling onto a model timeline and
|
||||
applying loudness normalization are the reader's job.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from data.alignment import model_frame_count
|
||||
from data.media import (
|
||||
AudioSpec,
|
||||
load_waveform,
|
||||
slice_audio,
|
||||
validate_visual,
|
||||
waveform_stats,
|
||||
)
|
||||
from preprocess.asd.schema import (
|
||||
FACE_SIZE,
|
||||
FORMAT_NAME,
|
||||
FORMAT_VERSION,
|
||||
MODEL_FPS,
|
||||
SAMPLE_RATE,
|
||||
SourceRecord,
|
||||
)
|
||||
|
||||
|
||||
def _load_faces(record: SourceRecord, face_size: int = FACE_SIZE) -> np.ndarray:
|
||||
if record.visual_kind == "npy":
|
||||
return validate_visual(np.load(record.visual_path, allow_pickle=False))
|
||||
|
||||
paths = sorted(record.visual_path.glob("*.jpg"), key=lambda path: float(path.stem))
|
||||
if not paths:
|
||||
raise FileNotFoundError(f"{record.sample_id}: no JPEG faces")
|
||||
frames = []
|
||||
for path in paths:
|
||||
frame = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
|
||||
if frame is None:
|
||||
raise ValueError(f"{record.sample_id}: could not decode {path}")
|
||||
frames.append(
|
||||
cv2.resize(frame, (face_size, face_size), interpolation=cv2.INTER_LINEAR)
|
||||
)
|
||||
return np.stack(frames).astype(np.uint8, copy=False)
|
||||
|
||||
|
||||
def _pcm16(waveform: torch.Tensor) -> np.ndarray:
|
||||
values = waveform.squeeze(0).clamp(-1.0, 1.0)
|
||||
return (values * 32767.0).round().to(torch.int16).cpu().numpy()
|
||||
|
||||
|
||||
def native_audio_samples(source_frames: int, source_fps: float) -> int:
|
||||
"""Audio length covering exactly the span of the stored face frames."""
|
||||
return round(source_frames / source_fps * SAMPLE_RATE)
|
||||
|
||||
|
||||
class ShardWriter:
|
||||
def __init__(self, output_dir: Path, max_shard_bytes: int = 1_000_000_000):
|
||||
self.output_dir = output_dir
|
||||
self.max_shard_bytes = max_shard_bytes
|
||||
self.shard_index = 0
|
||||
self.samples_written = 0
|
||||
self._reset()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _reset(self) -> None:
|
||||
self.faces: list[np.ndarray] = []
|
||||
self.audio: list[np.ndarray] = []
|
||||
self.labels: list[np.ndarray] = []
|
||||
self.entries: list[dict] = []
|
||||
self.byte_count = 0
|
||||
self.face_count = 0
|
||||
self.audio_count = 0
|
||||
self.label_count = 0
|
||||
|
||||
def add_arrays(
|
||||
self,
|
||||
*,
|
||||
dataset: str,
|
||||
sample_id: str,
|
||||
faces: np.ndarray,
|
||||
waveform: torch.Tensor,
|
||||
labels,
|
||||
source_fps: float,
|
||||
start_sec: float = 0.0,
|
||||
) -> None:
|
||||
faces = validate_visual(faces)
|
||||
labels = np.asarray(labels, dtype=np.uint8)
|
||||
if len(faces) != len(labels):
|
||||
raise ValueError(
|
||||
f"{sample_id}: {len(faces)} faces but {len(labels)} labels"
|
||||
)
|
||||
source_frames = len(labels)
|
||||
if source_frames == 0:
|
||||
raise ValueError(f"{sample_id}: sample is empty")
|
||||
|
||||
waveform = slice_audio(
|
||||
waveform,
|
||||
SAMPLE_RATE,
|
||||
target_samples=native_audio_samples(source_frames, source_fps),
|
||||
)
|
||||
stats = waveform_stats(waveform)
|
||||
audio = _pcm16(waveform)
|
||||
|
||||
sample_bytes = faces.nbytes + audio.nbytes + labels.nbytes
|
||||
if self.entries and self.byte_count + sample_bytes > self.max_shard_bytes:
|
||||
self.flush()
|
||||
|
||||
self.entries.append(
|
||||
{
|
||||
"dataset": dataset,
|
||||
"sample_id": sample_id,
|
||||
"source_fps": source_fps,
|
||||
"source_frames": source_frames,
|
||||
"start_sec": start_sec,
|
||||
# Reported for convenience; the loader recomputes it.
|
||||
"model_frames": model_frame_count(source_frames, source_fps, MODEL_FPS),
|
||||
"face_offset": self.face_count,
|
||||
"audio_offset": self.audio_count,
|
||||
"audio_samples": len(audio),
|
||||
"label_offset": self.label_count,
|
||||
"audio_mean": stats.mean,
|
||||
"audio_rms": stats.rms,
|
||||
"audio_peak": stats.peak,
|
||||
"speaking_ratio": float(labels.mean()),
|
||||
}
|
||||
)
|
||||
self.faces.append(faces)
|
||||
self.audio.append(audio)
|
||||
self.labels.append(labels)
|
||||
self.byte_count += sample_bytes
|
||||
self.face_count += len(faces)
|
||||
self.audio_count += len(audio)
|
||||
self.label_count += len(labels)
|
||||
|
||||
def add(self, record: SourceRecord) -> None:
|
||||
faces = _load_faces(record)
|
||||
waveform, sample_rate = load_waveform(
|
||||
record.audio_path,
|
||||
AudioSpec(sample_rate=SAMPLE_RATE, mono=True),
|
||||
npy_sample_rate=record.npy_audio_sample_rate,
|
||||
)
|
||||
waveform = slice_audio(waveform, sample_rate, start_time=record.start_sec)
|
||||
self.add_arrays(
|
||||
dataset=record.dataset,
|
||||
sample_id=record.sample_id,
|
||||
faces=faces,
|
||||
waveform=waveform,
|
||||
labels=record.labels,
|
||||
source_fps=record.source_fps,
|
||||
start_sec=record.start_sec,
|
||||
)
|
||||
|
||||
def flush(self) -> None:
|
||||
if not self.entries:
|
||||
return
|
||||
shard_dir = self.output_dir / f"shard-{self.shard_index:06d}"
|
||||
shard_dir.mkdir(parents=False, exist_ok=False)
|
||||
np.save(shard_dir / "faces.npy", np.concatenate(self.faces), allow_pickle=False)
|
||||
np.save(shard_dir / "audio.npy", np.concatenate(self.audio), allow_pickle=False)
|
||||
np.save(
|
||||
shard_dir / "labels.npy", np.concatenate(self.labels), allow_pickle=False
|
||||
)
|
||||
with (shard_dir / "index.json").open("w") as handle:
|
||||
json.dump(self.entries, handle, separators=(",", ":"))
|
||||
self.samples_written += len(self.entries)
|
||||
self.shard_index += 1
|
||||
self._reset()
|
||||
|
||||
def close(self) -> None:
|
||||
self.flush()
|
||||
manifest = {
|
||||
"format": FORMAT_NAME,
|
||||
"version": FORMAT_VERSION,
|
||||
"timeline": "native",
|
||||
"audio_sample_rate": SAMPLE_RATE,
|
||||
"audio_normalization": None,
|
||||
"face_shape": [FACE_SIZE, FACE_SIZE],
|
||||
"face_color": "grayscale",
|
||||
"face_dtype": "uint8",
|
||||
"audio_dtype": "int16",
|
||||
"label_dtype": "uint8",
|
||||
"shards": self.shard_index,
|
||||
"samples": self.samples_written,
|
||||
}
|
||||
with (self.output_dir / "dataset.json").open("w") as handle:
|
||||
json.dump(manifest, handle, indent=2)
|
||||
|
||||
|
||||
def write_records(
|
||||
records: Iterable[SourceRecord], output_dir: Path, max_shard_bytes: int
|
||||
) -> int:
|
||||
writer = ShardWriter(output_dir, max_shard_bytes)
|
||||
for record in records:
|
||||
writer.add(record)
|
||||
writer.close()
|
||||
return writer.samples_written
|
||||
Reference in New Issue
Block a user