Преглед изворни кода

Add controls to Gradio interface (#581)

Co-authored-by: Zebradots <[email protected]>
zebradots пре 1 година
родитељ
комит
cf962f5429
1 измењених фајлова са 15 додато и 1 уклоњено
  1. 15 1
      rembg/commands/s_command.py

+ 15 - 1
rembg/commands/s_command.py

@@ -264,7 +264,15 @@ def s_command(port: int, host: str, log_level: str, threads: int) -> None:
         def inference(input_path, model, cmd_args):
         def inference(input_path, model, cmd_args):
             output_path = "output.png"
             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:
             if cmd_args:
                 kwargs.update(json.loads(cmd_args))
                 kwargs.update(json.loads(cmd_args))
             kwargs["session"] = new_session(model, **kwargs)
             kwargs["session"] = new_session(model, **kwargs)
@@ -281,6 +289,12 @@ def s_command(port: int, host: str, log_level: str, threads: int) -> None:
             [
             [
                 gr.components.Image(type="filepath", label="Input"),
                 gr.components.Image(type="filepath", label="Input"),
                 gr.components.Dropdown(sessions_names, value="u2net", label="Models"),
                 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.Textbox(label="Arguments"),
             ],
             ],
             gr.components.Image(type="filepath", label="Output"),
             gr.components.Image(type="filepath", label="Output"),