Update app.py

app.py

@@ -28,6 +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)