VCardWizard / app.py
SarowarSaurav's picture
Update app.py
43c8e35 verified
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image
# Load the image processor and model
# Load model directly
from transformers import AutoImageProcessor, AutoModelForImageClassification
processor = AutoImageProcessor.from_pretrained("linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification")
model = AutoModelForImageClassification.from_pretrained("linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification")
# Define the function to process the image and make predictions
def classify_leaf_disease(image):
# Preprocess the image
inputs = processor(images=image, return_tensors="pt")
# Run the model on the image
with torch.no_grad():
outputs = model(**inputs)
# Get the predicted label and confidence score
probs = torch.softmax(outputs.logits, dim=1)
predicted_class_idx = probs.argmax(dim=1).item()
predicted_label = model.config.id2label[predicted_class_idx]
confidence_score = probs[0][predicted_class_idx].item()
# Format the output
return predicted_label, f"{confidence_score:.2f}", image
# Create Gradio Interface
interface = gr.Interface(
fn=classify_leaf_disease,
inputs=gr.Image(type="pil"),
outputs=[
gr.Textbox(label="Disease Name"),
gr.Textbox(label="Confidence Score"),
gr.Image(type="pil", label="Uploaded Image")
],
title="Leaf Disease Identification",
description="Upload an image of any plant leaf, and this model will identify the disease and show the confidence score."
)
# Launch the app
if __name__ == "__main__":
interface.launch()