scripts/run_detection.py

"""
run_detection.py — CLI entrypoint for the OWLv2 detection stage.

Usage:
    uv run python scripts/run_detection.py --image-dir data/raw --output-dir data/detections
    uv run python scripts/run_detection.py --help
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Optional

import click
from dotenv import load_dotenv

load_dotenv()  # picks up PYTORCH_ENABLE_MPS_FALLBACK and other vars

from autolabel.detect import run_detection
from autolabel.config import settings
from autolabel.utils import setup_logging


@click.command()
@click.option(
    "--image-dir",
    default=str(settings.raw_dir),
    show_default=True,
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    help="Directory containing input images.",
)
@click.option(
    "--output-dir",
    default=str(settings.detections_dir),
    show_default=True,
    type=click.Path(file_okay=False, path_type=Path),
    help="Directory to write per-image detection JSON files.",
)
@click.option(
    "--prompts",
    default=None,
    help="Comma-separated list of text prompts (overrides config defaults).",
)
@click.option(
    "--threshold",
    default=None,
    type=float,
    help="Score threshold override (0.0–1.0).",
)
@click.option(
    "--force",
    is_flag=True,
    default=False,
    help="Re-process images even if a detection JSON already exists.",
)
@click.option("--verbose", "-v", is_flag=True, default=False, help="Debug logging.")
def main(
    image_dir: Path,
    output_dir: Path,
    prompts: Optional[str],
    threshold: Optional[float],
    force: bool,
    verbose: bool,
) -> None:
    """Run OWLv2 open-vocabulary detection on IMAGE_DIR images."""
    setup_logging(logging.DEBUG if verbose else logging.INFO)

    prompt_list = None
    if prompts:
        prompt_list = [p.strip() for p in prompts.split(",") if p.strip()]

    if threshold is not None:
        settings.threshold = threshold

    run_detection(
        image_dir=image_dir,
        output_dir=output_dir,
        prompts=prompt_list,
        cfg=settings,
        force=force,
    )


if __name__ == "__main__":
    main()