Update app.py

app.py

@@ -55,22 +55,28 @@ def preprocess_image(image_path):
     # Get original shape
     original_shape = img.shape
     
-    # Resize to model input size (assuming 256x256)
-    img_resized = cv2.resize(img, (256, 256))
+    # Resize to model input size: 160x160 (NOT 256x256!)
+    img_resized = cv2.resize(img, (160, 160))
     
     # Normalize to [0, 1]
     img_normalized = img_resized.astype(np.float32) / 255.0
     
-    # Add batch and channel dimensions
-    img_input = np.expand_dims(img_normalized, axis=0)
-    img_input = np.expand_dims(img_input, axis=-1)
+    # Add batch and channel dimensions in channels_first format (NCHW)
+    # Model expects: (batch, channels, height, width) = (None, 1, 160, 160)
+    img_input = np.expand_dims(img_normalized, axis=0)  # (1, 160, 160)
+    img_input = np.expand_dims(img_input, axis=0)  # (1, 1, 160, 160)
     
     return img_input, original_shape
 
 def postprocess_mask(mask, original_shape):
     """Postprocess segmentation mask"""
-    # Remove batch dimension
-    mask = np.squeeze(mask)
+    # Mask comes in channels_first format: (batch, channels, height, width)
+    # Squeeze to remove batch and channel dimensions
+    mask = np.squeeze(mask)  # (160, 160)
+    
+    # If mask still has extra dimensions, squeeze again
+    while len(mask.shape) > 2:
+        mask = np.squeeze(mask)
     
     # Resize back to original shape
     mask_resized = cv2.resize(mask, (original_shape[1], original_shape[0]))