Add controls to Gradio interface (#581)

Co-authored-by: Zebradots <[email protected]>
danielgatis · Jan 31, 2024 · cf962f5 · cf962f5
1 parent 49d1686
commit cf962f5
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/rembg/commands/s_command.py b/rembg/commands/s_command.py
@@ -264,7 +264,15 @@ def gr_app(app):
         def inference(input_path, model, cmd_args):
             output_path = "output.png"
 
-            kwargs = {}
+            kwargs = {
+                "alpha_matting": a,
+                "alpha_matting_foreground_threshold": af,
+                "alpha_matting_background_threshold": ab,
+                "alpha_matting_erode_size": ae,
+                "only_mask": om,
+                "post_process_mask": ppm,
+            }
+
             if cmd_args:
                 kwargs.update(json.loads(cmd_args))
             kwargs["session"] = new_session(model, **kwargs)
@@ -281,6 +289,12 @@ def inference(input_path, model, cmd_args):
             [
                 gr.components.Image(type="filepath", label="Input"),
                 gr.components.Dropdown(sessions_names, value="u2net", label="Models"),
+                gr.components.Checkbox(value=True, label="Alpha matting"),
+                gr.components.Slider(value=240, minimum=0, maximum=255, label="Foreground threshold"),
+                gr.components.Slider(value=10, minimum=0, maximum=255, label="Background threshold"),
+                gr.components.Slider(value=40, minimum=0, maximum=255, label="Erosion size"),
+                gr.components.Checkbox(value=False, label="Only mask"),
+                gr.components.Checkbox(value=True, label="Post process mask"),
                 gr.components.Textbox(label="Arguments"),
             ],
             gr.components.Image(type="filepath", label="Output"),