Update app.py

app.py

@@ -25,13 +25,7 @@ def predict_segmentation(img):
     resized_img = cv2.resize(gray_img, im_size)
     # Predict the segmentation mask
     pred = learn.predict(resized_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)
+    output_image = pred[0]
     return output_image
 
 input_image = gr.inputs.Image()