patrickligardes commited on
Commit
13c380f
β€’
1 Parent(s): e7828cd

Update utils_mask.py

Browse files
Files changed (1) hide show
  1. utils_mask.py +13 -9
utils_mask.py CHANGED
@@ -79,22 +79,26 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
79
 
80
  if category == 'dresses':
81
  parse_mask = (parse_array == 7).astype(np.float32) + \
82
- (parse_array == 4).astype(np.float32) + \
83
- (parse_array == 5).astype(np.float32) + \
84
- (parse_array == 6).astype(np.float32)
85
 
86
  parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
87
-
88
- # Create a separate mask for the gap between legs
89
- leg_gap_mask = (parse_array == 12).astype(np.float32) + \
90
- (parse_array == 13).astype(np.float32)
91
 
92
  # Fill gaps between legs
93
- leg_gap_filled = cv2.dilate(leg_gap_mask.astype(np.uint8), np.ones((10, 10), np.uint8), iterations=10)
94
 
95
- # Add the gap-filled mask to the existing parse_mask
 
 
 
96
  parse_mask += leg_gap_filled
97
 
 
98
  elif category == 'upper_body':
99
  parse_mask = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
100
  parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \
 
79
 
80
  if category == 'dresses':
81
  parse_mask = (parse_array == 7).astype(np.float32) + \
82
+ (parse_array == 4).astype(np.float32) + \
83
+ (parse_array == 5).astype(np.float32) + \
84
+ (parse_array == 6).astype(np.float32)
85
 
86
  parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
87
+
88
+ # Create a separate mask for the leg parts
89
+ leg_mask = (parse_array == 12).astype(np.float32) + \
90
+ (parse_array == 13).astype(np.float32)
91
 
92
  # Fill gaps between legs
93
+ leg_gap_filled = cv2.dilate(leg_mask.astype(np.uint8), np.ones((10, 10), np.uint8), iterations=10)
94
 
95
+ # Ensure the filled leg gaps do not affect the upper body
96
+ leg_gap_filled = np.logical_and(leg_gap_filled, leg_mask).astype(np.float32)
97
+
98
+ # Add the filled leg gaps to the parse_mask
99
  parse_mask += leg_gap_filled
100
 
101
+
102
  elif category == 'upper_body':
103
  parse_mask = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
104
  parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \