import torch
import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification


extractor = AutoFeatureExtractor.from_pretrained("susnato/plant_disease_detection-beans")
model = AutoModelForImageClassification.from_pretrained("susnato/plant_disease_detection-beans")

labels = ['angular leaf spot', 'rust', 'healthy']

def classify(im):
  features = extractor(im, return_tensors='pt')
  logits = model(features["pixel_values"])[-1]
  probability = torch.nn.functional.softmax(logits, dim=-1)
  probs = probability[0].detach().numpy()
  confidences = {label: float(probs[i]) for i, label in enumerate(labels)}
  return confidences


block = gr.Blocks(theme="JohnSmith9982/small_and_pretty")
with block:
    gr.HTML(
        """
        <h1 align="center">Hackefest 2024-DETECÇÃO DE DOENÇAS EM PLANTAS<h1>
        """
    )
    with gr.Group():
        with gr.Row():
            gr.HTML(
                """
                <p style="color:black">
                    <h4 style="font-color:powderblue;">
                        <center>A detecção de doenças de plantas é uma tarefa crucial na agricultura para garantir a saúde e o rendimento das culturas. Com os avanços na tecnologia, particularmente no campo da visão computacional e do aprendizado de máquina, sistemas automatizados foram desenvolvidos para identificar e diagnosticar doenças de plantas com precisão e eficiência. Esses sistemas normalmente envolvem a captura de imagens de plantas e sua análise usando algoritmos para detectar sintomas de doenças como descoloração, lesões ou padrões de crescimento anormais. Ao aproveitar essas tecnologias, os agricultores podem identificar e tratar prontamente as plantas doentes, minimizando assim a perda de colheitas e aumentando a produtividade agrícola.</center>
                    </h4>
                </p>
                
                <p align="center">
                    <img src="https://huggingface.co/datasets/susnato/stock_images/resolve/main/merged.png">
                </p>
                """
            )

    with gr.Group():
        with gr.Row():
            gr.HTML(
                """
                <center><h3>Our Approach</h3></center>
                
                <p align="center">
                    <img src="https://huggingface.co/datasets/susnato/stock_images/resolve/main/diagram2.png">
                </p>
                """
            )

    with gr.Group():
        image = gr.Image(type='pil')
        outputs = gr.Label()
        button = gr.Button("Classify")

        button.click(classify,
                     inputs=[image],
                     outputs=[outputs],
                     )

    with gr.Group():
        gr.Examples(["ex3.jpg"],
            fn=classify,
            inputs=[image],
            outputs=[outputs],
            cache_examples=True
        )

block.launch(debug=False, share=False)