Update app.py

app.py

@@ -10,23 +10,32 @@ model = AutoModelForImageClassification.from_pretrained(model_name)
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
 def classify_image(image):
+    # Resize the input image to (224, 224)
+    image = image.resize((224, 224))
+    
     # Preprocess the image
     inputs = feature_extractor(images=image, return_tensors="pt")
+    
     # Make prediction
     with torch.no_grad():
         logits = model(**inputs).logits
+    
     # Retrieve the highest probability class label
     predicted_class_idx = logits.argmax(-1).item()
+    
     # Convert the index to the model's class label
     label = model.config.id2label[predicted_class_idx]
+    
     return label
 
-# Create Gradio interface
+# Create Gradio interface without specifying shape for the image input
+# This allows the input component to accept images of any size
 iface = gr.Interface(fn=classify_image, 
-                     inputs=gr.inputs.Image(shape=(224, 224)), 
-                     outputs="label",
+                     inputs=gr.Image(),  # Removed shape parameter
+                     outputs=gr.Label(),
                      title="Waste Classification with ViT",
                      description="Upload an image of waste, and the model will classify it.")
 
 # Launch the app
 iface.launch()
+