change background

app.py

@@ -1,7 +1,16 @@
+############################################################################################################
+#
+# Source from
+# https://github.com/eugenesiow/practical-ml/blob/master/notebooks/Remove_Image_Background_DeepLabV3.ipynb
+#
+############################################################################################################
+
 import os
+
 import cv2
-import numpy as np
 import torch
+import PIL.Image
+import numpy as np
 import gradio as gr
 import torchvision.transforms as transforms
 
@@ -18,6 +27,18 @@ def image_to_tensor(image):
     ])(image)
 
 
+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
+
+
 def make_transparent_foreground(image, mask):
     # split the image into channels
     b, g, r = cv2.split(np.array(image).astype('uint8'))
@@ -52,9 +73,13 @@ def makeMask(image):
     return np.where(mask, 255, background).astype(np.uint8)
 
 
-def predict(image):
+def predict(image, new_background=None):
     mask = makeMask(image)
-    return make_transparent_foreground(image, mask)
+    foreground = make_transparent_foreground(image, mask)
+    if new_background is not None:
+        foreground = PIL.Image.fromarray(foreground)
+        return custom_background(new_background, foreground)
+    return foreground
 
 
 title = "Zero Background"
@@ -78,9 +103,11 @@ This demo is running on a CPU, if you like this project please make us a donatio
 <center><img src='https://visitor-badge.glitch.me/badge?page_id=deoldify.visitor-badge' alt='visitor badge'></center>
 """
 
+
 demo = gr.Interface(
     predict, [
         gr.Image(type="pil", label="Image"),
+        gr.Image(type="pil", label="Optionally: Set a new background")
     ], [
         gr.Image(type="pil", label="Image alpha background")
     ],

requirements.txt

@@ -1,6 +1,5 @@
 torch>=2.0.1
-torchvision
-opencv-python
-
-
-
+torchvision~=0.15.2
+opencv-python~=4.7.0.72
+numpy~=1.24.3
+pillow~=9.5.0