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-PLANT DISEASE DETECTION<h1>
        """
    )
    with gr.Group():
        with gr.Row():
            gr.HTML(
                """
                <p style="color:black">
                    <h4 style="font-color:powderblue;">
                        <center>Plant disease detection is a crucial task in agriculture to ensure the health and yield of crops. With advancements in technology, particularly in the field of computer vision and machine learning, automated systems have been developed to identify and diagnose plant diseases accurately and efficiently. These systems typically involve capturing images of plants and analyzing them using algorithms to detect symptoms of diseases such as discoloration, lesions, or abnormal growth patterns. By leveraging such technologies, farmers can promptly identify and treat diseased plants, thereby minimizing crop loss and increasing agricultural productivity</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)