import gradio as gr
from ultralytics import YOLO
import tempfile

# Charger le modèle YOLOv8
model = YOLO("best.pt")

def detect_image(image):
    results = model(image)
    save_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
    results[0].plot(save_path)
    return save_path

def detect_video(video):
    # Traiter la vidéo
    results = model(video, save=True, project=tempfile.gettempdir(), name="yolov8_gradio")
    # YOLOv8 sauvegarde déjà la vidéo traitée, on renvoie son chemin
    return results[0].path

# Interface avec Blocks
with gr.Blocks() as demo:
    gr.Markdown("""# YOLOv8 Object Detection Demo  
    Upload une image, une vidéo, ou utilise ta caméra.  
    Bienvenue sur notre assistant de conduite routière **RoadEyes**.  
    RoadEyes est une IA de computer vision créée par **BlackBenAI**.  
    Il identifie en direct les panneaux de signalisation présents sur la route avec leur nom.  
    """)

    with gr.Tab("Image"):
        inp = gr.Image(type="pil")
        out = gr.Image(type="filepath")
        inp.change(fn=detect_image, inputs=inp, outputs=out)

    with gr.Tab("Vidéo"):
        inp_video = gr.Video()
        out_video = gr.Video()
        inp_video.change(fn=detect_video, inputs=inp_video, outputs=out_video)

    with gr.Tab("Caméra"):
        inp_cam = gr.Image(type="pil", sources=["webcam"])
        out_cam = gr.Image(type="filepath")
        inp_cam.change(fn=detect_image, inputs=inp_cam, outputs=out_cam)

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