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


@spaces.GPU(duration=120)
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()