Update utils_mask.py

utils_mask.py

@@ -77,25 +77,32 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
     arms_left = (parse_array == 14).astype(np.float32)
     arms_right = (parse_array == 15).astype(np.float32)
 
+
     if category == 'dresses':
         # Initial dress mask for the upper body
         parse_mask_upper = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
-        parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \
-                                        (parse_array == label_map["pants"]).astype(np.float32)
-        parser_mask_fixed += parser_mask_fixed_lower_cloth
-
+        
         # Create a mask for the legs (including skirts and pants)
         parse_mask_legs = np.logical_or.reduce((parse_array == label_map["skirt"],
                                                 parse_array == label_map["pants"],
                                                 parse_array == label_map["left_leg"],
                                                 parse_array == label_map["right_leg"])).astype(np.float32)
-
+        
         # Dilate the leg mask to ensure coverage and fill gaps
         parse_mask_legs = cv2.dilate(parse_mask_legs.astype(np.uint8), np.ones((6, 6), np.uint8), iterations=6)
+        
+        # Include legs mask in the fixed lower clothing mask
+        parser_mask_fixed_lower_cloth = np.logical_or(parse_mask_legs, 
+                                                    (parse_array == label_map["skirt"]).astype(np.float32) +
+                                                    (parse_array == label_map["pants"]).astype(np.float32)).astype(np.float32)
+
+        # Add the fixed lower clothing mask to parser_mask_fixed
+        parser_mask_fixed += parser_mask_fixed_lower_cloth
 
         # Combine the upper body mask with the leg mask
         parse_mask = np.maximum(parse_mask_upper, parse_mask_legs)
 
+
     elif category == 'upper_body':
         parse_mask = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
         parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \