Update utils_mask.py

utils_mask.py

@@ -152,11 +152,27 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
     parser_mask_fixed = np.logical_or(parser_mask_fixed, parse_head)
     parse_mask = cv2.dilate(parse_mask, np.ones((5, 5), np.uint16), iterations=5)
     if category == 'dresses' or category == 'upper_body':
+        # Create a mask for the neck
         neck_mask = (parse_array == 18).astype(np.float32)
         neck_mask = cv2.dilate(neck_mask, np.ones((5, 5), np.uint16), iterations=1)
         neck_mask = np.logical_and(neck_mask, np.logical_not(parse_head))
+        
+        # Add padding between the legs
+        legs_mask = (parse_array == 19).astype(np.float32)
+        legs_mask = cv2.dilate(legs_mask, np.ones((5, 5), np.uint16), iterations=1)
+        legs_padding_size = 10  # Adjust the padding size as needed
+        legs_padding = cv2.dilate(legs_mask, np.ones((legs_padding_size, legs_padding_size), np.uint16), iterations=1)
+        
+        # Combine the legs mask with padding with existing masks
+        parse_mask += legs_padding
+        
+        # Combine the neck mask with existing masks
         parse_mask = np.logical_or(parse_mask, neck_mask)
+        
+        # Dilate the arm mask
         arm_mask = cv2.dilate(np.logical_or(im_arms_left, im_arms_right).astype('float32'), np.ones((5, 5), np.uint16), iterations=4)
+        
+        # Combine the arm mask with existing masks
         parse_mask += np.logical_or(parse_mask, arm_mask)
 
     parse_mask = np.logical_and(parser_mask_changeable, np.logical_not(parse_mask))