Update utils_mask.py

utils_mask.py

@@ -51,7 +51,7 @@ def refine_mask(mask):
 
     return refine_mask
 
-def get_mask_location(model_type, category, model_parse: Image.Image, keypoint: dict, width=384,height=512):
+def get_mask_location(model_type, category, model_parse: Image.Image, keypoint: dict, width=384, height=512):
     im_parse = model_parse.resize((width, height), Image.NEAREST)
     parse_array = np.array(im_parse)
 
@@ -60,7 +60,7 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
     elif model_type == 'dc':
         arm_width = 45
     else:
-        raise ValueError("model_type must be \'hd\' or \'dc\'!")
+        raise ValueError("model_type must be 'hd' or 'dc'!")
 
     parse_head = (parse_array == 1).astype(np.float32) + \
                  (parse_array == 3).astype(np.float32) + \
@@ -82,7 +82,6 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
                      (parse_array == 4).astype(np.float32) + \
                      (parse_array == 5).astype(np.float32) + \
                      (parse_array == 6).astype(np.float32)
-
         parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
 
     elif category == 'upper_body':
@@ -91,6 +90,7 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
                                         (parse_array == label_map["pants"]).astype(np.float32)
         parser_mask_fixed += parser_mask_fixed_lower_cloth
         parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+
     elif category == 'lower_body':
         parse_mask = (parse_array == 6).astype(np.float32) + \
                      (parse_array == 12).astype(np.float32) + \
@@ -100,9 +100,29 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
                              (parse_array == 14).astype(np.float32) + \
                              (parse_array == 15).astype(np.float32)
         parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+
+    elif category == 'full_body':
+        # Combine lower_body and upper_body
+        parse_mask_upper = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
+        parse_mask_lower = (parse_array == 6).astype(np.float32) + \
+                           (parse_array == 12).astype(np.float32) + \
+                           (parse_array == 13).astype(np.float32) + \
+                           (parse_array == 5).astype(np.float32)
+        parse_mask = parse_mask_upper + parse_mask_lower
+        
+        parser_mask_fixed += (parse_array == label_map["upper_clothes"]).astype(np.float32) + \
+                             (parse_array == label_map["skirt"]).astype(np.float32) + \
+                             (parse_array == label_map["pants"]).astype(np.float32) + \
+                             (parse_array == 14).astype(np.float32) + \
+                             (parse_array == 15).astype(np.float32)
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+
     else:
         raise NotImplementedError
 
+    # Rest of the function logic remains the same...
+
+
     # Load pose points
     pose_data = keypoint["pose_keypoints_2d"]
     pose_data = np.array(pose_data)