put back background in try-on output

app.py

@@ -67,8 +67,31 @@ print("########################Setup done!########################")
 print(f"####Using {device}.#####")
 u2net = u2net_load.model(model_name = 'u2netp')
 
+
+def composite_background(img_mask, person_image_path, tryon_image_path):
+    """Put background back on the person image after tryon."""
+    person = np.array(Image.open(person_image_path))
+    # tryon image
+    tryon = np.array(Image.open(tryon_image_path))
+    # persom image mask from rembg
+    p_mask = np.array(img_mask)
+    # make binary mask
+    p_mask = np.where(p_mask>0, 1, 0)
+    # invert mask
+    p_mask_inv = np.logical_not(p_mask)
+    # make bg without person
+    background = person * np.stack((p_mask_inv, p_mask_inv, p_mask_inv), axis=2)
+    # make tryon image without background
+    tryon_nobg = tryon * np.stack((p_mask, p_mask, p_mask), axis=2)
+    tryon_nobg = tryon_nobg.astype("uint8")
+    # composite 
+    tryon_with_bg = np.add(tryon_nobg, background)
+    tryon_with_bg_pil = Image.fromarray(np.uint8(tryon_with_bg)).convert('RGB')
+    tryon_with_bg_pil.save("results/test/try-on/tryon_with_bg.png")
+
+
 # Main inference function
-def inference(clothing_image, person_image, remove_bg):
+def inference(clothing_image, person_image, remove_bg, retrieve_bg):
     """
     Do try-on!
     """
@@ -101,6 +124,8 @@ def inference(clothing_image, person_image, remove_bg):
         img = remove(img, alpha_matting=True, alpha_matting_erode_size=15)
         print("Removing background from person image..")
     img = ImageOps.fit(img, (192, 256), Image.BICUBIC).convert("RGB")
+    # Get binary from person image, used in def_composite_background
+    img_mask = remove(img, alpha_matting=True, alpha_matting_erode_size=15, only_mask=True)
     img_path = os.path.join('Data_preprocessing/test_img', img_name)
     img.save(img_path)
     resize_time = time.time()
@@ -126,7 +151,15 @@ def inference(clothing_image, person_image, remove_bg):
     os.system("python test.py --name decavtonfifapretrain")
     tryon_image = Image.open("results/test/try-on/person.png")
     print("Size of image is: ", tryon_image.size)
-    return os.path.join("results/test/try-on/person.png")
+    
+    # Return try-on with background added back on the person image
+    if retrieve_bg == "yes":
+        composite_background(img_mask, 'Data_preprocessing/test_img/person.png',
+                     'results/test/try-on/person.png')
+        return os.path.join("results/test/try-on/tryon_with_bg.png")
+    # Return only try-on result
+    else:
+        return os.path.join("results/test/try-on/person.png")
 
 
 title = "Virtual Dressing Room"
@@ -136,9 +169,10 @@ thumbnail = None # "./pathtothumbnail.png"
 
 gr.Interface(
     inference,
-    [gr.inputs.Image(type='filepath', label="Clothing Image"), 
+    [gr.inputs.Image(type='filepath', label="Clothing Image"),
      gr.inputs.Image(type='filepath', label="Person Image"),
-     gr.inputs.Radio(choices=["yes","no"], default="no", label="Remove background from person image?")],
+     gr.inputs.Radio(choices=["yes","no"], default="no", label="Remove background from the person image?"),
+     gr.inputs.Radio(choices=["yes","no"], default="no", label="Retrieve original background from the person image?")],
     gr.outputs.Image(type="filepath", label="Predicted Output"),
     examples=[["./examples/1/cloth.jpg", "./examples/1/person.jpg"],
               ["./examples/2/cloth.jpg", "./examples/2/person.jpg"]],