src/longitrack_inference/cli.py

from __future__ import annotations

import argparse
from pathlib import Path


REPOSITORY_ROOT = Path(__file__).resolve().parents[2]


def parse_options() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run point-prompted LongiTrack tracking inference.")
    parser.add_argument("--images-path", type=Path, required=True, help="Input image directory.")
    parser.add_argument("--output-path", type=Path, required=True, help="Output directory.")
    parser.add_argument("--tracking-path", type=Path, required=True, help="Path to tracking.json.")
    parser.add_argument("--version", default="1.0", help="Model version to use for inference (default: 1.0).")
    parser.add_argument("--mode", choices=("automatic", "manual"), default="automatic", help="Tracking mode.")
    parser.add_argument("--device", choices=("cuda", "cpu", "mps"), default="cuda", help="Inference device.")
    return parser.parse_args()


def main() -> int:
    options = parse_options()
    model_path = REPOSITORY_ROOT / f"LongiTrack_v{options.version}"
    dataset_json_path = model_path / "dataset.json"
    if not model_path.is_dir() or not dataset_json_path.is_file():
        raise SystemExit(f"Unknown or incomplete model version: {options.version}")

    from torch import device
    from longiseg.inference.tracking_inference import predict

    predict(
        images_path=options.images_path,
        labels_path=None,
        output_path=options.output_path,
        model_path=model_path,
        tracking_path=options.tracking_path,
        dataset_json_path=dataset_json_path,
        folds=None,
        device=device(options.device),
        mode=options.mode,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())