import gradio as gr
from ultralytics import YOLO
import cv2
import PIL.Image as Image
import numpy as np
import os
from huggingface_hub import hf_hub_download
import spaces

token = os.getenv("ACE_TOKEN")
repo_id = "LexBwmn/ACE_LAB"

try:
    model_path = hf_hub_download(
        repo_id=repo_id,
        filename="model.pt",
        token=token
    )
    model = YOLO(model_path)
except Exception as e:
    print(f"Error loading private model: {e}")
    model = None

@spaces.GPU
def predict(img):
    global model
    if model is None:
        return None
    
    try:
        img_array = np.array(img)
        results = model(img_array, conf=0.466, imgsz=640)     
        res_plotted = results[0].plot()
        res_rgb = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
        return Image.fromarray(res_rgb)
    except Exception as e:
        print(f"CRITICAL ERROR DURING SUBMIT: {e}")
        return None

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Brain MRI"), 
    outputs=gr.Image(type="pil", label="Detection Results"),
    title="ACE-V1.1 Lab Test",
    description="Secure Inference Test for Brain Tumor Detection."
)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1)
    demo.launch(ssr_mode=False)