Spaces:
Running
on
Zero
Running
on
Zero
Update utils_mask.py
Browse files- utils_mask.py +13 -3
utils_mask.py
CHANGED
@@ -77,6 +77,7 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
|
|
77 |
arms_left = (parse_array == 14).astype(np.float32)
|
78 |
arms_right = (parse_array == 15).astype(np.float32)
|
79 |
|
|
|
80 |
if category == 'dresses':
|
81 |
# Initial dress mask for the upper body
|
82 |
parse_mask_upper = np.logical_or((parse_array == label_map["upper_clothes"]), (parse_array == label_map["dress"])).astype(np.float32)
|
@@ -91,10 +92,17 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
|
|
91 |
parse_array == label_map["right_leg"])).astype(np.float32)
|
92 |
|
93 |
# Dilate the leg mask to ensure coverage and fill gaps
|
94 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
95 |
|
96 |
-
# Combine the upper body mask with the leg mask
|
97 |
-
parse_mask = np.maximum(parse_mask_upper, parse_mask_legs)
|
98 |
|
99 |
|
100 |
elif category == 'upper_body':
|
@@ -176,3 +184,5 @@ def get_mask_location(model_type, category, model_parse: Image.Image, keypoint:
|
|
176 |
mask_gray = Image.fromarray(inpaint_mask.astype(np.uint8) * 127)
|
177 |
|
178 |
return mask, mask_gray
|
|
|
|
|
|
77 |
arms_left = (parse_array == 14).astype(np.float32)
|
78 |
arms_right = (parse_array == 15).astype(np.float32)
|
79 |
|
80 |
+
|
81 |
if category == 'dresses':
|
82 |
# Initial dress mask for the upper body
|
83 |
parse_mask_upper = np.logical_or((parse_array == label_map["upper_clothes"]), (parse_array == label_map["dress"])).astype(np.float32)
|
|
|
92 |
parse_array == label_map["right_leg"])).astype(np.float32)
|
93 |
|
94 |
# Dilate the leg mask to ensure coverage and fill gaps
|
95 |
+
parse_mask_legs_dilated = cv2.dilate(parse_mask_legs.astype(np.uint8), np.ones((6, 6), np.uint8), iterations=6)
|
96 |
+
|
97 |
+
# Combine the upper body mask with the dilated leg mask
|
98 |
+
parse_mask_total = np.maximum(parse_mask_upper, parse_mask_legs_dilated)
|
99 |
+
|
100 |
+
# Apply further processing as needed (e.g., additional masks or refinements)
|
101 |
+
|
102 |
+
parse_mask = np.logical_and(parser_mask_changeable, np.logical_not(parse_mask_total))
|
103 |
+
|
104 |
+
|
105 |
|
|
|
|
|
106 |
|
107 |
|
108 |
elif category == 'upper_body':
|
|
|
184 |
mask_gray = Image.fromarray(inpaint_mask.astype(np.uint8) * 127)
|
185 |
|
186 |
return mask, mask_gray
|
187 |
+
|
188 |
+
|