Spaces:
Runtime error
Runtime error
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 | |