import gradio as gr
from inference_example import AgribotInference

# 1) Charger le modèle UNE FOIS au démarrage
predictor = AgribotInference(model_path="agribot_models.pkl")

# 2) Fonction utilisée par l’UI Gradio
def infer(image):
    """
    image: PIL Image (Gradio gère ça pour toi)
    """
    # On adapte légèrement: le predict actuel prend un chemin de fichier,
    # alors on peut soit modifier la classe, soit créer une variante.
    # Ici, on réécrit une petite fonction dédiée.
    image = image.convert("RGB")
    img_tensor = predictor.transform(image).unsqueeze(0)
    features = predictor.extract_features(img_tensor)
    prediction_idx = predictor.model.predict(features)[0]
    probabilities = predictor.model.predict_proba(features)[0]
    confidence = float(probabilities.max() * 100)
    diagnosis = predictor.class_labels[prediction_idx]
    return f"{diagnosis} ({confidence:.2f}%)"

demo = gr.Interface(
    fn=infer,
    inputs=gr.Image(type="pil", label="Upload maize leaf image"),
    outputs=gr.Textbox(label="Diagnosis"),
    title="AgriBot - Maize Leaf Disease Diagnosis",
    description="Upload a maize leaf image to get a disease diagnosis."
)

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