Update app.py

app.py

@@ -13,7 +13,7 @@ mp_face_detection = mp.solutions.face_detection.FaceDetection(
 # Initialize model and labels
 @st.cache_resource
 def load_model():
-    id2label = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}
+    id2label = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}  # Moved inside load_model
     label2id = {v: k for k, v in id2label.items()}
     
     model = SwinForImageClassification.from_pretrained(
@@ -25,9 +25,10 @@ def load_model():
     
     model.load_state_dict(torch.load('swin.pth', map_location='cpu'))
     model.eval()
-    return model, AutoFeatureExtractor.from_pretrained("microsoft/swin-tiny-patch4-window7-224")
+    return model, AutoFeatureExtractor.from_pretrained("microsoft/swin-tiny-patch4-window7-224"), id2label  # Return id2label
 
-model, feature_extractor = load_model()
+# Load model components and labels
+model, feature_extractor, id2label = load_model()  # Receive id2label here
 
 glasses_recommendations = {
     "Heart": "Rimless (tanpa bingkai bawah)",
@@ -44,11 +45,15 @@ def preprocess_image(image):
         detection = results.detections[0]
         bbox = detection.location_data.relative_bounding_box
         h, w, _ = image.shape
-        # In the face cropping section
         x1 = max(0, int(bbox.xmin * w))
         y1 = max(0, int(bbox.ymin * h))
         x2 = min(w, int((bbox.xmin + bbox.width) * w))
         y2 = min(h, int((bbox.ymin + bbox.height) * h))
+        
+        # Add validation check
+        if (x2 <= x1) or (y2 <= y1) or (x2 - x1 < 10) or (y2 - y1 < 10):
+            raise ValueError("Invalid face crop dimensions")
+            
         image = image[y1:y2, x1:x2]
     else:
         raise ValueError("No face detected")
@@ -56,16 +61,17 @@ def preprocess_image(image):
     image = cv2.resize(image, (224, 224))
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
-    # Use feature extractor with proper batching
     inputs = feature_extractor(
-        images=[image],  # Pass as a list of images
+        images=[image],
         return_tensors="pt",
-        padding=True  # Enable padding as a safety measure
+        padding=True
     )
     return inputs['pixel_values']
 
 def predict_face_shape(image):
-    image_tensor = preprocess_image(image).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    # Force CPU usage on Hugging Face Spaces
+    device = torch.device("cpu")
+    image_tensor = preprocess_image(image).to(device)
     
     with torch.no_grad():
         outputs = model(image_tensor)
@@ -87,6 +93,7 @@ if uploaded_file is not None:
     
     try:
         with st.spinner('Analyzing...'):
+            # Convert PIL image to OpenCV format correctly
             prediction = predict_face_shape(cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))
             recommendation = glasses_recommendations[prediction]