change

app.py

@@ -79,7 +79,10 @@ with gr.Blocks() as demo:
             color = np.concatenate([np.random.random(3), np.array([1.0])], axis=0)
             h, w = mask.shape[-2:]
             mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
-            mask_images.append(mask_image)
+            np_img = mask_image;
+            np_img = np.squeeze(np_img, axis=2)  # axis=2 is channel dimension 
+            pil_img = Image.fromarray(np_img, 'RGB')
+            mask_images.append(pil_img)
         
         return np.stack(mask_images)