Update app.py

app.py

@@ -159,7 +159,7 @@ def predict(images, resolution, weights_file):
             save_paths.append(save_file_path)
 
         # Apply the prediction mask to the original image
-        pred = torch.nn.functional.interpolate(preds, size=image_shape, mode='bilinear', align_corners=True).numpy()
+        pred = torch.nn.functional.interpolate(preds, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
         image_pil = image_pil.resize(pred.shape[::-1])
         pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
         image_masked = (pred * np.array(image_pil)).astype(np.uint8)