using crop if original image AR is not 2x3

app.py

@@ -124,18 +124,26 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 # Standard size of shein images
-WIDTH = int(4160/5)
-HEIGHT = int(6240/5)
+#WIDTH = int(4160/5)
+#HEIGHT = int(6240/5)
 # Standard size on which model is trained
-#WIDTH = int(768)
-#HEIGHT = int(1024)
+WIDTH = int(768)
+HEIGHT = int(1024)
 POSE_WIDTH = int(WIDTH/2)  # int(WIDTH/2)
 POSE_HEIGHT = int(HEIGHT/2)  #int(HEIGHT/2)
 
 CATEGORY = "upper_body" # "lower_body"
 
+def is_cropping_required(width, height):
+    # If aspect ratio is 1.5, which is same as standard 2x3 ( 768x1024 ), then no need to crop, else crop
+    aspect_ratio = round(height/width, 2)
+    if aspect_ratio == 2:
+        return False
+    return True
+
+
 @spaces.GPU
-def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked_crop,denoise_steps,seed):
+def start_tryon(human_img_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'
     
@@ -143,9 +151,12 @@ def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked
     pipe.to(device)
     pipe.unet_encoder.to(device)
     
-    #human_img_orig = dict["background"].convert("RGB")   # ImageEditor
-    human_img_orig = dict.convert("RGB")     # Image
+    if isinstance(human_img_dict, dict):        
+        human_img_orig = human_img_dict["background"].convert("RGB")   # ImageEditor
+    else:
+        human_img_orig = dict.convert("RGB")     # Image
 
+    """
     # Derive HEIGHT & WIDTH such that width is not more than 1000. This will cater to both Shein images (4160x6240) of 3:4 AR and model standard images ( 768x1024 ) of 2:3 AR
     WIDTH, HEIGHT = human_img_orig.size
     division_factor = math.ceil(WIDTH/1000)
@@ -153,9 +164,14 @@ def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked
     HEIGHT = int(HEIGHT/division_factor)
     POSE_WIDTH = int(WIDTH/2)
     POSE_HEIGHT = int(HEIGHT/2)
+    """
+    # is_checked_crop as True if original AR is not same as 2x3 as expected by model
+    w, h = human_img_orig.size
+    is_checked_crop = is_cropping_required(w, h)
 
     garm_img= garm_img.convert("RGB").resize((WIDTH,HEIGHT))
     if is_checked_crop:
+        # This will crop the image to make it Aspect Ratio of 3 x 4. And then at the end revert it back to original dimentions
         width, height = human_img_orig.size
         target_width = int(min(width, height * (3 / 4)))
         target_height = int(min(height, width * (4 / 3)))
@@ -179,7 +195,7 @@ def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked
         mask = mask.resize((WIDTH, HEIGHT))
         logging.info("Mask location on model identified")
     else:
-        mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((WIDTH, HEIGHT)))
+        mask = pil_to_binary_mask(human_img_dict['layers'][0].convert("RGB").resize((WIDTH, HEIGHT)))
         # mask = transforms.ToTensor()(mask)
         # mask = mask.unsqueeze(0)
     mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
@@ -191,7 +207,7 @@ def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked
      
     
 
-    args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
+    args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', device))
     # verbosity = getattr(args, "verbosity", None)
     pose_img = args.func(args,human_img_arg)    
     pose_img = pose_img[:,:,::-1]    
@@ -282,15 +298,15 @@ human_list = os.listdir(os.path.join(example_path,"human"))
 human_list_path = [os.path.join(example_path,"human",human) for human in human_list]
 
 human_ex_list = []
-human_ex_list = human_list_path # Image
-""" if using ImageEditor instead of Image while taking input, use this - ImageEditor
+#human_ex_list = human_list_path # Image
+#""" if using ImageEditor instead of Image while taking input, use this - ImageEditor
 for ex_human in human_list_path:
     ex_dict= {}
     ex_dict['background'] = ex_human
     ex_dict['layers'] = None
     ex_dict['composite'] = None
     human_ex_list.append(ex_dict)
-"""
+#"""
 ##default human
 
 
@@ -303,8 +319,8 @@ with image_blocks as demo:
         with gr.Column():
             # changing from ImageEditor to Image to allow easy passing of data through API
             # instead of passing {"dictionary": <>} ( which is failing ), we can directly pass the image
-            #imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
-            imgs = gr.Image(sources='upload', type='pil',label='Human. Mask with pen or use auto-masking')
+            imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
+            #imgs = gr.Image(sources='upload', type='pil',label='Human. Mask with pen or use auto-masking')
             with gr.Row():
                 is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
             with gr.Row():