import os
import sys
import torch
import gradio as gr
from PIL import Image

# ตั้งค่าให้รันบน CPU ของ Hugging Face ได้เสถียร
torch.cuda.is_available = lambda: False
torch.set_default_tensor_type(torch.BFloat16Tensor)

current_dir = os.getcwd()
sys.path.append(os.path.join(current_dir, "sam3"))

from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor

# โหลดโมเดลเตรียมไว้
model = build_sam3_image_model()
processor = Sam3Processor(model, device="cpu")

def predict(input_img, text_prompt):
    if input_img is None:
        return {"error": "No image provided"}
    
    image = input_img.convert("RGB")
    
    with torch.no_grad(), torch.autocast(device_type="cpu", dtype=torch.bfloat16):
        # ขั้นตอน 1: ฝังรูปเข้าไปในโมเดล
        inference_state = processor.set_image(image)
        # ขั้นตอน 2: ประมวลผลตาม Text Prompt ที่ส่งมา
        output = processor.set_text_prompt(
            state=inference_state, 
            prompt=text_prompt
        )
    
    # ส่งค่ากลับเป็นพิกัด (แปลงเป็น float เพื่อให้ส่งผ่านเน็ตได้)
    return {
        "masks": output["masks"].float().cpu().numpy().tolist(),
        "boxes": output["boxes"].float().cpu().numpy().tolist(),
        "scores": output["scores"].float().cpu().numpy().tolist()
    }

# สร้างหน้าเว็บ API
demo = gr.Interface(fn=predict, inputs=[gr.Image(type="pil"), "text"], outputs="json")
demo.launch()