Update model/cloth_masker.py

model/cloth_masker.py

@@ -84,6 +84,8 @@ MASK_CLOTH_PARTS = {
     'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits'],
     'inner': ['Upper-clothes'],
     'outer': ['Coat',]
+    'bags': ['Bag'],  # New category for bags
+    'footwear': ['Left-shoe', 'Right-shoe'],  # New category for footwear
 }
 MASK_DENSE_PARTS = {
     'upper': ['torso', 'big arms', 'forearms'],
@@ -91,6 +93,8 @@ MASK_DENSE_PARTS = {
     'overall': ['torso', 'thighs', 'legs', 'big arms', 'forearms'],
     'inner': ['torso'],
     'outer': ['torso', 'big arms', 'forearms']
+    'bags': [],  # No dense parts for bags
+    'footwear': ['left foot', 'right foot'],  # New category for footwear
 }
     
 schp_public_protect_parts = ['Hat', 'Hair', 'Sunglasses', 'Left-shoe', 'Right-shoe', 'Bag', 'Glove', 'Scarf']
@@ -186,45 +190,44 @@ class AutoMasker:
     
     @staticmethod
     def cloth_agnostic_mask(
-        densepose_mask: Image.Image,
-        schp_lip_mask: Image.Image,
-        schp_atr_mask: Image.Image,
-        part: str='overall',
-        **kwargs
-    ):
-        assert part in ['upper', 'lower', 'overall', 'inner', 'outer'], f"part should be one of ['upper', 'lower', 'overall', 'inner', 'outer'], but got {part}"
-        w, h = densepose_mask.size
-        
-        dilate_kernel = max(w, h) // 250
-        dilate_kernel = dilate_kernel if dilate_kernel % 2 == 1 else dilate_kernel + 1
-        dilate_kernel = np.ones((dilate_kernel, dilate_kernel), np.uint8)
-        
-        kernal_size = max(w, h) // 15
-        kernal_size = kernal_size if kernal_size % 2 == 1 else kernal_size + 1
-        
-        densepose_mask = np.array(densepose_mask)
-        schp_lip_mask = np.array(schp_lip_mask)
-        schp_atr_mask = np.array(schp_atr_mask)
-        
-        # Strong Protect Area (Hands, Face, Accessory, Feet)
-        hands_protect_area = part_mask_of(['hands', 'feet'], densepose_mask, DENSE_INDEX_MAP)
-        hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
-        hands_protect_area = hands_protect_area & \
-            (part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) | \
-             part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_lip_mask, LIP_MAPPING))
-        face_protect_area = part_mask_of('Face', schp_lip_mask, LIP_MAPPING)
+    densepose_mask: Image.Image,
+    schp_lip_mask: Image.Image,
+    schp_atr_mask: Image.Image,
+    part: str='overall',
+    **kwargs
+):
+    assert part in ['upper', 'lower', 'bags', 'footwear', 'inner', 'outer'], f"part should be one of ['upper', 'lower', 'bags', 'footwear', 'inner', 'outer'], but got {part}"
+    w, h = densepose_mask.size
+    
+    dilate_kernel = max(w, h) // 250
+    dilate_kernel = dilate_kernel if dilate_kernel % 2 == 1 else dilate_kernel + 1
+    dilate_kernel = np.ones((dilate_kernel, dilate_kernel), np.uint8)
+    
+    kernal_size = max(w, h) // 15
+    kernal_size = kernal_size if kernal_size % 2 == 1 else kernal_size + 1
+    
+    densepose_mask = np.array(densepose_mask)
+    schp_lip_mask = np.array(schp_lip_mask)
+    schp_atr_mask = np.array(schp_atr_mask)
+    
+    # Strong Protect Area (Hands, Face, Accessory, Feet)
+    hands_protect_area = part_mask_of(['hands', 'feet'], densepose_mask, DENSE_INDEX_MAP)
+    hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
+    hands_protect_area = hands_protect_area & \
+        (part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) | \
+         part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_lip_mask, LIP_MAPPING))
+    face_protect_area = part_mask_of('Face', schp_lip_mask, LIP_MAPPING)
+
+    strong_protect_area = hands_protect_area | face_protect_area 
 
-        strong_protect_area = hands_protect_area | face_protect_area 
+    # Weak Protect Area (Hair, Irrelevant Clothes, Body Parts)
+    body_protect_area = part_mask_of(PROTECT_BODY_PARTS[part], schp_lip_mask, LIP_MAPPING) | part_mask_of(PROTECT_BODY_PARTS[part], schp_atr_mask, ATR_MAPPING)
+    hair_protect_area = part_mask_of(['Hair'], schp_lip_mask, LIP_MAPPING) | \
+        part_mask_of(['Hair'], schp_atr_mask, ATR_MAPPING)
+    cloth_protect_area = part_mask_of(PROTECT_CLOTH_PARTS[part]['LIP'], schp_lip_mask, LIP_MAPPING) | \
+        part_mask_of(PROTECT_CLOTH_PARTS[part]['ATR'], schp_atr_mask, ATR_MAPPING)
+    accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses
 
-        # Weak Protect Area (Hair, Irrelevant Clothes, Body Parts)
-        body_protect_area = part_mask_of(PROTECT_BODY_PARTS[part], schp_lip_mask, LIP_MAPPING) | part_mask_of(PROTECT_BODY_PARTS[part], schp_atr_mask, ATR_MAPPING)
-        hair_protect_area = part_mask_of(['Hair'], schp_lip_mask, LIP_MAPPING) | \
-            part_mask_of(['Hair'], schp_atr_mask, ATR_MAPPING)
-        cloth_protect_area = part_mask_of(PROTECT_CLOTH_PARTS[part]['LIP'], schp_lip_mask, LIP_MAPPING) | \
-            part_mask_of(PROTECT_CLOTH_PARTS[part]['ATR'], schp_atr_mask, ATR_MAPPING)
-        accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks']), schp_lip_mask, LIP_MAPPING) | \
-            part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING) 
-        weak_protect_area = body_protect_area | cloth_protect_area | hair_protect_area | strong_protect_area | accessory_protect_area
         
         # Mask Area
         strong_mask_area = part_mask_of(MASK_CLOTH_PARTS[part], schp_lip_mask, LIP_MAPPING) | \