valmun's picture
Create app.py
3e8c6dd verified
raw history blame
No virus
2.23 kB
import gradio as gr
from transformers import ViTForImageClassification, ViTImageProcessor
import torch
# Define the pretrained model
model_name = "Treelar/vit-b16-plant_village"
# Load the ViT model and the image processor
model = ViTForImageClassification.from_pretrained(model_name)
image_processor = ViTImageProcessor.from_pretrained(model_name)
def predict(image):
# Convert the image into the model's required format
inputs = image_processor(images=image, return_tensors="pt")
# Disable gradient calculation to make the process efficient
with torch.no_grad():
outputs = model(**inputs) # Gets the model's output for the image
logits = outputs.logits # Output scores from the model
# Convert the logits to a probability using the softmax function
probability = torch.nn.functional.softmax(logits, dim=1)
top_probability, top_index = probability.max(1) # Gets the highest probability with its respective index
# Gets the disease label from the model using the probability's index
label_index = top_index.item()
label = model.config.id2label[label_index]
# Split the label into leaf category and disease name
label_parts = label.split("___")
leaf_category = label_parts[0]
disease_name = label_parts[1]
# Calculate the percentage breakdown of predicted diseases
percentage_breakdown = {disease: round(float(probability[0, index]) * 100, 2) for index, disease in enumerate(model.config.label2id)}
return leaf_category.capitalize(), disease_name.replace('_', ' ').capitalize(), percentage_breakdown
# Gradio interface setup with separate boxes for Leaf Type, Identified Disease, and Percentage Breakdown
interface = gr.Interface(
fn=predict,
inputs=gr.Image(label="Upload the Image"),
outputs=[
gr.Textbox(label="Leaf Type:"),
gr.Textbox(label="Identified Disease:"),
gr.JSON(label="Percentage Breakdown:")
],
title="Plant Disease Identifier",
description="Once the image has been uploaded and submitted, the disease of the plant will be determined. The percentage breakdown shows the probability of each disease."
)
interface.launch(debug=True) # Start server and launch the UI