import gradio as gr
from PIL import Image, ImageOps
from transformers import CLIPProcessor, CLIPModel
import torch


MODEL_ID = "EduFalcao/CropVision-CLIP"


processor = CLIPProcessor.from_pretrained(MODEL_ID)
model     = CLIPModel.from_pretrained(MODEL_ID)


HF_LABELS = [
    "Grape leaf with Black rot",
    "Grape leaf with Esca (Black Measles)",
    "Grape leaf with Leaf blight (Isariopsis Leaf Spot)",
    "Healthy Grape leaf"
]
MAP = {
    "Grape leaf with Black rot": "Black Rot",
    "Grape leaf with Esca (Black Measles)": "ESCA",
    "Grape leaf with Leaf blight (Isariopsis Leaf Spot)": "Leaf Blight",
    "Healthy Grape leaf": "Healthy"
}

def predict(image: Image.Image):

    img = ImageOps.exif_transpose(image).convert("RGB")
    img = img.resize((224,224))


    inputs = processor(text=HF_LABELS, images=img, return_tensors="pt", padding=True)
    outputs = model(**inputs)
    probs   = outputs.logits_per_image.softmax(dim=1)[0].tolist()


    mapping = { MAP[HF_LABELS[i]]: probs[i] for i in range(len(probs)) }
    best    = max(mapping, key=mapping.get)
    prob_lines = "\n".join(f"{cls}: {mapping[cls]:.2f}"
                           for cls in ["Healthy","Leaf Blight","Black Rot","ESCA"])
    return best, prob_lines


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Carrega uma folha"),
    outputs=[
        gr.Textbox(label="Classe prevista"),
        gr.Textbox(label="Probabilidades entre Classes")
    ],
    title="Modelo CropVision",
    description="Neste modelo vamos classificar folhas de vinhas em Healthy, Leaf Blight, Black Rot ou ESCA"
)

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