init app

.gitignore

@@ -0,0 +1 @@
+.idea/

app.py

@@ -0,0 +1,63 @@
+import gradio as gr
+import torch
+from carvekit.api.interface import Interface
+from carvekit.ml.wrap.fba_matting import FBAMatting
+from carvekit.ml.wrap.tracer_b7 import TracerUniversalB7
+from carvekit.pipelines.postprocessing import MattingMethod
+from carvekit.pipelines.preprocessing import PreprocessingStub
+from carvekit.trimap.generator import TrimapGenerator
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Check doc strings for more information
+seg_net = TracerUniversalB7(device=device, batch_size=1)
+
+fba = FBAMatting(device=device,
+                 input_tensor_size=2048,
+                 batch_size=1)
+
+trimap = TrimapGenerator()
+
+preprocessing = PreprocessingStub()
+
+postprocessing = MattingMethod(matting_module=fba,
+                               trimap_generator=trimap,
+                               device=device)
+
+interface = Interface(pre_pipe=preprocessing,
+                      post_pipe=postprocessing,
+                      seg_pipe=seg_net)
+
+
+def predict(image):
+    return interface([image])[0]
+
+
+footer = r"""
+<center>
+<b>
+Demo based on <a href='https://github.com/OPHoperHPO/image-background-remove-tool'>CarveKit</a>
+</b>
+</center>
+"""
+
+with gr.Blocks(title="CarveKit") as app:
+    gr.HTML("<center><h1>Image Remove Background</h1></center>")
+    with gr.Row():
+        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])
+
+    with gr.Row():
+        gr.HTML(footer)
+
+app.launch(share=False, debug=True, enable_queue=True, show_error=True)

requirements.txt

@@ -0,0 +1,4 @@
+pillow~=9.5.0
+torch>=2.0.1
+gradio~=3.35.2
+carvekit~=4.1.0