Spaces:
Running on Zero
Running on Zero
| """SAM 3 panoptic concept-segmentation API (ZeroGPU). Self-contained.""" | |
| import base64 | |
| import io | |
| import os | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import torch | |
| from PIL import Image | |
| from transformers import Sam3Model, Sam3Processor | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| MODEL_ID = "facebook/sam3" | |
| # Built at import on CPU; moved to CUDA inside the @spaces.GPU function. | |
| processor = Sam3Processor.from_pretrained(MODEL_ID, token=HF_TOKEN) | |
| model = Sam3Model.from_pretrained(MODEL_ID, token=HF_TOKEN) | |
| model.eval() | |
| def _encode_mask(mask_bool: np.ndarray) -> str: | |
| arr = (mask_bool.astype(np.uint8)) * 255 | |
| buf = io.BytesIO() | |
| Image.fromarray(arr, mode="L").save(buf, format="PNG") | |
| return base64.b64encode(buf.getvalue()).decode("ascii") | |
| def api_panoptic(image, concepts, conf, mask_threshold=0.5): | |
| if image is None: | |
| return {"error": "no image provided"} | |
| image = image.convert("RGB") | |
| W, H = image.size | |
| concept_list = [c.strip() for c in (concepts or "").split(",") if c.strip()] | |
| device = "cuda" | |
| model.to(device) | |
| detections = [] | |
| for concept in concept_list: | |
| inputs = processor(images=image, text=concept, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| target_sizes = (inputs["original_sizes"].tolist() | |
| if "original_sizes" in inputs else [[H, W]]) | |
| res = processor.post_process_instance_segmentation( | |
| outputs, threshold=float(conf), mask_threshold=float(mask_threshold), | |
| target_sizes=target_sizes)[0] | |
| # NOTE (verify on live Space): expected keys masks/scores/boxes. | |
| masks, scores = res["masks"], res["scores"] | |
| boxes = res.get("boxes") | |
| for i in range(len(scores)): | |
| m = masks[i] | |
| m = m.cpu().numpy() if hasattr(m, "cpu") else np.asarray(m) | |
| mb = m > 0.5 if m.dtype != bool else m | |
| box = (boxes[i].cpu().numpy().tolist() | |
| if boxes is not None else [0, 0, 0, 0]) | |
| detections.append({ | |
| "label": concept, "score": float(scores[i]), | |
| "box": box, "mask_png_b64": _encode_mask(mb.astype(bool)), | |
| }) | |
| return {"version": "3", "model": MODEL_ID, "width": W, "height": H, | |
| "detections": detections} | |
| with gr.Blocks(title="SAM3 Panoptic") as demo: | |
| gr.Markdown("# SAM 3 Panoptic API\nUpload an image, enter comma-separated concepts.") | |
| with gr.Row(): | |
| inp = gr.Image(type="pil", label="Image") | |
| out = gr.JSON(label="Detections") | |
| txt = gr.Textbox(label="Concepts (comma-separated)", | |
| value="person, car, road, sky, building, tree") | |
| conf = gr.Slider(0.0, 1.0, value=0.4, step=0.05, label="Confidence (lower = more detail)") | |
| mthr = gr.Slider(0.05, 0.95, value=0.5, step=0.05, label="Mask threshold") | |
| gr.Button("Segment").click(api_panoptic, [inp, txt, conf, mthr], out, | |
| api_name="api_panoptic") | |
| if __name__ == "__main__": | |
| demo.queue().launch() | |