Update app.py

app.py

@@ -1,20 +1,21 @@
 import gradio as gr
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
 import torch
+import numpy as np
 
 # Load the pre-trained model and preprocessor (feature extractor)
 model_name = "jjuarez/Vit_waste_image_class"
-model = AutoModelForImageClassification.from_pretrained(model_name)
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
 
 def classify_image(image):
-    # Convert PIL Image to NumPy array if not already done by Gradio
+    # Convert the PIL Image to a format compatible with the feature extractor
     image = np.array(image)
 
-    # Let the feature extractor handle resizing and normalization
+    # Preprocess the image and prepare it for the model
     inputs = feature_extractor(images=image, return_tensors="pt")
-    
+
     # Make prediction
     with torch.no_grad():
         outputs = model(**inputs)
@@ -28,9 +29,9 @@ def classify_image(image):
 
     return label
 
-# Create Gradio interface without specifying shape for the image input
-iface = gr.Interface(fn=classify_image,
-                     inputs=gr.Image(),  # Removed shape parameter, accepts image of any size
+# Create Gradio interface
+iface = gr.Interface(fn=classify_image, 
+                     inputs=gr.Image(),  # Accepts image of any size
                      outputs=gr.Label(),
                      title="Waste Classification with ViT",
                      description="Upload an image of waste, and the model will classify it.")