from hashlib import sha1 from pathlib import Path import cv2 import gradio as gr import numpy as np from PIL import Image from paddleseg.cvlibs import manager, Config from paddleseg.utils import load_entire_model manager.BACKBONES._components_dict.clear() manager.TRANSFORMS._components_dict.clear() import ppmatting as ppmatting from ppmatting.core import predict from ppmatting.utils import estimate_foreground_ml model_names = [ "modnet-mobilenetv2", "ppmatting-512", "ppmatting-1024", "ppmatting-2048", "modnet-hrnet_w18", "modnet-resnet50_vd", ] model_dict = { name: None for name in model_names } last_result = { "cache_key": None, "algorithm": None, } def image_matting( image: np.ndarray, result_type: str, bg_color: str, algorithm: str, morph_op: str, morph_op_factor: float, ) -> np.ndarray: image = np.ascontiguousarray(image) cache_key = sha1(image).hexdigest() if cache_key == last_result["cache_key"] and algorithm == last_result["algorithm"]: alpha = last_result["alpha"] else: cfg = Config(f"configs/{algorithm}.yml") if model_dict[algorithm] is not None: model = model_dict[algorithm] else: model = cfg.model load_entire_model(model, f"models/{algorithm}.pdparams") model.eval() model_dict[algorithm] = model transforms = ppmatting.transforms.Compose(cfg.val_transforms) alpha = predict( model, transforms=transforms, image=image, ) last_result["cache_key"] = cache_key last_result["algorithm"] = algorithm last_result["alpha"] = alpha alpha = (alpha * 255).astype(np.uint8) kernel = np.ones((5, 5), np.uint8) if morph_op == "dilate": alpha = cv2.dilate(alpha, kernel, iterations=int(morph_op_factor)) else: alpha = cv2.erode(alpha, kernel, iterations=int(morph_op_factor)) alpha = (alpha / 255).astype(np.float32) image = (image / 255.0).astype("float32") fg = estimate_foreground_ml(image, alpha) if result_type == "Remove BG": result = np.concatenate((fg, alpha[:, :, None]), axis=-1) elif result_type == "Replace BG": bg_r = int(bg_color[1:3], base=16) bg_g = int(bg_color[3:5], base=16) bg_b = int(bg_color[5:7], base=16) bg = np.zeros_like(fg) bg[:, :, 0] = bg_r / 255. bg[:, :, 1] = bg_g / 255. bg[:, :, 2] = bg_b / 255. result = alpha[:, :, None] * fg + (1 - alpha[:, :, None]) * bg result = np.clip(result, 0, 1) else: result = alpha return result def main(): with gr.Blocks() as app: gr.Markdown("Image Matting Powered By AI") with gr.Row(variant="panel"): image_input = gr.Image() image_output = gr.Image() with gr.Row(variant="panel"): result_type = gr.Radio( label="Mode", show_label=True, choices=[ "Remove BG", "Replace BG", "Generate Mask", ], value="Remove BG", ) bg_color = gr.ColorPicker( label="BG Color", show_label=True, value="#000000", ) algorithm = gr.Dropdown( label="Algorithm", show_label=True, choices=model_names, value="modnet-hrnet_w18" ) with gr.Row(variant="panel"): morph_op = gr.Radio( label="Post-process", show_label=True, choices=[ "Dilate", "Erode", ], value="Dilate", ) morph_op_factor = gr.Slider( label="Factor", show_label=True, minimum=0, maximum=20, value=0, step=1, ) run_button = gr.Button("Run") run_button.click( image_matting, inputs=[ image_input, result_type, bg_color, algorithm, morph_op, morph_op_factor, ], outputs=image_output, ) app.launch() if __name__ == "__main__": main()