dhayananthb's picture
Create app.py
0efb630 verified
raw
history blame contribute delete
602 Bytes
from PIL import Image, UnidentifiedImageError
from transformers import ViTFeatureExtractor, ViTForImageClassification
feature_extractor = ViTFeatureExtractor.from_pretrained('wambugu71/crop_leaf_diseases_vit')
model = ViTForImageClassification.from_pretrained(
'wambugu1738/crop_leaf_diseases_vit',
ignore_mismatched_sizes=True
)
image = Image.open('<image_path>')
inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])