Update app.py

app.py

@@ -54,13 +54,23 @@ def preprocess_image(image):
     
     image = cv2.resize(image, (224, 224))
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    return feature_extractor(images=image, return_tensors="pt")['pixel_values']
+    
+    # Use feature extractor with proper batching
+    inputs = feature_extractor(
+        images=[image],  # Pass as a list of images
+        return_tensors="pt",
+        padding=True  # Enable padding as a safety measure
+    )
+    return inputs['pixel_values']
 
-def predict(image):
-    image_tensor = preprocess_image(image)
+def predict_face_shape(image):
+    image_tensor = preprocess_image(image).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    
     with torch.no_grad():
         outputs = model(image_tensor)
-        return id2label[torch.argmax(outputs.logits).item()]
+        predicted_class_idx = torch.argmax(outputs.logits, dim=1).item()
+    
+    return id2label[predicted_class_idx]
 
 # Streamlit UI
 st.title("Face Shape & Glasses Recommender")