extract_word_clips.py

"""
extract_word_clips.py

Scans LRS2/LRS3-style annotation files for a target word list and extracts
matching video clips from the full-sentence utterances.

Expects each dataset laid out the standard Oxford VGG way:
    <root>/<speaker_id>/<utterance_id>.mp4
    <root>/<speaker_id>/<utterance_id>.txt

Each .txt annotation is expected to contain a word-alignment table like:

    Text:  THE PEOPLE HAVE VOTED FOR CHANGE

    Conf:  1

    WORD START END ASDSCORE
    THE    0.20 0.32 1.0000
    PEOPLE 0.32 0.68 1.0000
    ...

IMPORTANT: before running this on the full dataset, open one real .txt file
from your download and confirm it matches this layout. Column spacing/labels
have varied slightly across Oxford's dataset releases, and it's a five-minute
check that saves you from silently parsing zero words.

Requires ffmpeg installed and available on PATH.
"""

import csv
import subprocess
from pathlib import Path

# ---------------- CONFIG: edit these before running ----------------

TARGET_WORDS = {
    "START", "STOP", "YES", "NO",
    # ... fill in the rest of your 20 words, UPPERCASE, no punctuation
}

DATASETS = [
    # Add one entry per dataset split you have access to.
    # "fps" should match the actual frame rate of that dataset's video (LRS2/LRS3 are 25fps).
    {"name": "LRS2_main", "root": "/path/to/lrs2/main", "fps": 25},
    {"name": "LRS2_pretrain", "root": "/path/to/lrs2/pretrain", "fps": 25},
    {"name": "LRS3_trainval", "root": "/path/to/lrs3/trainval", "fps": 25},
    {"name": "LRS3_pretrain", "root": "/path/to/lrs3/pretrain", "fps": 25},
]

OUTPUT_DIR = Path("/path/to/output/clips")
MANIFEST_PATH = Path("/path/to/output/manifest.csv")

# Match this to whatever frame count your pretrained frontend expects.
# LRW uses 29 frames (~1.16s at 25fps) centered on the word.
CLIP_NUM_FRAMES = 29

# ---------------------------------------------------------------------


def parse_annotation(txt_path):
    """Returns a list of (WORD, start_sec, end_sec) from one annotation file."""
    entries = []
    in_table = False
    with open(txt_path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            line = line.strip()
            if line.upper().startswith("WORD") and "START" in line.upper():
                in_table = True
                continue
            if in_table and line:
                parts = line.split()
                if len(parts) >= 3:
                    word, start, end = parts[0], parts[1], parts[2]
                    try:
                        entries.append((word.upper(), float(start), float(end)))
                    except ValueError:
                        continue
    return entries


def build_word_index(datasets, target_words):
    """
    Phase 1: text-only pass. Walks every annotation file and records where
    each target word occurs. Does not touch any video files.
    """
    index = []
    for ds in datasets:
        root = Path(ds["root"])
        txt_files = list(root.rglob("*.txt"))
        print(f"[{ds['name']}] scanning {len(txt_files)} annotation files...")
        for txt_path in txt_files:
            video_path = txt_path.with_suffix(".mp4")
            if not video_path.exists():
                continue
            for word, start, end in parse_annotation(txt_path):
                if word in target_words:
                    index.append({
                        "dataset": ds["name"],
                        "fps": ds["fps"],
                        "video_path": str(video_path),
                        "speaker_id": video_path.parent.name,
                        "word": word,
                        "start": start,
                        "end": end,
                    })
    return index


def extract_clip(entry, num_frames, out_path):
    """
    Phase 2: only runs on matched instances. Seeks to the word's midpoint
    and cuts a fixed-length window of frames around it.
    """
    fps = entry["fps"]
    center_time = (entry["start"] + entry["end"]) / 2
    half_window_sec = (num_frames / 2) / fps
    clip_start = max(0.0, center_time - half_window_sec)

    cmd = [
        "ffmpeg", "-y",
        "-ss", f"{clip_start:.3f}",
        "-i", entry["video_path"],
        "-frames:v", str(num_frames),
        "-an",
        str(out_path),
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print("Phase 1: indexing target words across all annotation files...")
    index = build_word_index(DATASETS, TARGET_WORDS)
    print(f"Found {len(index)} total instances across {len(TARGET_WORDS)} target words.\n")

    counts = {}
    for e in index:
        counts[e["word"]] = counts.get(e["word"], 0) + 1
    print("Per-word instance counts:")
    for w in sorted(TARGET_WORDS):
        print(f"  {w}: {counts.get(w, 0)}")

    print("\nPhase 2: extracting matched clips (only these instances, not the full dataset)...")

    with open(MANIFEST_PATH, "w", newline="") as mf:
        writer = csv.writer(mf)
        writer.writerow(["word", "dataset", "speaker_id", "source_video", "start", "end", "output_path"])
        for i, entry in enumerate(index):
            out_name = f"{entry['word']}_{entry['dataset']}_{entry['speaker_id']}_{i}.mp4"
            out_path = OUTPUT_DIR / out_name
            try:
                extract_clip(entry, CLIP_NUM_FRAMES, out_path)
                writer.writerow([entry["word"], entry["dataset"], entry["speaker_id"],
                                  entry["video_path"], entry["start"], entry["end"], str(out_path)])
            except subprocess.CalledProcessError:
                print(f"  failed on {entry['video_path']} @ {entry['start']}s, skipping")

            if (i + 1) % 200 == 0:
                print(f"  extracted {i + 1}/{len(index)}...")

    print(f"\nDone. Manifest written to {MANIFEST_PATH}")


if __name__ == "__main__":
    main()