leonelhs commited on
Commit
7edd531
1 Parent(s): aa50dd4

upload examples

Browse files
Files changed (5) hide show
  1. .gitignore +3 -1
  2. examples/01.jpg +0 -0
  3. examples/02.jpg +0 -0
  4. examples/03.jpg +0 -0
  5. playground.py +184 -0
.gitignore CHANGED
@@ -146,9 +146,11 @@ cython_debug/
146
 
147
  # PyCharm
148
  .idea/
149
- examples/
150
  task_make
151
  task_upload
152
  setup.py
 
 
 
153
 
154
 
 
146
 
147
  # PyCharm
148
  .idea/
 
149
  task_make
150
  task_upload
151
  setup.py
152
+ playground.py
153
+
154
+
155
 
156
 
examples/01.jpg ADDED
examples/02.jpg ADDED
examples/03.jpg ADDED
playground.py ADDED
@@ -0,0 +1,184 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import PIL.Image
2
+ import cv2
3
+ import gradio as gr
4
+ import huggingface_hub
5
+ import numpy as np
6
+ import onnxruntime as rt
7
+ from PIL import ImageOps
8
+ from carvekit.trimap.generator import TrimapGenerator
9
+ from pymatting import estimate_alpha_cf, estimate_foreground_ml, stack_images, load_image
10
+
11
+ providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
12
+ model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
13
+ rmbg_model = rt.InferenceSession(model_path, providers=providers)
14
+
15
+ trimapGenerator = TrimapGenerator()
16
+
17
+
18
+ # def custom_background(background, foreground):
19
+ # foreground = ImageOps.contain(foreground, background.size)
20
+ # x = (background.size[0] - foreground.size[0]) // 2
21
+ # y = (background.size[1] - foreground.size[1]) // 2
22
+ # background.paste(foreground, (x, y), foreground)
23
+ # return background
24
+
25
+ def custom_background(background: PIL.Image.Image, foreground: np.ndarray):
26
+ final_foreground = PIL.Image.fromarray(foreground)
27
+ x = (background.size[0] - final_foreground.size[0]) / 2
28
+ y = (background.size[1] - final_foreground.size[1]) / 2
29
+ box = (x, y, final_foreground.size[0] + x, final_foreground.size[1] + y)
30
+ crop = background.crop(box)
31
+ final_image = crop.copy()
32
+ # put the foreground in the centre of the background
33
+ paste_box = (0, final_image.size[1] - final_foreground.size[1], final_image.size[0], final_image.size[1])
34
+ final_image.paste(final_foreground, paste_box, mask=final_foreground)
35
+ return np.array(final_image)
36
+
37
+
38
+ def get_mask(img, s=1024):
39
+ img = (img / 255).astype(np.float32)
40
+ h, w = h0, w0 = img.shape[:-1]
41
+ h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
42
+ ph, pw = s - h, s - w
43
+ img_input = np.zeros([s, s, 3], dtype=np.float32)
44
+ img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h))
45
+ img_input = np.transpose(img_input, (2, 0, 1))
46
+ img_input = img_input[np.newaxis, :]
47
+ mask = rmbg_model.run(None, {'img': img_input})[0][0]
48
+ mask = np.transpose(mask, (1, 2, 0))
49
+ mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
50
+ mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
51
+ return mask
52
+
53
+
54
+ def change_background_color(image, color="blue"):
55
+ mask = get_mask(image)
56
+ image = (mask * image + 255 * (1 - mask)).astype(np.uint8)
57
+ mask = (mask * 255).astype(np.uint8)
58
+ image = np.concatenate([image, mask], axis=2, dtype=np.uint8)
59
+ image = PIL.Image.fromarray(image)
60
+ background = PIL.Image.new('RGB', image.size, color)
61
+ background.paste(image, (0, 0), image)
62
+ return background
63
+
64
+
65
+ def generate_trimap(probs, size=7, conf_threshold=0.95):
66
+ """
67
+ This function creates a trimap based on simple dilation algorithm
68
+ Inputs [3]: an image with probabilities of each pixel being the foreground, size of dilation kernel,
69
+ foreground confidence threshold
70
+ Output : a trimap
71
+ """
72
+ mask = (probs > 0.05).astype(np.uint8) * 255
73
+ pixels = 2 * size + 1
74
+ kernel = np.ones((pixels, pixels), np.uint8)
75
+ dilation = cv2.dilate(mask, kernel, iterations=1)
76
+ remake = np.zeros_like(mask)
77
+ remake[dilation == 255] = 127 # Set every pixel within dilated region as probably foreground.
78
+ remake[probs > conf_threshold] = 255 # Set every pixel with large enough probability as definitely foreground.
79
+ return remake
80
+
81
+
82
+ def image2gray(image):
83
+ image = PIL.Image.fromarray(image).convert("L")
84
+ return np.array(image) / 255.0
85
+
86
+
87
+ def paste(img_orig, alpha):
88
+ img_ = img_orig.astype(np.float32) / 255
89
+ alpha_ = cv2.resize(alpha, (img_.shape[1], img_.shape[0]), cv2.INTER_LANCZOS4)
90
+ fg_alpha = np.concatenate([img_, alpha_[:, :, np.newaxis]], axis=2)
91
+ cv2.imwrite("new_back.png", (fg_alpha * 255).astype(np.uint8))
92
+
93
+
94
+ def predict(image, new_background):
95
+ mask = get_mask(image)
96
+ mask = (mask * 255).astype(np.uint8)
97
+ mask = mask.repeat(3, axis=2)
98
+
99
+ trimap = generate_trimap(mask)
100
+ trimap = image2gray(trimap)
101
+ # trimap = load_image("images/trimaps/lemur_trimap.png", "GRAY")
102
+
103
+ original = PIL.Image.fromarray(image)
104
+ # mask = image2gray(mask)
105
+ mask = PIL.Image.fromarray(mask).convert("L")
106
+ trimap = trimapGenerator(original_image=original, mask=mask)
107
+ trimap = np.array(trimap) / 255.0
108
+
109
+ foreground = image / 255
110
+ alpha = estimate_alpha_cf(foreground, trimap)
111
+ foreground = estimate_foreground_ml(foreground, alpha)
112
+ cutout = stack_images(foreground, alpha)
113
+ cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
114
+
115
+ if new_background is not None:
116
+ return mask, trimap, custom_background(new_background, cutout)
117
+ return alpha, trimap, cutout
118
+
119
+
120
+ # contours
121
+ def serendipity(image, new_background):
122
+ mask = get_mask(image)
123
+ mask = 255 - mask
124
+ image = (mask * image + 255 * (1 - mask)).astype(np.uint8)
125
+ mask = (mask * 255).astype(np.uint8)
126
+ image = np.concatenate([image, mask], axis=2, dtype=np.uint8)
127
+ return mask, image
128
+
129
+
130
+ def negative(image, new_background):
131
+ mask = get_mask(image)
132
+ image = (mask * image + 255 * (1 - mask)).astype(np.uint8)
133
+ image = 255 - image
134
+ mask = (mask * 255).astype(np.uint8)
135
+ image = np.concatenate([image, mask], axis=2, dtype=np.uint8)
136
+ return mask, image
137
+
138
+
139
+ def checkit(image, new_background):
140
+ mask = get_mask(image)
141
+ mask = 255 - mask
142
+ image = (mask / image - 255 / (1 + mask)).astype(np.uint8)
143
+ mask = (mask * 255).astype(np.uint8)
144
+ mask = 255 - mask
145
+ image = np.concatenate([image, mask], axis=2, dtype=np.uint8)
146
+ mask = mask.repeat(3, axis=2)
147
+ # if new_background is not None:
148
+ # foreground = PIL.Image.fromarray(image)
149
+ # return mask, custom_background(new_background, foreground)
150
+ return mask, image
151
+
152
+
153
+ footer = r"""
154
+ <center>
155
+ <b>
156
+ Demo based on <a href='https://github.com/SkyTNT/anime-segmentation'>SkyTNT Anime Segmentation</a>
157
+ </b>
158
+ </center>
159
+ """
160
+
161
+ with gr.Blocks(title="Face Shine") as app:
162
+ gr.HTML("<center><h1>Anime Remove Background</h1></center>")
163
+ with gr.Row():
164
+ with gr.Column():
165
+ input_img = gr.Image(type="numpy", image_mode="RGB", label="Input image")
166
+ new_img = gr.Image(type="pil", image_mode="RGBA", label="Custom background")
167
+ run_btn = gr.Button(variant="primary")
168
+ with gr.Column():
169
+ with gr.Accordion(label="Image mask", open=False):
170
+ output_mask = gr.Image(type="numpy", label="mask")
171
+ output_trimap = gr.Image(type="numpy", label="trimap")
172
+ output_img = gr.Image(type="numpy", label="result")
173
+
174
+ run_btn.click(predict, [input_img, new_img], [output_mask, output_trimap, output_img])
175
+
176
+ with gr.Row():
177
+ examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
178
+ examples = gr.Dataset(components=[input_img], samples=examples_data)
179
+ examples.click(lambda x: x[0], [examples], [input_img])
180
+
181
+ with gr.Row():
182
+ gr.HTML(footer)
183
+
184
+ app.launch(share=False, debug=True, enable_queue=True, show_error=True)