|
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() |
|
|