prithivMLmods commited on
Commit
bccd6e8
·
verified ·
1 Parent(s): b6c1f2b

Upload 5 files

Browse files
Files changed (5) hide show
  1. app.py +130 -0
  2. src/__init__.py +0 -0
  3. src/core.py +466 -0
  4. src/helper.py +86 -0
  5. src/st_style.py +42 -0
app.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import pandas as pd
3
+ import streamlit as st
4
+ import os
5
+ from datetime import datetime
6
+ from PIL import Image
7
+ from streamlit_drawable_canvas import st_canvas
8
+ from io import BytesIO
9
+ from copy import deepcopy
10
+
11
+ from src.core import process_inpaint
12
+
13
+
14
+ def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
15
+ if fmt not in ["jpg", "png"]:
16
+ raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
17
+
18
+ pil_format = "JPEG" if fmt == "jpg" else "PNG"
19
+ file_format = "jpg" if fmt == "jpg" else "png"
20
+ mime = "image/jpeg" if fmt == "jpg" else "image/png"
21
+
22
+ buf = BytesIO()
23
+ pil_image.save(buf, format=pil_format)
24
+
25
+ return st.download_button(
26
+ label=label,
27
+ data=buf.getvalue(),
28
+ file_name=f'{filename}.{file_format}',
29
+ mime=mime,
30
+ )
31
+
32
+ if "button_id" not in st.session_state:
33
+ st.session_state["button_id"] = ""
34
+ if "color_to_label" not in st.session_state:
35
+ st.session_state["color_to_label"] = {}
36
+
37
+ if 'reuse_image' not in st.session_state:
38
+ st.session_state.reuse_image = None
39
+ def set_image(img):
40
+ st.session_state.reuse_image = img
41
+
42
+ st.title("Magic-Eraser")
43
+
44
+ st.image(open("assets/magic-eraser.png", "rb").read())
45
+
46
+ st.markdown(
47
+ """
48
+ You don't have to worry about mastering photo editing techniques to remove an object from your photo. **Simply mark over the areas you want to erase, and our AI will take care of the rest.**
49
+ """
50
+ )
51
+ uploaded_file = st.file_uploader("Choose image", accept_multiple_files=False, type=["png", "jpg", "jpeg"])
52
+
53
+ if uploaded_file is not None:
54
+
55
+ if st.session_state.reuse_image is not None:
56
+ img_input = Image.fromarray(st.session_state.reuse_image)
57
+ else:
58
+ bytes_data = uploaded_file.getvalue()
59
+ img_input = Image.open(BytesIO(bytes_data)).convert("RGBA")
60
+
61
+ #resize
62
+ max_size = 2000
63
+ img_width, img_height = img_input.size
64
+ if img_width > max_size or img_height > max_size:
65
+ if img_width > img_height:
66
+ new_width = max_size
67
+ new_height = int((max_size / img_width) * img_height)
68
+ else:
69
+ new_height = max_size
70
+ new_width = int((max_size / img_height) * img_width)
71
+ img_input = img_input.resize((new_width, new_height))
72
+
73
+ stroke_width = st.slider("Brush size", 1, 100, 50)
74
+
75
+ st.write("**Now draw (brush) the part of image that you want to remove.**")
76
+
77
+ canvas_bg = deepcopy(img_input)
78
+ aspect_ratio = canvas_bg.width / canvas_bg.height
79
+ streamlit_width = 720
80
+
81
+ if canvas_bg.width > streamlit_width:
82
+ canvas_bg = canvas_bg.resize((streamlit_width, int(streamlit_width / aspect_ratio)))
83
+
84
+ canvas_result = st_canvas(
85
+ stroke_color="rgba(255, 0, 255, 1)",
86
+ stroke_width=stroke_width,
87
+ background_image=canvas_bg,
88
+ width=canvas_bg.width,
89
+ height=canvas_bg.height,
90
+ drawing_mode="freedraw",
91
+ key="compute_arc_length",
92
+ )
93
+
94
+ if canvas_result.image_data is not None:
95
+ im = np.array(Image.fromarray(canvas_result.image_data.astype(np.uint8)).resize(img_input.size))
96
+ background = np.where(
97
+ (im[:, :, 0] == 0) &
98
+ (im[:, :, 1] == 0) &
99
+ (im[:, :, 2] == 0)
100
+ )
101
+ drawing = np.where(
102
+ (im[:, :, 0] == 255) &
103
+ (im[:, :, 1] == 0) &
104
+ (im[:, :, 2] == 255)
105
+ )
106
+ im[background]=[0,0,0,255]
107
+ im[drawing]=[0,0,0,0] # RGBA
108
+
109
+ reuse = False
110
+
111
+ if st.button('Submit'):
112
+
113
+ with st.spinner("AI is doing the magic!"):
114
+ output = process_inpaint(np.array(img_input), np.array(im)) #TODO Put button here
115
+ img_output = Image.fromarray(output).convert("RGB")
116
+
117
+ st.write("AI has finished the job!")
118
+ st.image(img_output)
119
+ # reuse = st.button('Edit again (Re-use this image)', on_click=set_image, args=(inpainted_img, ))
120
+
121
+ uploaded_name = os.path.splitext(uploaded_file.name)[0]
122
+ image_download_button(
123
+ pil_image=img_output,
124
+ filename=uploaded_name,
125
+ fmt="jpg",
126
+ label="Download Image"
127
+ )
128
+
129
+ st.info("**TIP**: If the result is not perfect, you can download it then "
130
+ "upload then remove the artifacts.")
src/__init__.py ADDED
File without changes
src/core.py ADDED
@@ -0,0 +1,466 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import base64
2
+ import json
3
+ import os
4
+ import re
5
+ import time
6
+ import uuid
7
+ from io import BytesIO
8
+ from pathlib import Path
9
+ import cv2
10
+
11
+ # For inpainting
12
+
13
+ import numpy as np
14
+ import pandas as pd
15
+ import streamlit as st
16
+ from PIL import Image
17
+ from streamlit_drawable_canvas import st_canvas
18
+
19
+
20
+ import argparse
21
+ import io
22
+ import multiprocessing
23
+ from typing import Union
24
+
25
+ import torch
26
+
27
+ try:
28
+ torch._C._jit_override_can_fuse_on_cpu(False)
29
+ torch._C._jit_override_can_fuse_on_gpu(False)
30
+ torch._C._jit_set_texpr_fuser_enabled(False)
31
+ torch._C._jit_set_nvfuser_enabled(False)
32
+ except:
33
+ pass
34
+
35
+ from src.helper import (
36
+ download_model,
37
+ load_img,
38
+ norm_img,
39
+ numpy_to_bytes,
40
+ pad_img_to_modulo,
41
+ resize_max_size,
42
+ )
43
+
44
+ NUM_THREADS = str(multiprocessing.cpu_count())
45
+
46
+ os.environ["OMP_NUM_THREADS"] = NUM_THREADS
47
+ os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS
48
+ os.environ["MKL_NUM_THREADS"] = NUM_THREADS
49
+ os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS
50
+ os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS
51
+ if os.environ.get("CACHE_DIR"):
52
+ os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"]
53
+
54
+ #BUILD_DIR = os.environ.get("LAMA_CLEANER_BUILD_DIR", "./lama_cleaner/app/build")
55
+
56
+ # For Seam-carving
57
+
58
+ from scipy import ndimage as ndi
59
+
60
+ SEAM_COLOR = np.array([255, 200, 200]) # seam visualization color (BGR)
61
+ SHOULD_DOWNSIZE = True # if True, downsize image for faster carving
62
+ DOWNSIZE_WIDTH = 500 # resized image width if SHOULD_DOWNSIZE is True
63
+ ENERGY_MASK_CONST = 100000.0 # large energy value for protective masking
64
+ MASK_THRESHOLD = 10 # minimum pixel intensity for binary mask
65
+ USE_FORWARD_ENERGY = True # if True, use forward energy algorithm
66
+
67
+ device = torch.device("cpu")
68
+ model_path = "./pt/erase.pt"
69
+ model = torch.jit.load(model_path, map_location="cpu")
70
+ model = model.to(device)
71
+ model.eval()
72
+
73
+
74
+ ########################################
75
+ # UTILITY CODE
76
+ ########################################
77
+
78
+
79
+ def visualize(im, boolmask=None, rotate=False):
80
+ vis = im.astype(np.uint8)
81
+ if boolmask is not None:
82
+ vis[np.where(boolmask == False)] = SEAM_COLOR
83
+ if rotate:
84
+ vis = rotate_image(vis, False)
85
+ cv2.imshow("visualization", vis)
86
+ cv2.waitKey(1)
87
+ return vis
88
+
89
+ def resize(image, width):
90
+ dim = None
91
+ h, w = image.shape[:2]
92
+ dim = (width, int(h * width / float(w)))
93
+ image = image.astype('float32')
94
+ return cv2.resize(image, dim)
95
+
96
+ def rotate_image(image, clockwise):
97
+ k = 1 if clockwise else 3
98
+ return np.rot90(image, k)
99
+
100
+
101
+ ########################################
102
+ # ENERGY FUNCTIONS
103
+ ########################################
104
+
105
+ def backward_energy(im):
106
+ """
107
+ Simple gradient magnitude energy map.
108
+ """
109
+ xgrad = ndi.convolve1d(im, np.array([1, 0, -1]), axis=1, mode='wrap')
110
+ ygrad = ndi.convolve1d(im, np.array([1, 0, -1]), axis=0, mode='wrap')
111
+
112
+ grad_mag = np.sqrt(np.sum(xgrad**2, axis=2) + np.sum(ygrad**2, axis=2))
113
+
114
+ # vis = visualize(grad_mag)
115
+ # cv2.imwrite("backward_energy_demo.jpg", vis)
116
+
117
+ return grad_mag
118
+
119
+ def forward_energy(im):
120
+ """
121
+ Forward energy algorithm as described in "Improved Seam Carving for Video Retargeting"
122
+ by Rubinstein, Shamir, Avidan.
123
+ Vectorized code adapted from
124
+ https://github.com/axu2/improved-seam-carving.
125
+ """
126
+ h, w = im.shape[:2]
127
+ im = cv2.cvtColor(im.astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float64)
128
+
129
+ energy = np.zeros((h, w))
130
+ m = np.zeros((h, w))
131
+
132
+ U = np.roll(im, 1, axis=0)
133
+ L = np.roll(im, 1, axis=1)
134
+ R = np.roll(im, -1, axis=1)
135
+
136
+ cU = np.abs(R - L)
137
+ cL = np.abs(U - L) + cU
138
+ cR = np.abs(U - R) + cU
139
+
140
+ for i in range(1, h):
141
+ mU = m[i-1]
142
+ mL = np.roll(mU, 1)
143
+ mR = np.roll(mU, -1)
144
+
145
+ mULR = np.array([mU, mL, mR])
146
+ cULR = np.array([cU[i], cL[i], cR[i]])
147
+ mULR += cULR
148
+
149
+ argmins = np.argmin(mULR, axis=0)
150
+ m[i] = np.choose(argmins, mULR)
151
+ energy[i] = np.choose(argmins, cULR)
152
+
153
+ # vis = visualize(energy)
154
+ # cv2.imwrite("forward_energy_demo.jpg", vis)
155
+
156
+ return energy
157
+
158
+ ########################################
159
+ # SEAM HELPER FUNCTIONS
160
+ ########################################
161
+
162
+ def add_seam(im, seam_idx):
163
+ """
164
+ Add a vertical seam to a 3-channel color image at the indices provided
165
+ by averaging the pixels values to the left and right of the seam.
166
+ Code adapted from https://github.com/vivianhylee/seam-carving.
167
+ """
168
+ h, w = im.shape[:2]
169
+ output = np.zeros((h, w + 1, 3))
170
+ for row in range(h):
171
+ col = seam_idx[row]
172
+ for ch in range(3):
173
+ if col == 0:
174
+ p = np.mean(im[row, col: col + 2, ch])
175
+ output[row, col, ch] = im[row, col, ch]
176
+ output[row, col + 1, ch] = p
177
+ output[row, col + 1:, ch] = im[row, col:, ch]
178
+ else:
179
+ p = np.mean(im[row, col - 1: col + 1, ch])
180
+ output[row, : col, ch] = im[row, : col, ch]
181
+ output[row, col, ch] = p
182
+ output[row, col + 1:, ch] = im[row, col:, ch]
183
+
184
+ return output
185
+
186
+ def add_seam_grayscale(im, seam_idx):
187
+ """
188
+ Add a vertical seam to a grayscale image at the indices provided
189
+ by averaging the pixels values to the left and right of the seam.
190
+ """
191
+ h, w = im.shape[:2]
192
+ output = np.zeros((h, w + 1))
193
+ for row in range(h):
194
+ col = seam_idx[row]
195
+ if col == 0:
196
+ p = np.mean(im[row, col: col + 2])
197
+ output[row, col] = im[row, col]
198
+ output[row, col + 1] = p
199
+ output[row, col + 1:] = im[row, col:]
200
+ else:
201
+ p = np.mean(im[row, col - 1: col + 1])
202
+ output[row, : col] = im[row, : col]
203
+ output[row, col] = p
204
+ output[row, col + 1:] = im[row, col:]
205
+
206
+ return output
207
+
208
+ def remove_seam(im, boolmask):
209
+ h, w = im.shape[:2]
210
+ boolmask3c = np.stack([boolmask] * 3, axis=2)
211
+ return im[boolmask3c].reshape((h, w - 1, 3))
212
+
213
+ def remove_seam_grayscale(im, boolmask):
214
+ h, w = im.shape[:2]
215
+ return im[boolmask].reshape((h, w - 1))
216
+
217
+ def get_minimum_seam(im, mask=None, remove_mask=None):
218
+ """
219
+ DP algorithm for finding the seam of minimum energy. Code adapted from
220
+ https://karthikkaranth.me/blog/implementing-seam-carving-with-python/
221
+ """
222
+ h, w = im.shape[:2]
223
+ energyfn = forward_energy if USE_FORWARD_ENERGY else backward_energy
224
+ M = energyfn(im)
225
+
226
+ if mask is not None:
227
+ M[np.where(mask > MASK_THRESHOLD)] = ENERGY_MASK_CONST
228
+
229
+ # give removal mask priority over protective mask by using larger negative value
230
+ if remove_mask is not None:
231
+ M[np.where(remove_mask > MASK_THRESHOLD)] = -ENERGY_MASK_CONST * 100
232
+
233
+ seam_idx, boolmask = compute_shortest_path(M, im, h, w)
234
+
235
+ return np.array(seam_idx), boolmask
236
+
237
+ def compute_shortest_path(M, im, h, w):
238
+ backtrack = np.zeros_like(M, dtype=np.int_)
239
+
240
+
241
+ # populate DP matrix
242
+ for i in range(1, h):
243
+ for j in range(0, w):
244
+ if j == 0:
245
+ idx = np.argmin(M[i - 1, j:j + 2])
246
+ backtrack[i, j] = idx + j
247
+ min_energy = M[i-1, idx + j]
248
+ else:
249
+ idx = np.argmin(M[i - 1, j - 1:j + 2])
250
+ backtrack[i, j] = idx + j - 1
251
+ min_energy = M[i - 1, idx + j - 1]
252
+
253
+ M[i, j] += min_energy
254
+
255
+ # backtrack to find path
256
+ seam_idx = []
257
+ boolmask = np.ones((h, w), dtype=np.bool_)
258
+ j = np.argmin(M[-1])
259
+ for i in range(h-1, -1, -1):
260
+ boolmask[i, j] = False
261
+ seam_idx.append(j)
262
+ j = backtrack[i, j]
263
+
264
+ seam_idx.reverse()
265
+ return seam_idx, boolmask
266
+
267
+ ########################################
268
+ # MAIN ALGORITHM
269
+ ########################################
270
+
271
+ def seams_removal(im, num_remove, mask=None, vis=False, rot=False):
272
+ for _ in range(num_remove):
273
+ seam_idx, boolmask = get_minimum_seam(im, mask)
274
+ if vis:
275
+ visualize(im, boolmask, rotate=rot)
276
+ im = remove_seam(im, boolmask)
277
+ if mask is not None:
278
+ mask = remove_seam_grayscale(mask, boolmask)
279
+ return im, mask
280
+
281
+
282
+ def seams_insertion(im, num_add, mask=None, vis=False, rot=False):
283
+ seams_record = []
284
+ temp_im = im.copy()
285
+ temp_mask = mask.copy() if mask is not None else None
286
+
287
+ for _ in range(num_add):
288
+ seam_idx, boolmask = get_minimum_seam(temp_im, temp_mask)
289
+ if vis:
290
+ visualize(temp_im, boolmask, rotate=rot)
291
+
292
+ seams_record.append(seam_idx)
293
+ temp_im = remove_seam(temp_im, boolmask)
294
+ if temp_mask is not None:
295
+ temp_mask = remove_seam_grayscale(temp_mask, boolmask)
296
+
297
+ seams_record.reverse()
298
+
299
+ for _ in range(num_add):
300
+ seam = seams_record.pop()
301
+ im = add_seam(im, seam)
302
+ if vis:
303
+ visualize(im, rotate=rot)
304
+ if mask is not None:
305
+ mask = add_seam_grayscale(mask, seam)
306
+
307
+ # update the remaining seam indices
308
+ for remaining_seam in seams_record:
309
+ remaining_seam[np.where(remaining_seam >= seam)] += 2
310
+
311
+ return im, mask
312
+
313
+ ########################################
314
+ # MAIN DRIVER FUNCTIONS
315
+ ########################################
316
+
317
+ def seam_carve(im, dy, dx, mask=None, vis=False):
318
+ im = im.astype(np.float64)
319
+ h, w = im.shape[:2]
320
+ assert h + dy > 0 and w + dx > 0 and dy <= h and dx <= w
321
+
322
+ if mask is not None:
323
+ mask = mask.astype(np.float64)
324
+
325
+ output = im
326
+
327
+ if dx < 0:
328
+ output, mask = seams_removal(output, -dx, mask, vis)
329
+
330
+ elif dx > 0:
331
+ output, mask = seams_insertion(output, dx, mask, vis)
332
+
333
+ if dy < 0:
334
+ output = rotate_image(output, True)
335
+ if mask is not None:
336
+ mask = rotate_image(mask, True)
337
+ output, mask = seams_removal(output, -dy, mask, vis, rot=True)
338
+ output = rotate_image(output, False)
339
+
340
+ elif dy > 0:
341
+ output = rotate_image(output, True)
342
+ if mask is not None:
343
+ mask = rotate_image(mask, True)
344
+ output, mask = seams_insertion(output, dy, mask, vis, rot=True)
345
+ output = rotate_image(output, False)
346
+
347
+ return output
348
+
349
+
350
+ def object_removal(im, rmask, mask=None, vis=False, horizontal_removal=False):
351
+ im = im.astype(np.float64)
352
+ rmask = rmask.astype(np.float64)
353
+ if mask is not None:
354
+ mask = mask.astype(np.float64)
355
+ output = im
356
+
357
+ h, w = im.shape[:2]
358
+
359
+ if horizontal_removal:
360
+ output = rotate_image(output, True)
361
+ rmask = rotate_image(rmask, True)
362
+ if mask is not None:
363
+ mask = rotate_image(mask, True)
364
+
365
+ while len(np.where(rmask > MASK_THRESHOLD)[0]) > 0:
366
+ seam_idx, boolmask = get_minimum_seam(output, mask, rmask)
367
+ if vis:
368
+ visualize(output, boolmask, rotate=horizontal_removal)
369
+ output = remove_seam(output, boolmask)
370
+ rmask = remove_seam_grayscale(rmask, boolmask)
371
+ if mask is not None:
372
+ mask = remove_seam_grayscale(mask, boolmask)
373
+
374
+ num_add = (h if horizontal_removal else w) - output.shape[1]
375
+ output, mask = seams_insertion(output, num_add, mask, vis, rot=horizontal_removal)
376
+ if horizontal_removal:
377
+ output = rotate_image(output, False)
378
+
379
+ return output
380
+
381
+
382
+
383
+ def s_image(im,mask,vs,hs,mode="resize"):
384
+ im = cv2.cvtColor(im, cv2.COLOR_RGBA2RGB)
385
+ mask = 255-mask[:,:,3]
386
+ h, w = im.shape[:2]
387
+ if SHOULD_DOWNSIZE and w > DOWNSIZE_WIDTH:
388
+ im = resize(im, width=DOWNSIZE_WIDTH)
389
+ if mask is not None:
390
+ mask = resize(mask, width=DOWNSIZE_WIDTH)
391
+
392
+ # image resize mode
393
+ if mode=="resize":
394
+ dy = hs#reverse
395
+ dx = vs#reverse
396
+ assert dy is not None and dx is not None
397
+ output = seam_carve(im, dy, dx, mask, False)
398
+
399
+
400
+ # object removal mode
401
+ elif mode=="remove":
402
+ assert mask is not None
403
+ output = object_removal(im, mask, None, False, True)
404
+
405
+ return output
406
+
407
+
408
+ ##### Inpainting helper code
409
+
410
+ def run(image, mask):
411
+ """
412
+ image: [C, H, W]
413
+ mask: [1, H, W]
414
+ return: BGR IMAGE
415
+ """
416
+ origin_height, origin_width = image.shape[1:]
417
+ image = pad_img_to_modulo(image, mod=8)
418
+ mask = pad_img_to_modulo(mask, mod=8)
419
+
420
+ mask = (mask > 0) * 1
421
+ image = torch.from_numpy(image).unsqueeze(0).to(device)
422
+ mask = torch.from_numpy(mask).unsqueeze(0).to(device)
423
+
424
+ start = time.time()
425
+ with torch.no_grad():
426
+ inpainted_image = model(image, mask)
427
+
428
+ print(f"process time: {(time.time() - start)*1000}ms")
429
+ cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy()
430
+ cur_res = cur_res[0:origin_height, 0:origin_width, :]
431
+ cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
432
+ cur_res = cv2.cvtColor(cur_res, cv2.COLOR_BGR2RGB)
433
+ return cur_res
434
+
435
+
436
+ def get_args_parser():
437
+ parser = argparse.ArgumentParser()
438
+ parser.add_argument("--port", default=8080, type=int)
439
+ parser.add_argument("--device", default="cuda", type=str)
440
+ parser.add_argument("--debug", action="store_true")
441
+ return parser.parse_args()
442
+
443
+
444
+ def process_inpaint(image, mask):
445
+ image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
446
+ original_shape = image.shape
447
+ interpolation = cv2.INTER_CUBIC
448
+
449
+ #size_limit: Union[int, str] = request.form.get("sizeLimit", "1080")
450
+ #if size_limit == "Original":
451
+ size_limit = max(image.shape)
452
+ #else:
453
+ # size_limit = int(size_limit)
454
+
455
+ print(f"Origin image shape: {original_shape}")
456
+ image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
457
+ print(f"Resized image shape: {image.shape}")
458
+ image = norm_img(image)
459
+
460
+ mask = 255-mask[:,:,3]
461
+ mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
462
+ mask = norm_img(mask)
463
+
464
+ res_np_img = run(image, mask)
465
+
466
+ return cv2.cvtColor(res_np_img, cv2.COLOR_BGR2RGB)
src/helper.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+
4
+ from urllib.parse import urlparse
5
+ import cv2
6
+ import numpy as np
7
+ import torch
8
+ from torch.hub import download_url_to_file, get_dir
9
+
10
+ ERASE_MODEL_URL = os.environ.get(
11
+ "ERASE_MODEL_URL",
12
+ "https://huggingface.co/spaces/prithivMLmods/Magic-Eraser-Tool/blob/main/pt/erase.pt",
13
+ )
14
+
15
+ def download_model(url=ERASE_MODEL_URL):
16
+ parts = urlparse(url)
17
+ hub_dir = get_dir()
18
+ model_dir = os.path.join(hub_dir, "checkpoints")
19
+ if not os.path.isdir(model_dir):
20
+ os.makedirs(os.path.join(model_dir, "hub", "checkpoints"))
21
+ filename = os.path.basename(parts.path)
22
+ cached_file = os.path.join(model_dir, filename)
23
+ if not os.path.exists(cached_file):
24
+ sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
25
+ hash_prefix = None
26
+ download_url_to_file(url, cached_file, hash_prefix, progress=True)
27
+ return cached_file
28
+
29
+
30
+ def ceil_modulo(x, mod):
31
+ if x % mod == 0:
32
+ return x
33
+ return (x // mod + 1) * mod
34
+
35
+
36
+ def numpy_to_bytes(image_numpy: np.ndarray) -> bytes:
37
+ data = cv2.imencode(".jpg", image_numpy)[1]
38
+ image_bytes = data.tobytes()
39
+ return image_bytes
40
+
41
+
42
+ def load_img(img_bytes, gray: bool = False):
43
+ nparr = np.frombuffer(img_bytes, np.uint8)
44
+ if gray:
45
+ np_img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
46
+ else:
47
+ np_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
48
+ if len(np_img.shape) == 3 and np_img.shape[2] == 4:
49
+ np_img = cv2.cvtColor(np_img, cv2.COLOR_BGRA2RGB)
50
+ else:
51
+ np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)
52
+
53
+ return np_img
54
+
55
+
56
+ def norm_img(np_img):
57
+ if len(np_img.shape) == 2:
58
+ np_img = np_img[:, :, np.newaxis]
59
+ np_img = np.transpose(np_img, (2, 0, 1))
60
+ np_img = np_img.astype("float32") / 255
61
+ return np_img
62
+
63
+
64
+ def resize_max_size(
65
+ np_img, size_limit: int, interpolation=cv2.INTER_CUBIC
66
+ ) -> np.ndarray:
67
+ # Resize image's longer size to size_limit if longer size larger than size_limit
68
+ h, w = np_img.shape[:2]
69
+ if max(h, w) > size_limit:
70
+ ratio = size_limit / max(h, w)
71
+ new_w = int(w * ratio + 0.5)
72
+ new_h = int(h * ratio + 0.5)
73
+ return cv2.resize(np_img, dsize=(new_w, new_h), interpolation=interpolation)
74
+ else:
75
+ return np_img
76
+
77
+
78
+ def pad_img_to_modulo(img, mod):
79
+ channels, height, width = img.shape
80
+ out_height = ceil_modulo(height, mod)
81
+ out_width = ceil_modulo(width, mod)
82
+ return np.pad(
83
+ img,
84
+ ((0, 0), (0, out_height - height), (0, out_width - width)),
85
+ mode="symmetric",
86
+ )
src/st_style.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ button_style = """
2
+ <style>
3
+ div.stButton > button:first-child {
4
+ background-color: rgb(255, 75, 75);
5
+ color: rgb(255, 255, 255);
6
+ }
7
+ div.stButton > button:hover {
8
+ background-color: rgb(255, 75, 75);
9
+ color: rgb(255, 255, 255);
10
+ }
11
+ div.stButton > button:active {
12
+ background-color: rgb(255, 75, 75);
13
+ color: rgb(255, 255, 255);
14
+ }
15
+ div.stButton > button:focus {
16
+ background-color: rgb(255, 75, 75);
17
+ color: rgb(255, 255, 255);
18
+ }
19
+ .css-1cpxqw2:focus:not(:active) {
20
+ background-color: rgb(255, 75, 75);
21
+ border-color: rgb(255, 75, 75);
22
+ color: rgb(255, 255, 255);
23
+ }
24
+ """
25
+
26
+ style = """
27
+ <style>
28
+ #MainMenu {
29
+ visibility: hidden;
30
+ }
31
+ footer {
32
+ visibility: hidden;
33
+ }
34
+ header {
35
+ visibility: hidden;
36
+ }
37
+ </style>
38
+ """
39
+
40
+
41
+ def apply_prod_style(st):
42
+ return st.markdown(style, unsafe_allow_html=True)