fix plot image with mask

app.py

@@ -118,7 +118,7 @@ def process_big_images(
     outputs = segment(Image.fromarray(img))
     msk = post_processing(outputs, target_size)
 
-    return Mask(msk)
+    return img, Mask(msk)
 
 
 def image_with_mask(
@@ -156,9 +156,9 @@ def categories_map(
 
 
 def main(image):
-    img = Image.fromarray(image)
+    image = Image.fromarray(image)
 
-    mask = process_big_images(img)
+    img, mask = process_big_images(image)
     mask_colorized = colorize(mask)
     fig = image_with_mask(img, mask)