Update app.py

app.py

@@ -10,22 +10,26 @@ model = AutoModelForImageClassification.from_pretrained(model_name)
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
 def classify_image(image):
+    # Ensure image is in RGB format (in case it's a PNG with alpha channel, for example)
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
     # 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
-    
+        logits = model(**inputs.pixel_values).logits  # Ensure to access pixel_values
+
     # 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 without specifying shape for the image input