import gradio as gr
from ultralytics import YOLO
import os

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

def predict(image):
    # Faire l'inférence
    results = model.predict(source=image, conf=0.25)
    
    # Extraire les résultats pour un modèle de classification
    result = results[0]
    probs = result.probs
    
    # Créer un dictionnaire {Classe: Probabilité} pour Gradio
    # Ex: {"vide": 0.92, "plein": 0.08}
    labels = result.names
    confidences = {labels[i]: float(probs.data[i]) for i in range(len(labels))}
    
    return confidences

# Créer l'interface Gradio
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="NoThrowam - Détection de remplissage",
    description="Uploadez une photo de poubelle pour connaître son niveau de plénitude."
)

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