minsung
34a885bcf8
Scripts, prompt JSON tiers, usage docs, and README. Input images (data/) and segmentation outputs (output/) are gitignored.
190 lines
6.3 KiB
Python
190 lines
6.3 KiB
Python
"""Single-image text-prompt segmentation with SamGeo3 (meta backend).
|
||
|
||
Usage (from project root, with venv active):
|
||
python scripts/text_segment.py
|
||
python scripts/text_segment.py --image data/test_image.jpg --prompt person
|
||
python scripts/text_segment.py --model-id facebook/sam3.1 --confidence 0.4
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# Project root = parent of scripts/
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
DEFAULT_DATA = ROOT / "data"
|
||
DEFAULT_OUTPUT = ROOT / "output"
|
||
SAMPLE_URL = (
|
||
"https://raw.githubusercontent.com/facebookresearch/sam3/"
|
||
"refs/heads/main/assets/images/test_image.jpg"
|
||
)
|
||
|
||
|
||
def ensure_sample_image(path: Path) -> Path:
|
||
if path.is_file():
|
||
return path
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
print(f"Downloading sample image -> {path}")
|
||
try:
|
||
from samgeo import download_file
|
||
|
||
download_file(SAMPLE_URL, str(path))
|
||
except Exception:
|
||
import urllib.request
|
||
|
||
urllib.request.urlretrieve(SAMPLE_URL, str(path))
|
||
if not path.is_file():
|
||
raise FileNotFoundError(f"failed to obtain image: {path}")
|
||
return path
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
p = argparse.ArgumentParser(description="SamGeo3 text segmentation smoke test")
|
||
p.add_argument("--image", type=Path, default=DEFAULT_DATA / "test_image.jpg")
|
||
p.add_argument("--prompt", default="person", help="Text prompt for grounding")
|
||
p.add_argument("--model-id", default="facebook/sam3.1")
|
||
p.add_argument("--backend", default="meta", choices=["meta", "transformers"])
|
||
# sam3.1 meta scores on this sample often peak ~0.3–0.4; 0.5 can yield zero masks
|
||
p.add_argument("--confidence", type=float, default=0.3)
|
||
p.add_argument("--mask-threshold", type=float, default=0.5)
|
||
p.add_argument("--resolution", type=int, default=1008)
|
||
p.add_argument("--min-size", type=int, default=0, help="Filter tiny masks (pixels)")
|
||
p.add_argument("--max-size", type=int, default=None)
|
||
p.add_argument("--device", default=None, help="cuda | cpu (default: auto)")
|
||
p.add_argument(
|
||
"--checkpoint",
|
||
default=os.environ.get("SAM3_CHECKPOINT_PATH"),
|
||
help="Local .pt path (or set SAM3_CHECKPOINT_PATH)",
|
||
)
|
||
p.add_argument(
|
||
"--output-dir",
|
||
type=Path,
|
||
default=DEFAULT_OUTPUT,
|
||
help="Directory for mask / annotation outputs",
|
||
)
|
||
p.add_argument(
|
||
"--no-viz",
|
||
action="store_true",
|
||
help="Skip matplotlib annotation PNG (mask file still saved)",
|
||
)
|
||
return p.parse_args()
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
|
||
if args.backend == "transformers" and "sam3.1" in args.model_id:
|
||
print(
|
||
"ERROR: facebook/sam3.1 requires backend='meta'. "
|
||
"Use model-id facebook/sam3 for transformers.",
|
||
file=sys.stderr,
|
||
)
|
||
return 2
|
||
|
||
image_path = ensure_sample_image(args.image.resolve())
|
||
out_dir = args.output_dir.resolve()
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
stem = image_path.stem
|
||
safe_prompt = "".join(c if c.isalnum() or c in "-_" else "_" for c in args.prompt)
|
||
mask_path = out_dir / f"{stem}_{safe_prompt}_mask.png"
|
||
ann_path = out_dir / f"{stem}_{safe_prompt}_ann.png"
|
||
scores_path = out_dir / f"{stem}_{safe_prompt}_scores.npy"
|
||
|
||
print("=== SamGeo3 text segmentation ===")
|
||
print(f"image: {image_path}")
|
||
print(f"prompt: {args.prompt}")
|
||
print(f"model_id: {args.model_id}")
|
||
print(f"backend: {args.backend}")
|
||
print(f"output_dir: {out_dir}")
|
||
|
||
import torch
|
||
from samgeo import SamGeo3
|
||
|
||
if args.device is None:
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
else:
|
||
device = args.device
|
||
|
||
print(f"device: {device} | cuda={torch.cuda.is_available()}")
|
||
if device == "cuda" and not torch.cuda.is_available():
|
||
print("ERROR: --device cuda but CUDA is not available", file=sys.stderr)
|
||
return 1
|
||
|
||
init_kwargs = dict(
|
||
backend=args.backend,
|
||
model_id=args.model_id,
|
||
device=device,
|
||
confidence_threshold=args.confidence,
|
||
mask_threshold=args.mask_threshold,
|
||
resolution=args.resolution,
|
||
enable_segmentation=True,
|
||
enable_inst_interactivity=False,
|
||
)
|
||
if args.checkpoint and os.path.isfile(args.checkpoint):
|
||
init_kwargs["checkpoint_path"] = args.checkpoint
|
||
init_kwargs["load_from_HF"] = False
|
||
print(f"checkpoint: {args.checkpoint}")
|
||
|
||
print("\nLoading model...")
|
||
sam = SamGeo3(**init_kwargs)
|
||
|
||
print("set_image...")
|
||
sam.set_image(str(image_path))
|
||
|
||
print(f'generate_masks("{args.prompt}")...')
|
||
gen_kwargs = {"min_size": args.min_size}
|
||
if args.max_size is not None:
|
||
gen_kwargs["max_size"] = args.max_size
|
||
sam.generate_masks(args.prompt, **gen_kwargs)
|
||
|
||
n = len(sam.masks) if getattr(sam, "masks", None) is not None else 0
|
||
if n == 0:
|
||
print("No masks found. Try another prompt or lower --confidence.")
|
||
return 0
|
||
|
||
scores = getattr(sam, "scores", None)
|
||
if scores is not None and len(scores):
|
||
try:
|
||
vals = [float(s.item() if hasattr(s, "item") else s) for s in scores]
|
||
print(f"scores (n={len(vals)}): min={min(vals):.3f} max={max(vals):.3f}")
|
||
except Exception:
|
||
pass
|
||
|
||
print(f"Saving masks -> {mask_path}")
|
||
# PNG cannot store float score maps; save mask first, scores as .npy
|
||
sam.save_masks(str(mask_path), unique=True)
|
||
|
||
scores = getattr(sam, "scores", None)
|
||
if scores is not None and len(scores):
|
||
import numpy as np
|
||
|
||
score_vals = np.array(
|
||
[float(s.item() if hasattr(s, "item") else s) for s in scores],
|
||
dtype=np.float32,
|
||
)
|
||
np.save(str(scores_path), score_vals)
|
||
print(f"Saved per-object scores -> {scores_path}")
|
||
|
||
if not args.no_viz:
|
||
try:
|
||
print(f"Saving annotations -> {ann_path}")
|
||
sam.show_anns(output=str(ann_path))
|
||
except Exception as e:
|
||
print(f"[WARN] show_anns failed: {e}")
|
||
|
||
print(f"\nDone. Found {n} object(s).")
|
||
print(f" mask: {mask_path}")
|
||
if scores_path.is_file():
|
||
print(f" scores: {scores_path}")
|
||
if ann_path.is_file():
|
||
print(f" ann: {ann_path}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|