Update app.py

app.py

@@ -37,17 +37,21 @@ def segment_image(image):
     outputs = model(**inputs)
     logits = outputs.logits
     segmentation = torch.argmax(logits, dim=1).squeeze().cpu().numpy()
+
+    # Resize segmentation mask to match original image size
+    segmentation_resized = cv2.resize(segmentation, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
     
-    # Create a blank mask image
+    # Create a blank mask image with the same size as the original image
     segmented_image = np.zeros_like(image)
     
     # Color each part with unique colors
     for part, color in PART_COLORS.items():
-        mask = np.isin(segmentation, PART_LABELS[part])
+        mask = np.isin(segmentation_resized, PART_LABELS[part])
         segmented_image[mask] = color
     
     return segmented_image
 
+
 def estimate_pose(image):
     # Convert image from BGR (OpenCV) to RGB
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)