adding option to pass background image for adding background

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
 from src.unet_hacked_tryon import UNet2DConditionModel
+from src.background_processor import BackgroundProcessor
 from transformers import (
     CLIPImageProcessor,
     CLIPVisionModelWithProjection,
@@ -131,7 +132,7 @@ POSE_HEIGHT = int(HEIGHT/2)  #int(HEIGHT/2)
 CATEGORY = "upper_body" # "lower_body"
 
 @spaces.GPU
-def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed):
+def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked_crop,denoise_steps,seed):
     device = "cuda"
     # device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
     
@@ -249,9 +250,17 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
     if is_checked_crop:
         out_img = images[0].resize(crop_size)        
         human_img_orig.paste(out_img, (int(left), int(top)))    
-        return human_img_orig, mask_gray
+        final_image = human_img_orig
+        # return human_img_orig, mask_gray
     else:
-        return images[0], mask_gray
+        final_image = images[0]
+        # return images[0], mask_gray
+    
+    # apply background to final image
+    if background_img:
+       logging.info("Adding background")
+       final_image = BackgroundProcessor.add_background(final_image, background_img) 
+    return final_image, mask_gray
     # return images[0], mask_gray
 
 garm_list = os.listdir(os.path.join(example_path,"cloth"))
@@ -298,13 +307,23 @@ with image_blocks as demo:
                 inputs=garm_img,
                 examples_per_page=8,
                 examples=garm_list_path)
+            
+        with gr.Column():
+           background_img = gr.Image(label="Background", sources='upload', type="pil")
+
+        with gr.Column():
+            with gr.Row():
+               image_out = gr.Image(label="Output", elem_id="output-img", show_share_button=False) 
+            with gr.Row():
+               masked_img = gr.Image(label="Masked image output", elem_id="masked-img", show_share_button=False) 
+        """
         with gr.Column():
             # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
             masked_img = gr.Image(label="Masked image output", elem_id="masked-img", show_share_button=False)
-        with gr.Column():
+        with gr.Column():            
             # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
             image_out = gr.Image(label="Output", elem_id="output-img", show_share_button=False)
-
+        """
 
 
 
@@ -317,7 +336,7 @@ with image_blocks as demo:
 
 
 
-    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
+    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, background_img, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
 
             
 

src/background_processor.py

@@ -0,0 +1,38 @@
+from PIL import Image
+import numpy as np
+from preprocess.humanparsing.run_parsing import Parsing
+
+parsing_model = Parsing(0)
+
+class BackgroundProcessor:
+    @classmethod
+    def add_background(cls, human_img: Image, background_img: Image):
+        
+        human_img = human_img.convert("RGB")
+        width = human_img.width
+        height = human_img.height
+        
+        # Create mask image
+        parsed_img, _ = parsing_model(human_img)
+        mask_img = parsed_img.convert("L")
+        mask_img = mask_img.resize((width, height))
+        
+        background_img = background_img.convert("RGB")
+        background_img = background_img.resize((width, height))
+
+        # Convert to numpy arrays
+        human_np = np.array(human_img)
+        mask_np = np.array(mask_img)
+        background_np = np.array(background_img)
+
+        # Ensure mask is 3-channel (RGB) for compatibility
+        mask_np = np.stack((mask_np,) * 3, axis=-1)
+
+        # Apply the mask to human_img
+        human_with_background = np.where(mask_np > 0, human_np, background_np)
+
+        # Convert back to PIL Image
+        result_img = Image.fromarray(human_with_background.astype('uint8'))
+
+        # Return or save the result
+        return result_img