# ================================
# 0. PATCH pour huggingface_hub (contourne l'absence de HfFolder)
# ================================
import huggingface_hub
if not hasattr(huggingface_hub, 'HfFolder'):
    class HfFolder:
        _token = None
        @staticmethod
        def get_token():
            return HfFolder._token
        @staticmethod
        def save_token(token):
            HfFolder._token = token
    huggingface_hub.HfFolder = HfFolder

# ================================
# 1. PATCH pour contourner le bug de Gradio 4.44.0
#    (TypeError: argument of type 'bool' is not iterable)
# ================================
import gradio_client.utils

original_get_type = gradio_client.utils.get_type

def patched_get_type(schema):
    if isinstance(schema, bool):
        return "boolean"
    return original_get_type(schema)

gradio_client.utils.get_type = patched_get_type

# ================================
# 2. IMPORTS STANDARDS
# ================================
import gradio as gr
import numpy as np
from PIL import Image
from ultralytics import YOLO

# ================================
# 3. CHARGEMENT DU MODÈLE (YOLOv8 classification)
# ================================
MODEL_PATH = "best.pt"   # le modèle doit être présent à la racine
model = YOLO(MODEL_PATH)

# ================================
# 4. FONCTION DE PRÉDICTION
# ================================
def predict(image):
    """
    Prend une image PIL, retourne un dictionnaire {classe: probabilité}
    """
    # Redimensionner à 224x224 (taille utilisée lors de l'entraînement)
    image = image.resize((224, 224))
    # Convertir en numpy array
    img_np = np.array(image)
    # Faire la prédiction avec YOLO
    results = model.predict(img_np, verbose=False)
    probs = results[0].probs  # objet Probs
    # Les probabilités sont dans probs.data (tensor)
    prob_weed = float(probs.data[1].cpu().numpy())  # classe 1 = weed
    prob_crop = float(probs.data[0].cpu().numpy())  # classe 0 = crop
    return {"crop": prob_crop, "weed": prob_weed}

# ================================
# 5. INTERFACE GRADIO
# ================================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Chargez une image de champ"),
    outputs=gr.Label(num_top_classes=2, label="Prédiction"),
    title="🌾 Classification Culture / Mauvaise Herbe (YOLOv8)",
    description="Modèle YOLOv8-cls entraîné sur le dataset DeepWeeds pour distinguer les cultures des mauvaises herbes."
)

if __name__ == "__main__":
    iface.launch(share=True)