| |
| """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() |
|
|
