vivym's picture
bugfix
6655549
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()