import torch
import torch.nn as nn
import gradio as gr
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel
from ultralytics.nn.modules.conv import Conv
from PIL import Image

# ---- FIX for PyTorch 2.6+ ----
torch.serialization.add_safe_globals([DetectionModel, nn.Sequential, Conv])

# ---- Load trained YOLO model ----
model = YOLO("best.pt")   # your junk food model
model.to("cpu")           # Required for Hugging Face Spaces

# ---- Prediction function ----
def predict(image):
    results = model.predict(image, conf=0.25)
    annotated = results[0].plot()   # BGR numpy image
    
    # Convert BGR → RGB
    annotated = annotated[:, :, ::-1]
    
    return Image.fromarray(annotated)

# ---- Gradio Interface ----
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Food Image"),
    outputs=gr.Image(type="pil", label="Detection Result"),
    title="🍔 Junk Food Detection (YOLO)",
    description="Upload an image to detect junk food items like Pizza, Burger, Ice Cream, Fries, etc."
)


if __name__ == "__main__":
    iface.launch()