fix background paste

app.py

@@ -4,6 +4,7 @@ import huggingface_hub
 import onnxruntime as rt
 import numpy as np
 import cv2
+from PIL import ImageOps
 
 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
@@ -11,15 +12,11 @@ rmbg_model = rt.InferenceSession(model_path, providers=providers)
 
 
 def custom_background(background, foreground):
-    x = (background.size[0] - foreground.size[0]) / 2
-    y = (background.size[1] - foreground.size[1]) / 2
-    box = (x, y, foreground.size[0] + x, foreground.size[1] + y)
-    crop = background.crop(box)
-    final_image = crop.copy()
-    # put the foreground in the centre of the background
-    paste_box = (0, final_image.size[1] - foreground.size[1], final_image.size[0], final_image.size[1])
-    final_image.paste(foreground, paste_box, mask=foreground)
-    return final_image
+    foreground = ImageOps.contain(foreground, background.size)
+    x = (background.size[0] - foreground.size[0]) // 2
+    y = (background.size[1] - foreground.size[1]) // 2
+    background.paste(foreground, (x, y), foreground)
+    return background
 
 
 def get_mask(img, s=1024):
@@ -38,16 +35,16 @@ def get_mask(img, s=1024):
     return mask
 
 
-def predict(img, new_background):
-    mask = get_mask(img)
-    img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
+def predict(image, new_background):
+    mask = get_mask(image)
+    image = (mask * image + 255 * (1 - mask)).astype(np.uint8)
     mask = (mask * 255).astype(np.uint8)
-    img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
+    image = np.concatenate([image, mask], axis=2, dtype=np.uint8)
     mask = mask.repeat(3, axis=2)
     if new_background is not None:
-        foreground = PIL.Image.fromarray(img)
+        foreground = PIL.Image.fromarray(image)
         return mask, custom_background(new_background, foreground)
-    return mask, img
+    return mask, image
 
 
 footer = r"""