skytnt commited on
Commit
25c148e
1 Parent(s): 81cf577
Files changed (2) hide show
  1. app.py +2 -2
  2. examples/03.jpg +0 -0
app.py CHANGED
@@ -11,13 +11,13 @@ def get_mask(img, s=1024):
11
  h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
12
  ph, pw = s - h, s - w
13
  img_input = np.zeros([s, s, 3], dtype=np.float32)
14
- img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
15
  img_input = np.transpose(img_input, (2, 0, 1))
16
  img_input = img_input[np.newaxis, :]
17
  mask = rmbg_model.run(None, {'img': img_input})[0][0]
18
  mask = np.transpose(mask, (1, 2, 0))
19
  mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
20
- mask = cv2.resize(mask, (w0, h0), interpolation=cv2.INTER_CUBIC)[:, :, np.newaxis]
21
  return mask
22
 
23
 
11
  h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
12
  ph, pw = s - h, s - w
13
  img_input = np.zeros([s, s, 3], dtype=np.float32)
14
+ img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h))
15
  img_input = np.transpose(img_input, (2, 0, 1))
16
  img_input = img_input[np.newaxis, :]
17
  mask = rmbg_model.run(None, {'img': img_input})[0][0]
18
  mask = np.transpose(mask, (1, 2, 0))
19
  mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
20
+ mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
21
  return mask
22
 
23
 
examples/03.jpg CHANGED