Update model/cloth_masker.py
Browse files- model/cloth_masker.py +10 -10
model/cloth_masker.py
CHANGED
|
@@ -207,7 +207,10 @@ class AutoMasker:
|
|
| 207 |
schp_atr_mask = np.array(schp_atr_mask)
|
| 208 |
|
| 209 |
# Strong Protect Area (Hands, Face, Accessory, Feet)
|
| 210 |
-
|
|
|
|
|
|
|
|
|
|
| 211 |
hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
|
| 212 |
hands_protect_area = hands_protect_area & \
|
| 213 |
(part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) | \
|
|
@@ -222,8 +225,12 @@ class AutoMasker:
|
|
| 222 |
part_mask_of(['Hair'], schp_atr_mask, ATR_MAPPING)
|
| 223 |
cloth_protect_area = part_mask_of(PROTECT_CLOTH_PARTS[part]['LIP'], schp_lip_mask, LIP_MAPPING) | \
|
| 224 |
part_mask_of(PROTECT_CLOTH_PARTS[part]['ATR'], schp_atr_mask, ATR_MAPPING)
|
| 225 |
-
|
| 226 |
-
part_mask_of(accessory_parts,
|
|
|
|
|
|
|
|
|
|
|
|
|
| 227 |
weak_protect_area = body_protect_area | cloth_protect_area | hair_protect_area | strong_protect_area | accessory_protect_area
|
| 228 |
|
| 229 |
# Mask Area
|
|
@@ -237,14 +244,7 @@ class AutoMasker:
|
|
| 237 |
|
| 238 |
|
| 239 |
mask_area = (np.ones_like(densepose_mask) & (~weak_protect_area) & (~background_area)) | mask_dense_area
|
| 240 |
-
|
| 241 |
-
mask_area = hull_mask(mask_area * 255) // 255 # Convex Hull to expand the mask area
|
| 242 |
-
mask_area = mask_area & (~weak_protect_area)
|
| 243 |
-
mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
|
| 244 |
-
mask_area[mask_area < 25] = 0
|
| 245 |
-
mask_area[mask_area >= 25] = 1
|
| 246 |
mask_area = (mask_area | strong_mask_area) & (~strong_protect_area)
|
| 247 |
-
mask_area = cv2.dilate(mask_area, dilate_kernel, iterations=1)
|
| 248 |
|
| 249 |
return Image.fromarray(mask_area * 255)
|
| 250 |
|
|
|
|
| 207 |
schp_atr_mask = np.array(schp_atr_mask)
|
| 208 |
|
| 209 |
# Strong Protect Area (Hands, Face, Accessory, Feet)
|
| 210 |
+
if part == "overall":
|
| 211 |
+
hands_protect_area = part_mask_of(['hands',], densepose_mask, DENSE_INDEX_MAP)
|
| 212 |
+
else:
|
| 213 |
+
hands_protect_area = part_mask_of(['hands', 'feet'], densepose_mask, DENSE_INDEX_MAP)
|
| 214 |
hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
|
| 215 |
hands_protect_area = hands_protect_area & \
|
| 216 |
(part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) | \
|
|
|
|
| 225 |
part_mask_of(['Hair'], schp_atr_mask, ATR_MAPPING)
|
| 226 |
cloth_protect_area = part_mask_of(PROTECT_CLOTH_PARTS[part]['LIP'], schp_lip_mask, LIP_MAPPING) | \
|
| 227 |
part_mask_of(PROTECT_CLOTH_PARTS[part]['ATR'], schp_atr_mask, ATR_MAPPING)
|
| 228 |
+
if part == "overall":
|
| 229 |
+
accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Scarf', 'Socks']), schp_lip_mask, LIP_MAPPING) | \
|
| 230 |
+
part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING)
|
| 231 |
+
else:
|
| 232 |
+
accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks']), schp_lip_mask, LIP_MAPPING) | \
|
| 233 |
+
part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING)
|
| 234 |
weak_protect_area = body_protect_area | cloth_protect_area | hair_protect_area | strong_protect_area | accessory_protect_area
|
| 235 |
|
| 236 |
# Mask Area
|
|
|
|
| 244 |
|
| 245 |
|
| 246 |
mask_area = (np.ones_like(densepose_mask) & (~weak_protect_area) & (~background_area)) | mask_dense_area
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 247 |
mask_area = (mask_area | strong_mask_area) & (~strong_protect_area)
|
|
|
|
| 248 |
|
| 249 |
return Image.fromarray(mask_area * 255)
|
| 250 |
|