Update app.py

app.py

@@ -29,6 +29,11 @@ def segment_image(image):
     person_class_id = 15  # Adjust based on the model's class mapping
     mask = (segmentation == person_class_id).astype(np.uint8)
 
+    # Resize mask to match the input image size
+    mask = Image.fromarray(mask)
+    mask = mask.resize(image.size, resample=Image.NEAREST)
+    mask = np.array(mask)
+
     return mask
 
 # Apply style transfer to different layers