Update app.py

app.py

@@ -25,7 +25,7 @@ def predict_segmentation(img):
     resized_img = cv2.resize(gray_img, im_size)
     # Predict the segmentation mask
     pred = learn.predict(resized_img)
-    scaled_pred = (pred[0] * 255).astype(np.uint8)
+    scaled_pred = (pred[0].numpy() * 255).astype(np.uint8)
     output_image = PILImage.create(scaled_pred)
     return output_image