Spaces:
Runtime error
Runtime error
| import os | |
| import gradio as gr | |
| from modules import scripts_postprocessing | |
| from modules.paths import models_path | |
| from modules import shared | |
| models = [ | |
| "None", | |
| "silueta", | |
| "u2net", | |
| "u2net_human_seg", | |
| "isnet-general-use", | |
| "isnet-anime", | |
| # "u2netp", | |
| # "u2net_cloth_seg", | |
| # "sam", | |
| ] | |
| def dependencies(): | |
| import installer | |
| for pkg in ['pymatting', 'pooch', 'rembg']: | |
| if not installer.installed(pkg): | |
| installer.install(pkg, ignore=False) | |
| class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): | |
| name = "Rembg" | |
| order = 20000 | |
| model = None | |
| def ui(self): | |
| with gr.Row(): | |
| model = gr.Dropdown(label="Model", choices=models, value="None", elem_id="extras_rembg_model") | |
| only_mask = gr.Checkbox(label="Return mask", value=False, elem_id="extras_rembg_only_mask") | |
| postprocess_mask = gr.Checkbox(label="Postprocess mask", value=False, elem_id="extras_rembg_process_mask") | |
| alpha_matting = gr.Checkbox(label="Alpha matting", value=True, elem_id="extras_rembg_alpha") | |
| with gr.Row(visible=True) as alpha_mask_row: | |
| alpha_matting_erode_size = gr.Slider(label="Erode size", minimum=0, maximum=40, step=1, value=10, elem_id="extras_rembg_alpha_erode") | |
| alpha_matting_foreground_threshold = gr.Slider(label="Foreground threshold", minimum=0, maximum=255, step=1, value=240, elem_id="extras_rembg_alpha_foreground") | |
| alpha_matting_background_threshold = gr.Slider(label="Background threshold", minimum=0, maximum=255, step=1, value=10, elem_id="extras_rembg_alpha_background") | |
| alpha_matting.change(fn=lambda x: gr.update(visible=x), inputs=[alpha_matting], outputs=[alpha_mask_row]) | |
| return { | |
| "model": model, | |
| "only_mask": only_mask, | |
| "postprocess_mask": postprocess_mask, | |
| "alpha_matting": alpha_matting, | |
| "alpha_matting_foreground_threshold": alpha_matting_foreground_threshold, | |
| "alpha_matting_background_threshold": alpha_matting_background_threshold, | |
| "alpha_matting_erode_size": alpha_matting_erode_size, | |
| } | |
| def process(self, pp: scripts_postprocessing.PostprocessedImage, model, only_mask, postprocess_mask, alpha_matting, alpha_matting_foreground_threshold, alpha_matting_background_threshold, alpha_matting_erode_size): # pylint: disable=arguments-differ | |
| dependencies() | |
| try: | |
| import rembg | |
| except Exception as e: | |
| shared.log.error(f'Rembg load failed: {e}') | |
| return | |
| if not model or model == "None": | |
| return | |
| if "U2NET_HOME" not in os.environ: | |
| os.environ["U2NET_HOME"] = os.path.join(models_path, "Rembg") | |
| args = { | |
| 'data': pp.image, | |
| 'only_mask': only_mask, | |
| 'post_process_mask': postprocess_mask, | |
| 'bgcolor': None, | |
| 'alpha_matting': alpha_matting, | |
| 'alpha_matting_foreground_threshold': alpha_matting_foreground_threshold, | |
| 'alpha_matting_background_threshold': alpha_matting_background_threshold, | |
| 'alpha_matting_erode_size': alpha_matting_erode_size, | |
| 'session': rembg.new_session(model), | |
| } | |
| """ | |
| # sam is more of a segment-anything model as it needs set of labels an points | |
| if model == 'sam': | |
| import numpy as np | |
| args['input_points'] = np.array([[400, 350], [700, 400], [200, 400]]) # array of y,x | |
| args['input_labels'] = np.array([1, 1, 2]) | |
| """ | |
| pp.image = rembg.remove(**args) | |
| pp.info["Rembg"] = model | |