add trimap generator

app.py

@@ -29,8 +29,9 @@ interface = Interface(pre_pipe=preprocessing,
                       seg_pipe=seg_net)
 
 
-def generate_trimap(original, mask):
-    trimap(original_image=original, mask=mask)
+def generate_trimap(original):
+    mask = seg_net([original])
+    return trimap(original_image=original, mask=mask[0])
 
 
 def predict(image):
@@ -48,22 +49,35 @@ Demo based on <a href='https://github.com/OPHoperHPO/image-background-remove-too
 """
 
 with gr.Blocks(title="CarveKit") as app:
-    gr.Markdown("<center><h1><b>CarveKit</b></h1></center>")
-    gr.HTML(
-        "<center><h3>Automated high-quality background removal framework for an image using neural networks.</h3></center>")
-    with gr.Row().style(equal_height=False):
-        with gr.Column():
-            input_img = gr.Image(type="pil", label="Input image")
-            run_btn = gr.Button(variant="primary")
-        with gr.Column():
-            output_img = gr.Image(type="pil", label="result")
-
-    run_btn.click(predict, [input_img], [output_img])
 
-    with gr.Row():
-        examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
-        examples = gr.Dataset(components=[input_img], samples=examples_data)
-        examples.click(lambda x: x[0], [examples], [input_img])
+    gr.Markdown("<center><h1><b>CarveKit</b></h1></center>")
+    gr.HTML("<center><h3>High-quality image background removal</h3></center>")
+
+    with gr.Tabs() as tabs:
+        with gr.TabItem("Remove background", id=0):
+            with gr.Row().style(equal_height=False):
+                with gr.Column():
+                    input_img = gr.Image(type="pil", label="Input image")
+                    run_btn = gr.Button(variant="primary")
+                with gr.Column():
+                    output_img = gr.Image(type="pil", label="result")
+
+            run_btn.click(predict, [input_img], [output_img])
+
+        with gr.TabItem("Generate trimap", id=1):
+            with gr.Row().style(equal_height=False):
+                with gr.Column():
+                    trimap_input = gr.Image(type="pil", label="Input image")
+                    trimap_btn = gr.Button(variant="primary")
+                with gr.Column():
+                    trimap_output = gr.Image(type="pil", label="result")
+
+            trimap_btn.click(generate_trimap, [trimap_input], [trimap_output])
+
+    # with gr.Row():
+    #     examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
+    #     examples = gr.Dataset(components=[input_img], samples=examples_data)
+    #     examples.click(lambda x: x[0], [examples], [input_img])
 
     with gr.Row():
         gr.HTML(footer)