Update app.py

app.py

@@ -39,9 +39,8 @@ def predict_segmentation(img):
     gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     resized_img = cv2.resize(gray_img, im_size)
     pred = learn.predict(resized_img)
-    color_pred = pred[0].show(ctx=None, cmap='gray', alpha=None)
-    color_pred_array = color_pred_array.astype(np.uint8)
-    output_image = Image.fromarray(color_pred_array)
+    scaled_pred = (pred[0].numpy() * 255).astype(np.uint8)
+    output_image = PILImage.create(scaled_pred)
     # Save the image to a temporary file
     temp_file = 'temp.png'
     output_image.save(temp_file)