File size: 4,507 Bytes
4a582ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6655549
4a582ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fc3685f
4a582ec
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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()