tools/caption.py

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

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

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

Directory mode (input is a folder):
    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.py videos_dir/

    # Save to a specific JSON file
    caption.py videos_dir/ --output-path /path/to/captions.json

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

    # Custom prompt and model
    caption.py videos_dir/ --prompt "Describe the human motion only." --model gemini-2.0-flash
"""

import json
import os
import time
from pathlib import Path

import typer
from rich.console import Console

console = Console()

_DEFAULT_PROMPT = (
    "Describe this video in detail. The caption should include the main subject, motions, setting etc."
    "Be concise but thorough using several narrative sentences."
)


def _upload_and_wait(client, video_path: Path):
    """Upload a video to the Gemini File API and wait until it is ready."""
    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}"
        )

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


def caption_video(video_path: Path, client, model: str, prompt: str) -> str:
    """Upload a video to Gemini and return the generated caption string."""
    video_file = _upload_and_wait(client, video_path)
    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


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="Path to input video file or directory containing .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 {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.0-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,
    ),
) -> None:
    """Caption videos using the Gemini API.

    All captions are stored together in a single JSON file keyed by video filename stem.
    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.

    Examples:
        # Single file (writes video_captions.json next to input)
        caption.py video.mp4

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

        # Save to a specific JSON file
        caption.py videos_dir/ --output-path /path/to/captions.json

        # Custom prompt and model
        caption.py videos_dir/ --prompt "Describe the human motion only." --model gemini-2.0-flash

        # Re-process all videos
        caption.py videos_dir/ --override
    """
    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}[/]")

    # Resolve video list and output JSON path
    if input_path.is_file():
        video_files = [input_path]
        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}")
        json_path = output_path or input_path / "captions.json"

    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}[/]"
        )

    console.print(f"Found [bold]{len(video_files)}[/] video(s) → [bold green]{json_path}[/]")

    processed = skipped = 0
    for video_file in video_files:
        key = video_file.stem
        if key in captions and not override:
            console.print(f"[yellow]Skipping[/] {video_file.name} (already captioned)")
            skipped += 1
            continue

        console.print(f"Processing [bold blue]{video_file.name}[/]...")
        caption = caption_video(video_file, client=client, model=model, prompt=prompt)
        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",
        )

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

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


if __name__ == "__main__":
    app()