Update app.py

app.py

@@ -25,15 +25,8 @@ def predict_segmentation(img):
     resized_img = cv2.resize(gray_img, im_size)
     # Predict the segmentation mask
     pred = learn.predict(resized_img)
-    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)
+    scaled_pred = (pred[0] * 255).astype(np.uint8)
+    output_image = PILImage.create(scaled_pred)
     return output_image
 
 input_image = gr.inputs.Image()