iamomtiwari's picture
Update app.py
4b6d380 verified
raw
history blame contribute delete
884 Bytes
import gradio as gr
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification
import torch
# Load the image processor and model
processor = ViTImageProcessor.from_pretrained('wambugu1738/crop_leaf_diseases_vit')
model = ViTForImageClassification.from_pretrained(
'wambugu1738/crop_leaf_diseases_vit',
ignore_mismatched_sizes=True
)
# Define a function to make predictions
def classify_image(image):
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
return model.config.id2label[predicted_class_idx]
# Create the Gradio interface
app = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="numpy"), # Corrected input type
outputs="text"
)
# Launch the Gradio app with a public link
app.launch(share=True)