processing/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

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

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

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

console = Console()

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

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

Describe:
1. The number of visible people.
2. Their general appearance and clothing.
3. The environment or scene.
4. The broad activity, such as walking, dancing, exercising, or talking.
5. The background and the setting. Also the overall style or mood of the video if it is visually evident.

Do not:
- Describe exact limb positions, joint coordinates, or frame-by-frame motion.
- List a detailed sequence of poses.
- Transcribe speech or describe what the people are saying.
- Infer identities, relationships, occupations, ethnicity, intentions, or
  details that are not visually evident.
- Use subjective or cinematic language.
- Begin with “The video shows”.

If a detail is uncertain, omit it.

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


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.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).",
    ),
) -> 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

        # Caption up to 8 videos in parallel
        caption.py videos_dir/ --concurrency 8
    """
    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}[/]")

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

    if skipped:
        console.print(f"[yellow]Skipping {skipped} video(s)[/] already present in {json_path.name}")
    console.print(f"[bold]{len(pending)}[/] video(s) remaining to caption")

    processed = failed = 0
    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:
        task = 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, video_file, client, model, prompt): video_file
                for video_file in pending
            }
            for future in as_completed(futures):
                video_file = futures[future]
                key = video_file.stem
                try:
                    caption = future.result()
                except Exception as e:
                    console.print(f"  [bold red]✗[/] {video_file.name}: [red]{e}[/]")
                    failed += 1
                    progress.advance(task)
                    continue

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

    console.print(
        f"\n[bold green]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}[/]."
    )


if __name__ == "__main__":
    app()