Update app.py

app.py

@@ -1,5 +1,6 @@
 import cv2
 from fastai.vision.all import *
+import numpy as np
 import gradio as gr
 
 fnames = get_image_files("./albumentations/original")
@@ -27,7 +28,7 @@ def predict_segmentation(img):
     # Convert the predicted mask to a color image
     color_pred = pred[0].show(ctx=None, cmap='gray', alpha=None)
     # Convert the color image to a numpy array
-    color_pred_array = np.array(color_pred)
+    color_pred_array = color_pred_array.astype(np.uint8)
     # Convert the numpy array back to a PIL image
     output_image = Image.fromarray(color_pred_array)
     return output_image