import gradio as gr
import os
from roboflow import Roboflow
import cv2

# Configuration
api_key = os.getenv("ROBOFLOW_API_KEY")
rf = Roboflow(api_key=api_key)
project = rf.workspace().project("trash-detector-p2jn5-h9ffm")
model = project.version(3).model

def predict(image):
    # 1. Préparation de l'image
    cv2.imwrite("temp.jpg", cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
    
    # 2. Appel API
    response = model.predict("temp.jpg").json()
    
    # 3. Extraction chirurgicale selon ton JSON :
    # Le JSON est : {'predictions': [ { 'predictions': { 'Plastic': {...} } } ] }
    
    output = {}
    
    try:
        # On accède au premier élément de la liste principale
        data_principale = response['predictions'][0]
        
        # On accède au dictionnaire interne qui contient les classes
        details_classes = data_principale['predictions']
        
        for nom_classe, infos in details_classes.items():
            confiance = infos.get('confidence', 0)
            output[nom_classe] = float(confiance)
            
    except (KeyError, IndexError, TypeError):
        # En cas de problème, on affiche un message d'erreur propre
        return {"Erreur de lecture du JSON": 0.0}

    return output


# Interface
interface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(), 
    outputs=gr.Label(num_top_classes=3),
    title="Analyseur de Déchets",
    description="Déposez une photo pour identifier le type d'ordure."
)

interface.launch()