FIT-Check commited on
Commit
0459816
·
verified ·
1 Parent(s): 231818c

Update model/cloth_masker.py

Browse files
Files changed (1) hide show
  1. 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
- hands_protect_area = part_mask_of(['hands', 'feet'], densepose_mask, DENSE_INDEX_MAP)
 
 
 
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
- accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks']), schp_lip_mask, LIP_MAPPING) | \
226
- part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING)
 
 
 
 
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