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)
-    output_image = pred[0]
+    output_image = PILImage.create(pred[0])
     return output_image
 
 input_image = gr.inputs.Image()