Update app.py

app.py

@@ -25,7 +25,15 @@ def predict_segmentation(img):
     resized_img = cv2.resize(gray_img, im_size)
     # Predict the segmentation mask
     pred = learn.predict(resized_img)
-    output_image = PILImage.create(pred[0])
+    import matplotlib.pyplot as plt
+    from PIL import Image
+    fig, ax = plt.subplots()
+    ax.imshow(pred[0], cmap='gray')
+    ax.axis('off')
+    fig.canvas.draw()
+    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    output_image = Image.fromarray(data)
     return output_image
 
 input_image = gr.inputs.Image()