import gradio as gr
import json
from PIL import Image
from tensorflow.keras.models import load_model
from helpers import detect_image

model = load_model("yolo_model.keras", compile=False)
with open("anchors.json") as f: anchors = json.load(f)
with open("labels.json") as f: labels = json.load(f)

def predict(image):
    if image is None: return None
    return detect_image(image.convert("RGB"), model, anchors, labels)

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload image"),
    outputs=gr.Image(type="pil", label="Detections"),
    title="YOLO Object Detection",
    allow_flagging="never",
)

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