File size: 2,860 Bytes
d18e0a7 2b27efc d18e0a7 2b27efc d18e0a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model card's Python snippet on the demo image.
pip install -e code/
python examples/quickstart.py [image] [prompt ...]
Writes examples/out/quickstart_overlay.png (masks and boxes on the image) and examples/out/quickstart.json.
The default image and the output directory are found relative to this file, so it runs from any directory; an image
path given on the command line is relative to the current directory.
"""
import json
import os
import sys
import time
import numpy as np
from PIL import Image, ImageDraw
from tt_sam3 import Sam3
HERE = os.path.dirname(os.path.abspath(__file__))
image_path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "..", "media", "source_1.png")
prompts = sys.argv[2:] or ["car"]
out_dir = os.path.join(HERE, "out")
os.makedirs(out_dir, exist_ok=True)
# --- the model card snippet -----------------------------------------------------------------------------
with Sam3.from_pretrained(device_id=0, warmup_variants={"prompts": prompts}) as model: # warm: call 1 is fast too
results = model(image_path, prompts) # one Sam3Result per prompt
t0 = time.perf_counter()
model(image_path, prompts) # warm call (same image: the image part is reused)
warm_ms = (time.perf_counter() - t0) * 1e3
# ---------------------------------------------------------------------------------------------------------
for r in results:
print(r)
print(f"warm call (same image): {warm_ms:.1f} ms")
# overlay: one color per instance, mask blended at 50 %, box outline, score label
image = Image.open(image_path).convert("RGB")
canvas = np.asarray(image).astype(np.float32)
palette = np.array([[230, 25, 75], [60, 180, 75], [255, 225, 25], [0, 130, 200], [245, 130, 48], [145, 30, 180],
[70, 240, 240], [240, 50, 230], [210, 245, 60], [250, 190, 212]], dtype=np.float32)
k = 0
for r in results:
for m in r.masks:
canvas[m] = 0.5 * canvas[m] + 0.5 * palette[k % len(palette)]
k += 1
overlay = Image.fromarray(canvas.astype(np.uint8))
draw = ImageDraw.Draw(overlay)
k = 0
for r in results:
for box, score in zip(r.boxes, r.scores):
color = tuple(int(c) for c in palette[k % len(palette)])
draw.rectangle([float(v) for v in box], outline=color, width=3)
draw.text((float(box[0]) + 4, float(box[1]) + 4), f"{r.prompt} {score:.2f}", fill=color)
k += 1
overlay.save(os.path.join(out_dir, "quickstart_overlay.png"))
with open(os.path.join(out_dir, "quickstart.json"), "w") as f:
json.dump({"image": os.path.relpath(image_path, os.path.join(HERE, "..")), "warm_call_ms": round(warm_ms, 1),
"results": [r.to_dict() for r in results]}, f, indent=1)
print("wrote", os.path.join(out_dir, "quickstart_overlay.png"), "and quickstart.json")
|