Update

app.py

@@ -70,11 +70,11 @@ def binarize_mask(mask: np.ndarray) -> np.ndarray:
     return mask
 
 
-def update_trimap(foreground_mask_editor: dict, unknown_mask_editor: dict) -> np.ndarray:
-    foreground = foreground_mask_editor["layers"][0]
+def update_trimap(mask_editor: dict) -> np.ndarray:
+    foreground = mask_editor["layers"][0]
     foreground = binarize_mask(foreground)
 
-    unknown = unknown_mask_editor["layers"][0]
+    unknown = mask_editor["layers"][1]
     unknown = binarize_mask(unknown)
 
     trimap = np.zeros_like(foreground)
@@ -161,25 +161,17 @@ with gr.Blocks(css_paths="style.css") as demo:
                     with gr.Tab(label="Trimap"):
                         trimap = gr.Image(label="Trimap", type="pil", image_mode="L")
                     with gr.Tab(label="Draw trimap"):
-                        foreground_mask = gr.ImageEditor(
-                            label="Foreground",
+                        mask_editor = gr.ImageEditor(
+                            label="Masks",
                             type="numpy",
                             sources=("upload",),
                             transforms=(),
                             image_mode="L",
                             height=500,
                             brush=gr.Brush(default_color=("#00ff00", 0.6)),
-                            layers=gr.LayerOptions(allow_additional_layers=False, layers=["Foreground mask"]),
-                        )
-                        unknown_mask = gr.ImageEditor(
-                            label="Unknown",
-                            type="numpy",
-                            sources=("upload",),
-                            transforms=(),
-                            image_mode="L",
-                            height=500,
-                            brush=gr.Brush(default_color=("#00ff00", 0.6)),
-                            layers=gr.LayerOptions(allow_additional_layers=False, layers=["Unknown mask"]),
+                            layers=gr.LayerOptions(
+                                allow_additional_layers=False, layers=["Foreground mask", "Unknown mask"]
+                            ),
                         )
                         generate_trimap_button = gr.Button("Generate trimap")
                 apply_background_replacement = gr.Checkbox(label="Replace background", value=False)
@@ -219,14 +211,14 @@ with gr.Blocks(css_paths="style.css") as demo:
         outputs=image,
         api_name=False,
     ).then(
-        fn=lambda image: (image, image),
+        fn=lambda image: image,
         inputs=image,
-        outputs=[foreground_mask, unknown_mask],
+        outputs=mask_editor,
         api_name=False,
     )
     generate_trimap_button.click(
         fn=update_trimap,
-        inputs=[foreground_mask, unknown_mask],
+        inputs=mask_editor,
         outputs=trimap,
         api_name=False,
     )