Initial publish: SamGeo3 multi-prompt segmentation lab.
Scripts, prompt JSON tiers, usage docs, and README. Input images (data/) and segmentation outputs (output/) are gitignored.
This commit is contained in:
minsung
@@ -0,0 +1,189 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user