Update app.py

app.py

@@ -19,7 +19,6 @@ from preprocess.humanparsing.run_parsing import Parsing
 from preprocess.openpose.run_openpose import OpenPose
 from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation
 from torchvision.transforms.functional import to_pil_image
-from datasets import Dataset
 
 
 def pil_to_binary_mask(pil_image, threshold=0):
@@ -85,7 +84,7 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU
-def start_tryon(dict,garm_img,garment_des,cloth_type,is_checked,is_checked_crop,denoise_steps,seed):
+def start_tryon(img,garm_img,garment_des,cloth_type,is_checked,is_checked_crop,denoise_steps,seed):
     device = "cuda"
     
     openpose_model.preprocessor.body_estimation.model.to(device)
@@ -93,7 +92,7 @@ def start_tryon(dict,garm_img,garment_des,cloth_type,is_checked,is_checked_crop,
     pipe.unet_encoder.to(device)
 
     garm_img= garm_img.convert("RGB").resize((768,1024))
-    human_img_orig = dict["background"].convert("RGB")    
+    human_img_orig = img.copy() #dict["background"].convert("RGB")    
     
     if is_checked_crop:
         width, height = human_img_orig.size
@@ -116,7 +115,7 @@ def start_tryon(dict,garm_img,garment_des,cloth_type,is_checked,is_checked_crop,
         mask, mask_gray = get_mask_location('hd', cloth_type, model_parse, keypoints)
         mask = mask.resize((768,1024))
     else:
-        mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
+        mask = pil_to_binary_mask(img.resize((768, 1024)))
         # mask = transforms.ToTensor()(mask)
         # mask = mask.unsqueeze(0)
     mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
@@ -204,8 +203,6 @@ def main_(imgs,topwear_img,topwear_des,bottomwear_img,bottomwear_des,dress_img,d
         return start_tryon(imgs,bottomwear_img,bottomwear_des,"lower_body",is_checked,is_checked_crop,denoise_steps,seed)
     elif topwear_img!=None and bottomwear_img!=None:
         _, half_img, half_mask = start_tryon(imgs,topwear_img,topwear_des,"upper_body",is_checked,is_checked_crop,denoise_steps,seed)
-        IMG_data = [{"image": half_img}]  
-        img_dataset = Dataset.from_dict(IMG_data)
         return start_tryon(img_dataset,bottomwear_img,bottomwear_des,"lower_body",is_checked,is_checked_crop,denoise_steps,seed)
     
         
@@ -232,7 +229,8 @@ with image_blocks as demo:
     gr.HTML("<center><p>Upload an image of a person and images of the clothes✨</p></center>")
     with gr.Row():
         with gr.Column():
-            imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
+            inp_img = gr.ImageEditor(sources='upload',label='Person_Image', type="pil", label='Human. Mask with pen or use auto-masking',
+                                     image_mode='RGB', layers=False, interactive=True)
             with gr.Row():
                 is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
             with gr.Row():
@@ -276,7 +274,7 @@ with image_blocks as demo:
                 seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
 
 
-    try_button.click(fn=main_, inputs=[imgs,topwear_image,topwear_desc,bottomwear_image,bottomwear_desc,dress_image,dress_desc,is_checked,is_checked_crop,denoise_steps,seed],
+    try_button.click(fn=main_, inputs=[inp_img,topwear_image,topwear_desc,bottomwear_image,bottomwear_desc,dress_image,dress_desc,is_checked,is_checked_crop,denoise_steps,seed],
                      outputs=[image_in, image_out, masked_img], api_name='tryon')