code/src/extract_word_clips.py

"""All LRS2/LRS3 preprocessing in one CLI, as four subcommands forming the
data pipeline (run in this order):

  scan   Count word frequencies across LRS2/LRS3 alignment files (fast — no
         video decoding). Use it to choose a target vocabulary.
  fetch  Stream LRS3 shards from a HuggingFace dataset repo ONE AT A TIME:
         download shard -> clip target words -> delete shard. Bounded disk
         usage, resumable. Also handles the LRS2 tar via --files.
  clip   Cut per-word .mp4 clips out of local tars/directories using the
         per-word timestamps in the alignment files.
  roi    Convert per-word .mp4 clips into training-ready mouth-ROI .npy
         files (MediaPipe FaceLandmarker, 96x96 grayscale uint8, 29 frames).

Typical end-to-end flow:

    # 1. what vocabulary is available? (samples a few HF shards)
    python src/extract_word_clips.py fetch --repo Ainncy/LRS3 --shards 0-2 \\
        --scan_only --freq_out word_freq.csv

    # 2. sweep all LRS3 shards for the chosen words
    python src/extract_word_clips.py fetch --repo Ainncy/LRS3 --shards 0-99 \\
        --words_file vocab.txt --out_root data/raw_clips

    # 3. same words out of the local LRS2 tar
    python src/extract_word_clips.py clip --lrs2 data/raw/lrs2.tar \\
        --words_file vocab.txt --out_root data/raw_clips

    # 4. mouth-ROI extraction for training
    python src/extract_word_clips.py roi --in_root data/raw_clips \\
        --out_root data/lrs_processed
"""
from __future__ import annotations

import argparse
import csv
import glob
import hashlib
import json
import os
import sys
import tempfile
from collections import Counter

if __package__ in (None, ""):  # allow `python src/extract_word_clips.py`
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import cv2
import numpy as np
from tqdm import tqdm

from src.lrs_alignment import (DEFAULT_STOPWORDS, count_word_frequencies,
                               parse_word_timings)
from src.lrs_source import LRSSource, open_sources
from src.preprocess import extract_mouth_clip

VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv")
SPLITS = ("train", "val", "test")


# =========================== shared helpers =================================

def add_filter_args(ap: argparse.ArgumentParser) -> None:
    ap.add_argument("--min_duration", type=float, default=0.2)
    ap.add_argument("--max_duration", type=float, default=2.0)
    ap.add_argument("--min_conf", type=float, default=0.0)


def add_vocab_args(ap: argparse.ArgumentParser) -> None:
    ap.add_argument("--words", default="", help="comma-separated word list")
    ap.add_argument("--words_file", default="",
                    help="file with one word per line (CSV rows also work — "
                        "the first comma-separated field is taken)")
    ap.add_argument("--top_n", type=int, default=0,
                    help="auto-select the N most frequent content words")
    ap.add_argument("--min_count", type=int, default=1)
    ap.add_argument("--max_per_word", type=int, default=0,
                    help="cap clips kept per word, 0 = unlimited")


def add_output_args(ap: argparse.ArgumentParser) -> None:
    ap.add_argument("--out_root", default="data/raw_clips")
    ap.add_argument("--scratch",
                    default=os.path.join(tempfile.gettempdir(),
                                         "lrs_clip_scratch"))
    ap.add_argument("--val_pct", type=int, default=5)
    ap.add_argument("--test_pct", type=int, default=5)
    ap.add_argument("--window_frames", type=int, default=29,
                    help="cut a fixed window of this many frames centered on "
                        "the word (LRW convention; keeps co-articulation "
                        "context). 0 = cut the exact word span instead")
    ap.add_argument("--overwrite", action="store_true")


def load_vocab(args, sources_fn=None) -> set:
    if args.words:
        return {w.strip().upper() for w in args.words.split(",") if w.strip()}
    if args.words_file:
        with open(args.words_file) as f:
            return {ln.split(",")[0].strip().upper() for ln in f
                    if ln.strip() and not ln.lower().startswith("word")}
    if getattr(args, "top_n", 0) and sources_fn is not None:
        print(f"Scanning to auto-select top {args.top_n} words ...")
        counts = count_word_frequencies(sources_fn(), args.min_duration,
                                        args.max_duration, args.min_conf)
        for w in DEFAULT_STOPWORDS:
            counts.pop(w, None)
        counts = {w: c for w, c in counts.items() if c >= args.min_count}
        ranked = sorted(counts.items(), key=lambda kv: -kv[1])[:args.top_n]
        vocab = {w for w, _ in ranked}
        print(f"Auto-selected {len(vocab)} words: {sorted(vocab)}")
        return vocab
    raise SystemExit("Provide --words, --words_file, or --top_n")


def assign_split(speaker: str, val_pct: int, test_pct: int) -> str:
    """Speaker-disjoint split: every clip of a speaker lands in the same
    split, so test speakers are never seen in training (no identity leakage
    inflating accuracy)."""
    h = int(hashlib.md5(speaker.encode()).hexdigest(), 16) % 100
    if h < test_pct:
        return "test"
    if h < test_pct + val_pct:
        return "val"
    return "train"


def cut_segment(cap, fps: float, start: float, end: float,
                window_frames: int = 29):
    """Cut frames for one word. window_frames > 0 cuts a fixed-length window
    centered on the word's midpoint (LRW convention — retains lip motion
    context and avoids heavy padding on short words); 0 cuts the exact span."""
    if window_frames > 0:
        center_f = int(round((start + end) / 2 * fps))
        start_f = max(center_f - window_frames // 2, 0)
        n = window_frames
    else:
        start_f = max(int(round(start * fps)), 0)
        n = max(int(round(end * fps)) - start_f, 1)
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_f)
    frames = []
    for _ in range(n):
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(frame)
    return frames


def write_clip(frames, fps: float, out_path: str) -> bool:
    if len(frames) < 2:
        return False
    h, w = frames[0].shape[:2]
    writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"),
                             fps, (w, h))
    for f in frames:
        writer.write(f)
    writer.release()
    return True


def _manifest_writer(out_root: str):
    """Append-mode manifest of every clip's provenance — enables audits and
    speaker-aware analyses later (which video, speaker, and timestamps each
    training clip came from)."""
    os.makedirs(out_root, exist_ok=True)
    path = os.path.join(out_root, "manifest.csv")
    new = not os.path.exists(path)
    f = open(path, "a", newline="")
    w = csv.writer(f)
    if new:
        w.writerow(["word", "split", "speaker", "source", "member",
                    "start", "end", "output"])
    return f, w


def clip_source(source: LRSSource, vocab: set, kept_counts: dict,
                out_root: str, scratch: str, min_duration: float = 0.2,
                max_duration: float = 2.0, min_conf: float = 0.0,
                max_per_word: int = 0, val_pct: int = 5, test_pct: int = 5,
                window_frames: int = 29, overwrite: bool = False) -> int:
    """Clip every target-word occurrence out of one LRSSource.

    Mutates kept_counts in place (shared across sources so --max_per_word
    caps globally). Returns the number of clips written.
    """
    written = 0
    os.makedirs(scratch, exist_ok=True)
    mf, manifest = _manifest_writer(out_root)
    src_name = os.path.basename(source.path)
    try:
        with source:
            entries = list(source.entries())
            for entry in tqdm(entries, desc=src_name):
                text = source.read_text(entry.txt_member)
                rows = parse_word_timings(text)
                if not rows:
                    continue
                targets = [
                    (w, s, e) for w, s, e, conf in rows
                    if w in vocab
                    and min_duration <= (e - s) <= max_duration
                    and conf >= min_conf
                    and (max_per_word == 0 or kept_counts[w] < max_per_word)
                ]
                if not targets:
                    continue

                video_path = source.extract_video(entry.mp4_member, scratch)
                try:
                    cap = cv2.VideoCapture(video_path)
                    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
                    for idx, (word, start, end) in enumerate(targets):
                        if max_per_word and kept_counts[word] >= max_per_word:
                            continue
                        split = assign_split(entry.speaker, val_pct, test_pct)
                        out_dir = os.path.join(out_root, word, split)
                        out_path = os.path.join(out_dir,
                                                f"{entry.key}_{idx}.mp4")
                        if os.path.exists(out_path) and not overwrite:
                            continue
                        frames = cut_segment(cap, fps, start, end,
                                             window_frames)
                        os.makedirs(out_dir, exist_ok=True)
                        if write_clip(frames, fps, out_path):
                            kept_counts[word] += 1
                            written += 1
                            manifest.writerow(
                                [word, split, entry.speaker, src_name,
                                 entry.mp4_member, f"{start:.2f}",
                                 f"{end:.2f}", out_path])
                    cap.release()
                finally:
                    source.cleanup_video(video_path)
    finally:
        mf.close()
    return written


def report_counts(kept_counts: dict, total: int, out_root: str) -> None:
    print(f"\nDone. Wrote {total} word clips to {out_root}")
    for word in sorted(kept_counts, key=lambda w: -kept_counts[w]):
        if kept_counts[word]:
            print(f"  {word:<20s}{kept_counts[word]:>8d}")
    print(f"\nNext: python src/extract_word_clips.py roi "
          f"--in_root {out_root} --out_root data/lrs_processed")


# ============================== scan =========================================

def cmd_scan(args) -> None:
    sources = list(open_sources(args.lrs2, args.lrs3_glob))
    if not sources:
        raise SystemExit("Nothing to scan — provide --lrs2 and/or --lrs3_glob")
    print(f"Scanning {len(sources)} source(s) ...")
    counts = count_word_frequencies(sources, args.min_duration,
                                    args.max_duration, args.min_conf)
    if not args.include_stopwords:
        for w in DEFAULT_STOPWORDS:
            counts.pop(w, None)
    counts = {w: c for w, c in counts.items() if c >= args.min_count}

    ranked = sorted(counts.items(), key=lambda kv: -kv[1])
    print(f"\n{'word':<20s}{'count':>10s}")
    for word, c in ranked[:args.top]:
        print(f"{word:<20s}{c:>10d}")
    print(f"\n{len(counts)} distinct words pass filters; total occurrences "
          f"kept: {sum(counts.values())}")
    if args.out:
        with open(args.out, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["word", "count"])
            w.writerows(ranked)
        print(f"Wrote full table to {args.out}")


# ============================== clip =========================================

def cmd_clip(args) -> None:
    if not args.lrs2 and not args.lrs3_glob:
        raise SystemExit("Provide --lrs2 and/or --lrs3_glob")
    vocab = load_vocab(
        args, sources_fn=lambda: list(open_sources(args.lrs2, args.lrs3_glob)))
    kept_counts = {w: 0 for w in vocab}
    sources = list(open_sources(args.lrs2, args.lrs3_glob))
    print(f"Clipping from {len(sources)} source(s) into {args.out_root} ...")
    total = 0
    for source in sources:
        total += clip_source(
            source, vocab, kept_counts, args.out_root, args.scratch,
            min_duration=args.min_duration, max_duration=args.max_duration,
            min_conf=args.min_conf, max_per_word=args.max_per_word,
            val_pct=args.val_pct, test_pct=args.test_pct,
            window_frames=args.window_frames, overwrite=args.overwrite)
    report_counts(kept_counts, total, args.out_root)


# ============================== fetch ========================================

def _load_state(out_root: str) -> dict:
    path = os.path.join(out_root, ".processed_shards.json")
    if os.path.isfile(path):
        with open(path) as f:
            return json.load(f)
    return {"done": [], "kept_counts": {}}


def _save_state(out_root: str, state: dict) -> None:
    os.makedirs(out_root, exist_ok=True)
    with open(os.path.join(out_root, ".processed_shards.json"), "w") as f:
        json.dump(state, f, indent=2)


def _parse_shards(spec: str) -> list:
    out = []
    for part in spec.split(","):
        part = part.strip()
        if "-" in part:
            a, b = part.split("-")
            out.extend(range(int(a), int(b) + 1))
        elif part:
            out.append(int(part))
    return sorted(set(out))


def cmd_fetch(args) -> None:
    if bool(args.repo) == bool(args.from_dir):
        raise SystemExit("Pick exactly one of --repo or --from_dir")

    vocab = set()
    if not args.scan_only:
        vocab = load_vocab(args)
        print(f"Vocabulary: {len(vocab)} words")

    if args.from_dir:
        worklist = [(os.path.basename(p), p) for p in
                    sorted(glob.glob(os.path.join(args.from_dir, "*.tar")))]
        if not worklist:
            raise SystemExit(f"No .tar files in {args.from_dir}")
    elif args.files:
        worklist = [(f.strip(), None) for f in args.files.split(",")
                    if f.strip()]
    else:
        worklist = [(args.name_template.format(n=n), None)
                    for n in _parse_shards(args.shards)]

    state = _load_state(args.out_root)
    freq: Counter = Counter()
    kept_counts = {w: int(state["kept_counts"].get(w, 0)) for w in vocab}

    for shard_name, local_path in worklist:
        if shard_name in state["done"] and not args.scan_only:
            print(f"[skip] {shard_name} (already processed)")
            continue

        downloaded = False
        if local_path is None:
            print(f"[fetch] {shard_name} ...")
            try:
                from huggingface_hub import hf_hub_download
                local_path = hf_hub_download(
                    repo_id=args.repo, filename=shard_name,
                    repo_type="dataset", local_dir=args.download_dir)
                downloaded = True
            except Exception as e:  # noqa: BLE001
                print(f"[error] could not fetch {shard_name}: {e}")
                continue

        try:
            source = LRSSource(local_path)
            if args.scan_only:
                freq.update(count_word_frequencies(
                    [source], args.min_duration, args.max_duration,
                    args.min_conf))
            else:
                n = clip_source(
                    source, vocab, kept_counts, args.out_root, args.scratch,
                    min_duration=args.min_duration,
                    max_duration=args.max_duration, min_conf=args.min_conf,
                    max_per_word=args.max_per_word, val_pct=args.val_pct,
                    test_pct=args.test_pct,
                    window_frames=args.window_frames,
                    overwrite=args.overwrite)
                print(f"[done] {shard_name}: {n} clips")
                state["done"].append(shard_name)
                state["kept_counts"] = kept_counts
                _save_state(args.out_root, state)
        finally:
            if downloaded and not args.keep_shards:
                try:
                    os.remove(local_path)
                except OSError:
                    pass

    if args.scan_only:
        for w in DEFAULT_STOPWORDS:
            freq.pop(w, None)
        ranked = sorted(freq.items(), key=lambda kv: -kv[1])
        with open(args.freq_out, "w", newline="") as f:
            wtr = csv.writer(f)
            wtr.writerow(["word", "count"])
            wtr.writerows(ranked)
        print(f"\nTop 30 of {len(ranked)} words (full table -> "
              f"{args.freq_out}):")
        for w, c in ranked[:30]:
            print(f"  {w:<20s}{c:>8d}")
    else:
        report_counts(kept_counts, sum(kept_counts.values()), args.out_root)


# ============================== roi ==========================================

def cmd_roi(args) -> None:
    words = sorted(d for d in os.listdir(args.in_root)
                   if os.path.isdir(os.path.join(args.in_root, d)))
    if getattr(args, "words", ""):
        keep = {w.strip().upper() for w in args.words.split(",") if w.strip()}
        words = [w for w in words if w in keep]
    if not words:
        raise SystemExit(f"No word folders under {args.in_root}")

    total, failed = 0, 0
    for word in words:
        for split in SPLITS:
            src_dir = os.path.join(args.in_root, word, split)
            if not os.path.isdir(src_dir):
                continue
            dst_dir = os.path.join(args.out_root, word, split)
            os.makedirs(dst_dir, exist_ok=True)
            vids = [f for f in os.listdir(src_dir)
                    if f.lower().endswith(VIDEO_EXTS)]
            for name in tqdm(vids, desc=f"{word}/{split}", leave=False):
                out_path = os.path.join(dst_dir,
                                        os.path.splitext(name)[0] + ".npy")
                if os.path.exists(out_path) and not args.overwrite:
                    continue
                try:
                    clip = extract_mouth_clip(
                        os.path.join(src_dir, name),
                        out_size=args.out_size, num_frames=args.num_frames)
                    np.save(out_path,
                            (clip * 255.0).round().astype(np.uint8))
                    total += 1
                except Exception as e:  # noqa: BLE001
                    failed += 1
                    print(f"  [skip] {name}: {e}")

    print(f"Done. Wrote {total} ROI clips, {failed} failed, "
          f"to {args.out_root}")


# ============================== main =========================================

def main() -> None:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("scan", help="word-frequency report (local sources)")
    p.add_argument("--lrs2", default="", help="LRS2 tar file or extracted dir")
    p.add_argument("--lrs3_glob", default="",
                   help='glob for local LRS3 shards, e.g. "data/pretrain.*.tar"')
    add_filter_args(p)
    p.add_argument("--min_count", type=int, default=1)
    p.add_argument("--top", type=int, default=100)
    p.add_argument("--include_stopwords", action="store_true")
    p.add_argument("--out", default="", help="optional CSV output path")
    p.set_defaults(func=cmd_scan)

    p = sub.add_parser("clip", help="cut word clips from local tars/dirs")
    p.add_argument("--lrs2", default="")
    p.add_argument("--lrs3_glob", default="")
    add_vocab_args(p)
    add_filter_args(p)
    add_output_args(p)
    p.set_defaults(func=cmd_clip)

    p = sub.add_parser("fetch",
                       help="stream shards from a HF dataset repo, one at "
                           "a time (download -> clip -> delete, resumable)")
    p.add_argument("--repo", default="", help="HF dataset repo id, ORG/NAME")
    p.add_argument("--from_dir", default="",
                   help="process already-downloaded shards instead (kept)")
    p.add_argument("--shards", default="0-99",
                   help="shard numbers, e.g. '0-99' or '0-3,17'")
    p.add_argument("--name_template", default="pretrain.{n:02d}.tar")
    p.add_argument("--files", default="",
                   help="comma-separated explicit repo filenames to fetch "
                       "(overrides --shards; use for the LRS2 tar, e.g. "
                       "'LRS2-dataset/lrs2.tar.gz')")
    p.add_argument("--keep_shards", action="store_true")
    p.add_argument("--scan_only", action="store_true")
    p.add_argument("--freq_out", default="word_freq.csv")
    p.add_argument("--download_dir",
                   default=os.path.join(tempfile.gettempdir(),
                                        "lrs_shard_downloads"))
    add_vocab_args(p)
    add_filter_args(p)
    add_output_args(p)
    p.set_defaults(func=cmd_fetch)

    p = sub.add_parser("roi", help="word clips -> mouth-ROI .npy for training")
    p.add_argument("--in_root", required=True)
    p.add_argument("--out_root", required=True)
    p.add_argument("--words", default="",
                   help="only process these word folders (comma-separated); "
                       "used to shard ROI extraction across parallel workers")
    p.add_argument("--out_size", type=int, default=96)
    p.add_argument("--num_frames", type=int, default=29)
    p.add_argument("--overwrite", action="store_true")
    p.set_defaults(func=cmd_roi)

    args = ap.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()