leonelhs commited on
Commit
aa50dd4
1 Parent(s): 8318e85

fix background paste

Browse files
Files changed (1) hide show
  1. app.py +12 -15
app.py CHANGED
@@ -4,6 +4,7 @@ import huggingface_hub
4
  import onnxruntime as rt
5
  import numpy as np
6
  import cv2
 
7
 
8
  providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
9
  model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
@@ -11,15 +12,11 @@ rmbg_model = rt.InferenceSession(model_path, providers=providers)
11
 
12
 
13
  def custom_background(background, foreground):
14
- x = (background.size[0] - foreground.size[0]) / 2
15
- y = (background.size[1] - foreground.size[1]) / 2
16
- box = (x, y, foreground.size[0] + x, foreground.size[1] + y)
17
- crop = background.crop(box)
18
- final_image = crop.copy()
19
- # put the foreground in the centre of the background
20
- paste_box = (0, final_image.size[1] - foreground.size[1], final_image.size[0], final_image.size[1])
21
- final_image.paste(foreground, paste_box, mask=foreground)
22
- return final_image
23
 
24
 
25
  def get_mask(img, s=1024):
@@ -38,16 +35,16 @@ def get_mask(img, s=1024):
38
  return mask
39
 
40
 
41
- def predict(img, new_background):
42
- mask = get_mask(img)
43
- img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
44
  mask = (mask * 255).astype(np.uint8)
45
- img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
46
  mask = mask.repeat(3, axis=2)
47
  if new_background is not None:
48
- foreground = PIL.Image.fromarray(img)
49
  return mask, custom_background(new_background, foreground)
50
- return mask, img
51
 
52
 
53
  footer = r"""
 
4
  import onnxruntime as rt
5
  import numpy as np
6
  import cv2
7
+ from PIL import ImageOps
8
 
9
  providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
10
  model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
 
12
 
13
 
14
  def custom_background(background, foreground):
15
+ foreground = ImageOps.contain(foreground, background.size)
16
+ x = (background.size[0] - foreground.size[0]) // 2
17
+ y = (background.size[1] - foreground.size[1]) // 2
18
+ background.paste(foreground, (x, y), foreground)
19
+ return background
 
 
 
 
20
 
21
 
22
  def get_mask(img, s=1024):
 
35
  return mask
36
 
37
 
38
+ def predict(image, new_background):
39
+ mask = get_mask(image)
40
+ image = (mask * image + 255 * (1 - mask)).astype(np.uint8)
41
  mask = (mask * 255).astype(np.uint8)
42
+ image = np.concatenate([image, mask], axis=2, dtype=np.uint8)
43
  mask = mask.repeat(3, axis=2)
44
  if new_background is not None:
45
+ foreground = PIL.Image.fromarray(image)
46
  return mask, custom_background(new_background, foreground)
47
+ return mask, image
48
 
49
 
50
  footer = r"""