caption_transfer_reference.py

"""
Caption videos using the Gemini API and save all results to a single JSON file.

Manifest mode (input is a dataset dir with manifest.csv, or the manifest.csv itself):
    Captions the *original, unstylized* clip behind every target video — the manifest's
    `original_org` column — and rewrites dataset.json so each entry's `caption` describes
    the video content instead of the style-transfer instruction.  The instruction is kept
    under `style_prompt` so nothing is lost.

    # Caption the source clips and update dataset.json
    caption_transfer_reference.py style-transfer-reference/

    # Smoke test on 5 samples first
    caption_transfer_reference.py style-transfer-reference/ --limit 5

    # Caption only, leave dataset.json alone
    caption_transfer_reference.py style-transfer-reference/ --no-update-dataset

    # Rewrite dataset.json from an existing captions file without calling Gemini
    caption_transfer_reference.py style-transfer-reference/ --update-only

    # Caption the stylized target clips instead of the unstylized originals
    caption_transfer_reference.py style-transfer-reference/ --source-column original_video

Single-file mode (input is a video file):
    # Caption one video — writes {name}_captions.json next to the input
    caption_transfer_reference.py video.mp4

    # Save to a specific JSON file
    caption_transfer_reference.py video.mp4 --output-path /path/to/captions.json

Directory mode (input is a folder without a manifest, or with --ignore-manifest):
    Processes all .mp4 files in the folder.
    Without --output-path: writes captions.json inside the input folder.
    With --output-path: writes to the specified JSON file.

    # Caption all videos in a folder
    caption_transfer_reference.py videos_dir/

    # Re-process all videos, overriding existing captions
    caption_transfer_reference.py videos_dir/ --override

    # Custom prompt and model
    caption_transfer_reference.py videos_dir/ --prompt "Describe the main action only." --model gemini-2.0-flash

    # Caption up to 8 videos in parallel (default is 4)
    caption_transfer_reference.py videos_dir/ --concurrency 8
"""

import csv
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path

import typer
from rich.console import Console
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)

console = Console()

# Where the raw OpenVE-3M clips live; manifest paths are relative to this root.
DEFAULT_VIDEO_BASE = Path("/mnt/data/xinyuy/datasets/OpenVE-3M/videos")

_DEFAULT_PROMPT = (
"""
You are annotating a short video for training a text-to-video model.

Write one concise and factual English caption of 20-35 words.

Describe:
1. The main subjects, whether people, animals, objects, or vehicles.
2. Their general appearance, such as colour, shape, size, or clothing.
3. The broad action or event taking place, and any notable change or camera motion over time.
4. The environment, background, and setting.

Do not:
- Describe fine-grained motion frame by frame or list a detailed sequence of movements.
- Transcribe speech, on-screen text, or describe what is being said.
- Infer identities, relationships, occupations, ethnicity, brands, locations,
  intentions, or details that are not visually evident.
- Use subjective or cinematic language.

If a detail is uncertain, omit it.

Return only the caption, without explanations or formatting.
"""
)


@dataclass(frozen=True)
class CaptionTask:
    """One video to caption. `key` is what the caption is stored under in the JSON."""

    key: str
    path: Path
    label: str  # human-readable name for logs (sample name in manifest mode)


def _upload_and_wait(client, video_path: Path, verbose: bool):
    """Upload a video to the Gemini File API and wait until it is ready."""
    if verbose:
        console.print(f"  Uploading [blue]{video_path.name}[/]...")
    video_file = client.files.upload(file=str(video_path))

    while video_file.state.name == "PROCESSING":
        time.sleep(5)
        video_file = client.files.get(name=video_file.name)

    if video_file.state.name != "ACTIVE":
        raise RuntimeError(
            f"File upload failed for {video_path.name}: state={video_file.state.name}"
        )

    if verbose:
        console.print(f"  Upload complete (file id: [dim]{video_file.name}[/])")
    return video_file


def _caption_once(video_path: Path, client, model: str, prompt: str, verbose: bool) -> str:
    video_file = _upload_and_wait(client, video_path, verbose)
    try:
        response = client.models.generate_content(
            model=model,
            contents=[video_file, prompt],
        )
        return response.text.strip()
    finally:
        try:
            client.files.delete(name=video_file.name)
        except Exception:
            pass


def caption_video(
    video_path: Path, client, model: str, prompt: str, retries: int = 3, verbose: bool = False
) -> str:
    """Upload a video to Gemini and return the generated caption string.

    Bulk runs reliably hit transient rate-limit and 5xx errors, so each video gets a few
    retries with exponential backoff before it is counted as failed.
    """
    last_error: Exception | None = None
    for attempt in range(retries + 1):
        try:
            return _caption_once(video_path, client, model, prompt, verbose)
        except Exception as e:
            last_error = e
            if attempt == retries:
                break
            time.sleep(min(5 * 2**attempt, 60))
    raise last_error  # type: ignore[misc]


# ── Manifest handling ───────────────────────────────────────────────────────────
def _remap(rel_path: str) -> str:
    """Manifest paths say `global_style/`; the clips actually live in `global_style_new/`."""
    return rel_path.replace("global_style/", "global_style_new/", 1)


def resolve_source(video_base: Path, rel_path: str) -> Path | None:
    full = video_base / _remap(rel_path)
    if full.exists():
        return full
    fallback = video_base / rel_path
    return fallback if fallback.exists() else None


def load_manifest(manifest_path: Path) -> list[dict]:
    with manifest_path.open(newline="", encoding="utf-8") as fh:
        return list(csv.DictReader(fh))


def plan_manifest_tasks(
    rows: list[dict], video_base: Path, source_column: str
) -> tuple[list[CaptionTask], dict[str, str], int, list[str]]:
    """Map every manifest row to the source clip that should be captioned.

    Each row is paired with the source clip of *its own* target (`renamed_video`).  The
    style reference image is deliberately cut from a foreign clip of the same style
    (`reference_source_video`), so that donor's content must never end up in this row's
    caption — `check_donor_leakage` verifies that separation.

    Returns the deduplicated caption tasks, a {target media_path -> caption key} map used to
    rewrite dataset.json, the number of unresolvable rows, and any donor leaks found.
    """
    tasks: dict[str, CaptionTask] = {}
    media_to_key: dict[str, str] = {}
    missing = 0

    for row in rows:
        rel = (row.get(source_column) or "").strip()
        media = (row.get("renamed_video") or "").strip()
        if not rel or not media:
            missing += 1
            continue

        src = resolve_source(video_base, rel)
        if src is None:
            missing += 1
            continue

        # Key on the source clip's content hash so the cache survives a dataset rebuild,
        # where sample names (abstract_000, ...) get reshuffled but the clips do not.
        key = src.stem
        media_to_key[media] = key
        tasks.setdefault(key, CaptionTask(key=key, path=src, label=Path(media).stem))

    return list(tasks.values()), media_to_key, missing, check_donor_leakage(rows, media_to_key)


def check_donor_leakage(rows: list[dict], media_to_key: dict[str, str]) -> list[str]:
    """Flag rows whose caption source is the clip the style reference image was cut from.

    A hit means the caption would describe the reference image's video rather than the
    target's, which would leak the conditioning signal into the text.  The build pairs every
    target with a *foreign* donor, so this should always come back empty.
    """
    leaks = []
    for row in rows:
        media = (row.get("renamed_video") or "").strip()
        donor = (row.get("reference_source_video") or "").strip()
        if not media or not donor:
            continue
        key, donor_key = media_to_key.get(media), media_to_key.get(donor)
        if key is not None and key == donor_key:
            leaks.append(media)
    return leaks


def update_dataset_json(
    dataset_path: Path,
    media_to_key: dict[str, str],
    captions: dict[str, str],
    keep_style_prompt: bool,
) -> tuple[int, int]:
    """Rewrite dataset.json so `caption` describes the video instead of the style edit.

    The original style-transfer instruction moves to `style_prompt` (unless dropped).  Re-runs
    are idempotent: `style_prompt` is only filled from `caption` the first time, so a second
    pass never overwrites it with an already-replaced caption.
    """
    entries = json.loads(dataset_path.read_text(encoding="utf-8"))

    backup = dataset_path.with_suffix(dataset_path.suffix + ".bak")
    if not backup.exists():
        backup.write_text(json.dumps(entries, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
        console.print(f"Backed up original dataset to [cyan]{backup}[/]")

    updated = unresolved = 0
    for entry in entries:
        media = entry.get("media_path", "")
        key = media_to_key.get(media) or media_to_key.get(Path(media).name)
        caption = captions.get(key) if key else None

        if caption is None:
            unresolved += 1
            continue

        if keep_style_prompt:
            entry.setdefault("style_prompt", entry.get("caption", ""))
        else:
            entry.pop("style_prompt", None)
        # Assigning an existing key keeps `caption` first in the serialized entry.
        entry["caption"] = caption
        updated += 1

    dataset_path.write_text(
        json.dumps(entries, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    return updated, unresolved


app = typer.Typer(
    pretty_exceptions_enable=False,
    no_args_is_help=True,
    help="Caption videos using the Gemini API.",
)


@app.command()
def main(
    input_path: Path = typer.Argument(  # noqa: B008
        ...,
        help=(
            "Dataset directory containing manifest.csv, a manifest.csv path, "
            "a video file, or a directory of .mp4 files"
        ),
        exists=True,
    ),
    output_path: Path | None = typer.Option(  # noqa: B008
        None,
        "--output-path",
        "-o",
        help=(
            "Path to the output JSON file. "
            "Defaults to {dataset_dir}/source_captions.json (manifest mode), "
            "{input_dir}/captions.json (directory mode) "
            "or {video_stem}_captions.json (single-file mode)."
        ),
    ),
    override: bool = typer.Option(
        False,
        "--override",
        help="Re-caption videos even if they already have an entry in the JSON.",
    ),
    model: str = typer.Option(
        "gemini-2.5-flash",
        "--model",
        "-m",
        help="Gemini model to use for captioning.",
    ),
    prompt: str = typer.Option(
        _DEFAULT_PROMPT,
        "--prompt",
        "-p",
        help="Prompt sent to the model along with each video.",
    ),
    api_key: str | None = typer.Option(  # noqa: B008
        None,
        "--api-key",
        envvar="GEMINI_API_KEY",
        help="Gemini API key. Defaults to $GEMINI_API_KEY environment variable.",
        show_default=False,
    ),
    concurrency: int = typer.Option(
        4,
        "--concurrency",
        "-c",
        help="Number of videos to caption concurrently (parallel Gemini requests).",
    ),
    retries: int = typer.Option(
        3,
        "--retries",
        help="Retries per video on transient API errors, with exponential backoff.",
    ),
    video_base: Path = typer.Option(  # noqa: B008
        DEFAULT_VIDEO_BASE,
        "--video-base",
        help="Manifest mode: root the manifest's relative clip paths are resolved against.",
    ),
    source_column: str = typer.Option(
        "original_org",
        "--source-column",
        help=(
            "Manifest mode: column holding the clip to caption. "
            "'original_org' is the unstylized original; "
            "'original_video' is the stylized target's source path."
        ),
    ),
    dataset_json: Path | None = typer.Option(  # noqa: B008
        None,
        "--dataset-json",
        help="Manifest mode: dataset.json to rewrite. Defaults to {manifest_dir}/dataset.json.",
    ),
    update_dataset: bool = typer.Option(
        True,
        "--update-dataset/--no-update-dataset",
        help="Manifest mode: rewrite dataset.json captions after captioning.",
    ),
    update_only: bool = typer.Option(
        False,
        "--update-only",
        help="Manifest mode: rewrite dataset.json from the existing captions JSON, no API calls.",
    ),
    keep_style_prompt: bool = typer.Option(
        True,
        "--keep-style-prompt/--drop-style-prompt",
        help="Manifest mode: preserve the style-transfer instruction under `style_prompt`.",
    ),
    limit: int | None = typer.Option(
        None,
        "--limit",
        help="Caption at most N videos this run (useful for smoke tests).",
    ),
    ignore_manifest: bool = typer.Option(
        False,
        "--ignore-manifest",
        help="Treat a dataset directory as a plain folder of .mp4 files.",
    ),
    verbose: bool = typer.Option(
        False,
        "--verbose",
        "-v",
        help="Print per-video upload progress.",
    ),
) -> None:
    """Caption videos using the Gemini API.

    All captions are stored together in a single JSON file. Existing entries are skipped
    unless --override is set, so the run can be resumed. The JSON is written after each
    video so progress is not lost on interruption.

    Manifest mode captions the original clip behind each style-transfer target and rewrites
    dataset.json, whose `caption` field otherwise holds the style-editing instruction.

    Examples:
        # Caption the source clips of a style-transfer dataset and update dataset.json
        caption_transfer_reference.py style-transfer-reference/

        # Try 5 samples first
        caption_transfer_reference.py style-transfer-reference/ --limit 5

        # Rebuild dataset.json from captions already on disk
        caption_transfer_reference.py style-transfer-reference/ --update-only

        # Single file (writes video_captions.json next to input)
        caption_transfer_reference.py video.mp4

        # Directory (writes captions.json inside videos_dir/)
        caption_transfer_reference.py videos_dir/

        # Caption up to 8 videos in parallel
        caption_transfer_reference.py videos_dir/ --concurrency 8
    """
    # ── Resolve mode, work list and output paths ────────────────────────────────
    manifest_path: Path | None = None
    if not ignore_manifest:
        if input_path.is_file() and input_path.suffix.lower() == ".csv":
            manifest_path = input_path
        elif input_path.is_dir() and (input_path / "manifest.csv").exists():
            manifest_path = input_path / "manifest.csv"

    media_to_key: dict[str, str] = {}
    dataset_path: Path | None = None

    if manifest_path is not None:
        dataset_dir = manifest_path.parent
        rows = load_manifest(manifest_path)
        if source_column not in (rows[0] if rows else {}):
            raise typer.BadParameter(
                f"Column {source_column!r} not found in {manifest_path}. "
                f"Available: {', '.join(rows[0]) if rows else 'none'}"
            )

        tasks, media_to_key, missing, leaks = plan_manifest_tasks(
            rows, video_base, source_column
        )
        json_path = output_path or dataset_dir / "source_captions.json"
        dataset_path = dataset_json or dataset_dir / "dataset.json"

        console.print(f"Manifest [cyan]{manifest_path}[/] — [bold]{len(rows)}[/] row(s)")
        console.print(
            f"Captioning the [bold]{source_column}[/] clip of each target "
            f"([bold]{len(tasks)}[/] unique video(s) under {video_base})"
        )
        if missing:
            console.print(f"[yellow]{missing} row(s) skipped[/] — source clip not found on disk")

        # Never let a caption describe the clip its style reference image came from.
        if leaks:
            console.print(
                f"[bold red]Aborting:[/] {len(leaks)} row(s) would be captioned from the same "
                f"clip their reference image was cut from, e.g. {', '.join(leaks[:5])}"
            )
            raise typer.Exit(code=1)
        console.print(
            "[green]✓[/] verified: no caption source overlaps its reference image's clip"
        )
    elif input_path.is_file():
        tasks = [CaptionTask(key=input_path.stem, path=input_path, label=input_path.stem)]
        json_path = output_path or input_path.parent / f"{input_path.stem}_captions.json"
    else:
        video_files = sorted(input_path.glob("*.mp4"))
        if not video_files:
            raise typer.BadParameter(f"No .mp4 files found in {input_path}")
        tasks = [CaptionTask(key=v.stem, path=v, label=v.stem) for v in video_files]
        json_path = output_path or input_path / "captions.json"

    if update_only and manifest_path is None:
        raise typer.BadParameter("--update-only only applies in manifest mode.")

    json_path.parent.mkdir(parents=True, exist_ok=True)

    # Load existing captions so we can resume interrupted runs
    captions: dict[str, str] = {}
    if json_path.exists():
        captions = json.loads(json_path.read_text(encoding="utf-8"))
        console.print(
            f"Loaded [bold]{len(captions)}[/] existing caption(s) from [cyan]{json_path}[/]"
        )

    # ── Caption ─────────────────────────────────────────────────────────────────
    processed = failed = skipped = 0

    if update_only:
        console.print("[yellow]--update-only[/]: skipping captioning, reusing captions on disk")
    else:
        if api_key is None:
            api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            raise typer.BadParameter(
                "Gemini API key is required. Set GEMINI_API_KEY or pass --api-key."
            )

        try:
            from google import genai
        except ImportError as e:
            raise ImportError(
                "google-genai is not installed. Run: pip install google-genai"
            ) from e

        client = genai.Client(api_key=api_key)
        console.print(f"Using model [bold]{model}[/]")
        console.print(f"Found [bold]{len(tasks)}[/] video(s) → [bold green]{json_path}[/]")

        # Filter out already-captioned videos up front so the progress bar
        # only tracks work that actually remains to be done.
        pending: list[CaptionTask] = []
        for task_item in tasks:
            if task_item.key in captions and not override:
                skipped += 1
            else:
                pending.append(task_item)

        if skipped:
            console.print(
                f"[yellow]Skipping {skipped} video(s)[/] already present in {json_path.name}"
            )
        if limit is not None and len(pending) > limit:
            console.print(f"[yellow]--limit {limit}[/]: capping this run at {limit} video(s)")
            pending = pending[:limit]
        console.print(f"[bold]{len(pending)}[/] video(s) remaining to caption")

        write_lock = threading.Lock()

        def save_caption(key: str, caption: str) -> None:
            with write_lock:
                captions[key] = caption
                # Write after each video so progress survives interruption
                json_path.write_text(
                    json.dumps(captions, indent=2, ensure_ascii=False) + "\n",
                    encoding="utf-8",
                )

        with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            MofNCompleteColumn(),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
            console=console,
        ) as progress:
            bar = progress.add_task(
                f"Captioning videos ({concurrency} concurrent)", total=len(pending)
            )
            with ThreadPoolExecutor(max_workers=max(1, concurrency)) as executor:
                futures = {
                    executor.submit(
                        caption_video, t.path, client, model, prompt, retries, verbose
                    ): t
                    for t in pending
                }
                for future in as_completed(futures):
                    task_item = futures[future]
                    try:
                        caption = future.result()
                    except Exception as e:
                        console.print(f"  [bold red]✗[/] {task_item.label}: [red]{e}[/]")
                        failed += 1
                        progress.advance(bar)
                        continue

                    save_caption(task_item.key, caption)
                    console.print(
                        f"  [bold green]✓[/] {task_item.label}: "
                        f"[dim]{caption[:120]}{'...' if len(caption) > 120 else ''}[/]"
                    )
                    processed += 1
                    progress.advance(bar)

        console.print(
            f"\n[bold green]Captioning done.[/] Processed [bold]{processed}[/] video(s)"
            + (f", skipped [bold]{skipped}[/]" if skipped else "")
            + (f", [bold red]failed {failed}[/]" if failed else "")
            + f". Captions saved to [cyan]{json_path}[/]."
        )

    # ── Rewrite dataset.json ────────────────────────────────────────────────────
    if manifest_path is not None and update_dataset and dataset_path is not None:
        if not dataset_path.exists():
            console.print(f"[yellow]No dataset.json at {dataset_path}[/] — nothing to update.")
            return

        updated, unresolved = update_dataset_json(
            dataset_path, media_to_key, captions, keep_style_prompt
        )
        console.print(
            f"[bold green]dataset.json updated.[/] [bold]{updated}[/] entry(ies) now carry a "
            f"video caption"
            + (
                " (style instruction preserved as `style_prompt`)"
                if keep_style_prompt
                else " (style instruction dropped)"
            )
            + f" → [cyan]{dataset_path}[/]"
        )
        if unresolved:
            console.print(
                f"[yellow]{unresolved} entry(ies) still hold the style prompt[/] — no caption "
                f"yet. Re-run to fill them in."
            )


if __name__ == "__main__":
    app()