Update app.py

app.py

@@ -3,18 +3,18 @@ import gradio as gr
 import datasets
 import torch
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-from transformers import ViTFeatureExtractor, ViTForImageClassification
+# from transformers import ViTFeatureExtractor, ViTForImageClassification
 
 dataset = datasets.load_dataset('beans')
 
 extractor = AutoFeatureExtractor.from_pretrained("suresh-subramanian/beans-classification")
 model = AutoModelForImageClassification.from_pretrained("suresh-subramanian/beans-classification")
-feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+# feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
 
 labels = dataset['train'].features['labels'].names
 
 def classify(im):
-  features = feature_extractor(im, return_tensors='pt')
+  features = extractor(im, return_tensors='pt')
   with torch.no_grad():
     logits = model(features["pixel_values"])[-1]
   probability = torch.nn.functional.softmax(logits, dim=-1)