scripts/prepare_data.py

#!/usr/bin/env python3
"""Prepare training samples from raw talking-head videos.

Turns raw videos into the sample-directory format the training dataloader
consumes with on-the-fly preprocessing (see DATA.md): one directory per
81-frame clip containing

    sub_clip.mp4   aligned 512x512 face crop, 25 fps
    audio.wav      matching audio slice, 16 kHz mono PCM
    prompt.txt     the text prompt

plus train/val list files (one sample directory per line) that
``OMNIAVATAR_DATA_LIST`` / ``OMNIAVATAR_VAL_LIST`` point at.

Face detection + 512x512 affine alignment uses the exact same LatentSync
pipeline as inference (scripts/inference/_common.py), so training samples
match what the model sees at test time. Videos are converted to 25 fps CFR
first when needed; segments where face detection fails are skipped.

Example:
    python scripts/prepare_data.py \
        --input_dir /path/to/raw_videos \
        --output_dir ./data/v2v_training_data \
        --mask_path /path/to/mask.png \
        --val_count 2
"""
import argparse
import os
import subprocess
import sys
import tempfile

import cv2
import numpy as np
import torch

# --- Path setup (mirror the inference scripts) ----------------------------
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
LIPFORCING_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, ".."))
sys.path.insert(0, LIPFORCING_ROOT)
sys.path.insert(0, os.path.join(SCRIPT_DIR, "inference"))  # for `import _common`

from _common import _get_ffmpeg, load_image_processor, save_frames_as_video  # noqa: E402

FPS = 25  # training frame rate (OmniAvatar convention)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Prepare 512x512 aligned training samples from raw talking-head videos."
    )
    parser.add_argument("--input_dir", type=str, default=None,
                        help="Directory of raw videos (*.mp4/*.mov/*.avi/*.mkv)")
    parser.add_argument("--videos", type=str, nargs="+", default=None,
                        help="Explicit list of raw video paths (alternative to --input_dir)")
    parser.add_argument("--output_dir", type=str, required=True,
                        help="Where sample directories and list files are written")
    parser.add_argument("--mask_path", type=str, required=True,
                        help="Path to LatentSync mask.png (needed by the image processor)")
    parser.add_argument("--prompt", type=str, default="a person talking",
                        help="Prompt written to each sample's prompt.txt")
    parser.add_argument("--clip_frames", type=int, default=81,
                        help="Frames per training clip (81 = 21 VAE latents, the training length)")
    parser.add_argument("--max_clips_per_video", type=int, default=None,
                        help="Cap on clips extracted per video (default: all full clips)")
    parser.add_argument("--val_count", type=int, default=2,
                        help="Number of trailing sample dirs assigned to the val list")
    parser.add_argument("--train_list", type=str, default="train_list.txt",
                        help="Train list filename (written inside --output_dir)")
    parser.add_argument("--val_list", type=str, default="val_list.txt",
                        help="Val list filename (written inside --output_dir)")
    parser.add_argument("--skip_existing", action="store_true",
                        help="Skip sample dirs that already contain all three files")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device for face detection")
    return parser.parse_args()


def list_videos(args):
    if (args.input_dir is None) == (args.videos is None):
        raise ValueError("Provide exactly one of --input_dir or --videos")
    if args.videos is not None:
        return list(args.videos)
    exts = (".mp4", ".mov", ".avi", ".mkv")
    return sorted(
        os.path.join(args.input_dir, f)
        for f in os.listdir(args.input_dir)
        if f.lower().endswith(exts)
    )


def ensure_25fps(video_path):
    """Return a 25 fps CFR version of *video_path* (temp file if converted).

    Returns:
        (path, tmp_path_or_None) — tmp_path must be deleted by the caller.
    """
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()
    if abs(fps - FPS) < 0.05:
        return video_path, None
    print(f"  [fps] {fps:.2f} -> {FPS} CFR")
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp_path = tmp.name
    tmp.close()
    cmd = [
        _get_ffmpeg(), "-y", "-loglevel", "error", "-nostdin",
        "-i", video_path, "-r", str(FPS), "-an",
        "-c:v", "libx264", "-crf", "13", tmp_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        os.remove(tmp_path)
        raise RuntimeError(f"ffmpeg 25fps conversion failed:\n{result.stderr}")
    return tmp_path, tmp_path


def extract_audio_slice(video_path, start_s, dur_s, out_wav):
    """Extract a 16 kHz mono PCM slice [start_s, start_s+dur_s) from *video_path*."""
    cmd = [
        _get_ffmpeg(), "-y", "-loglevel", "error", "-nostdin",
        "-ss", f"{start_s:.4f}", "-t", f"{dur_s:.4f}", "-i", video_path,
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", out_wav,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg audio extraction failed:\n{result.stderr}")


def align_segment(frames, image_processor):
    """Affine-align a list of RGB frames to 512x512 face crops.

    Returns:
        [N, 512, 512, 3] uint8 array, or None if detection failed on any frame.
    """
    aligned = []
    for i, frame in enumerate(frames):
        try:
            face, _box, _matrix = image_processor.affine_transform(frame)
        except RuntimeError as e:
            print(f"    [skip segment] face detection failed on frame {i}: {e}")
            return None
        if isinstance(face, torch.Tensor):
            face = face.permute(1, 2, 0).cpu().numpy()
        aligned.append(np.asarray(face, dtype=np.uint8))
    return np.stack(aligned, axis=0)


def process_video(video_path, args, image_processor, sample_dirs):
    stem = os.path.splitext(os.path.basename(video_path))[0]
    print(f"\n=== {stem} ===")

    cfr_path, tmp_cfr = ensure_25fps(video_path)
    clip_dur = args.clip_frames / FPS
    n_clips = 0
    try:
        cap = cv2.VideoCapture(cfr_path)
        if not cap.isOpened():
            print(f"  [skip video] cannot open {cfr_path}")
            return
        # Reset the aligner's temporal smoothing for each new video.
        image_processor.restorer.p_bias = None

        seg_idx = 0
        frames = []
        while True:
            ret, frame = cap.read()
            if ret:
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            if len(frames) == args.clip_frames or (not ret and frames):
                if len(frames) < args.clip_frames:
                    break  # drop the trailing partial clip
                sample_dir = os.path.join(args.output_dir, f"{stem}_{seg_idx:04d}")
                seg_start_s = seg_idx * clip_dur
                seg_idx += 1

                done = all(
                    os.path.isfile(os.path.join(sample_dir, f))
                    for f in ("sub_clip.mp4", "audio.wav", "prompt.txt")
                )
                if args.skip_existing and done:
                    print(f"  [exists] {sample_dir}")
                    sample_dirs.append(sample_dir)
                    n_clips += 1
                elif (aligned := align_segment(frames, image_processor)) is not None:
                    os.makedirs(sample_dir, exist_ok=True)
                    save_frames_as_video(aligned, os.path.join(sample_dir, "sub_clip.mp4"), fps=FPS)
                    # Audio comes from the ORIGINAL file — the 25fps CFR
                    # conversion preserves wall-clock time, so offsets match.
                    extract_audio_slice(video_path, seg_start_s, clip_dur,
                                        os.path.join(sample_dir, "audio.wav"))
                    with open(os.path.join(sample_dir, "prompt.txt"), "w") as f:
                        f.write(args.prompt + "\n")
                    print(f"  [ok] {sample_dir} ({args.clip_frames} frames)")
                    sample_dirs.append(sample_dir)
                    n_clips += 1
                frames = []
                if args.max_clips_per_video and n_clips >= args.max_clips_per_video:
                    break
            if not ret:
                break
        cap.release()
    finally:
        if tmp_cfr is not None and os.path.exists(tmp_cfr):
            os.remove(tmp_cfr)
    if n_clips == 0:
        print(f"  [warn] no full {args.clip_frames}-frame clips extracted from {stem}")


def main():
    args = parse_args()
    videos = list_videos(args)
    if not videos:
        raise ValueError("No input videos found")
    print(f"Preparing {len(videos)} video(s) -> {args.output_dir}")
    os.makedirs(args.output_dir, exist_ok=True)

    image_processor = load_image_processor(args.mask_path, args.device)

    sample_dirs = []
    for video_path in videos:
        process_video(video_path, args, image_processor, sample_dirs)

    if not sample_dirs:
        raise RuntimeError("No samples were produced — check the input videos.")

    sample_dirs = [os.path.abspath(d) for d in sample_dirs]
    val_count = min(args.val_count, max(len(sample_dirs) - 1, 0))
    train_dirs = sample_dirs[:len(sample_dirs) - val_count] if val_count else sample_dirs
    val_dirs = sample_dirs[len(sample_dirs) - val_count:] if val_count else sample_dirs[-1:]

    train_list = os.path.join(args.output_dir, args.train_list)
    val_list = os.path.join(args.output_dir, args.val_list)
    with open(train_list, "w") as f:
        f.write("\n".join(train_dirs) + "\n")
    with open(val_list, "w") as f:
        f.write("\n".join(val_dirs) + "\n")

    print(f"\nWrote {len(train_dirs)} train / {len(val_dirs)} val samples")
    print(f"  OMNIAVATAR_DATA_LIST={train_list}")
    print(f"  OMNIAVATAR_VAL_LIST={val_list}")


if __name__ == "__main__":
    main()