scripts/infer_single_image.py

#!/usr/bin/env python3
"""Portable single-image inference for the bundled SAM3 breast-lesion model."""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

from PIL import Image


SCRIPT_DIR = Path(__file__).resolve().parent
BUNDLE_ROOT = SCRIPT_DIR.parent
for path in (SCRIPT_DIR, BUNDLE_ROOT / "runtime" / "sam3_repo"):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from sam3_buscot_runner import SAM3BuscotPredictor


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--prompt", default="breast lesion")
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--sam3_checkpoint", default=os.environ.get("SAM3_CHECKPOINT", str(BUNDLE_ROOT / "model" / "sam3_base.pt")))
    parser.add_argument("--lora_checkpoint", default=str(BUNDLE_ROOT / "model" / "best_model.pt"))
    args = parser.parse_args()

    predictor = SAM3BuscotPredictor(
        sam3_checkpoint=args.sam3_checkpoint,
        checkpoint_path=args.lora_checkpoint,
        prompt_type="semantic_text",
        prompt_text=args.prompt,
        encoder_trainable="lora",
        lora_rank=8,
        lora_alpha=16,
        threshold=args.threshold,
        device=args.device,
    )
    mask, details = predictor.predict(args.image)
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(mask * 255).save(output)
    print(f"Saved mask: {output}")
    print(details)


if __name__ == "__main__":
    main()