import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
import gradio as gr

# ── Charger le modèle ──
device = torch.device("cpu")
model = models.efficientnet_b0()
model.classifier = nn.Sequential(nn.Dropout(0.2), nn.Linear(1280, 2))
model.load_state_dict(torch.load("efficientnet_mines_V1.pth", map_location=device))
model.to(device)
model.eval()

classes = ['extractions_illegales', 'pas_extraction_illegale']

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

def predire(image):
    image = Image.fromarray(image).convert("RGB")
    tenseur = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(tenseur)
        proba = torch.softmax(output, dim=1)
        prediction = proba.argmax(1).item()
        confiance = proba[0][prediction].item() * 100

    if confiance < 70:
        return {
            "Classe": "INDETERMINE",
            "Confiance": f"{confiance:.1f}%",
            "Detail": "Image non reconnue. Veuillez uploader une image de site minier."
        }

    classe = classes[prediction]
    resultat = "ILLEGALE" if "illegale" in classe.lower() else "LEGALE"

    return {
        "Classe": resultat,
        "Confiance": f"{confiance:.1f}%",
        "Detail": classe
    }

interface = gr.Interface(
    fn=predire,
    inputs=gr.Image(label="Uploader une image de mine"),
    outputs=gr.JSON(label="Résultat"),
    title="Detection de Mines Illegales",
    description="Uploadez une image pour savoir si c'est une extraction illegale ou non."
)

interface.launch()