Update model/cloth_masker.py

model/cloth_masker.py

@@ -80,10 +80,10 @@ PROTECT_CLOTH_PARTS = {
 }
 MASK_CLOTH_PARTS = {
     'upper': ['Upper-clothes', 'Coat', 'Dress', 'Jumpsuits'],
-    'lower': ['Pants', 'Skirt', 'Dress', 'Jumpsuits'],
+    'lower': ['Pants', 'Skirt', 'Dress', ' Jumpsuits'],
     'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits'],
     'inner': ['Upper-clothes'],
-    'outer': ['Coat',],
+    'outer': ['Coat'],
     'bags': ['Bag'],  # New category for bags
     'footwear': ['Left-shoe', 'Right-shoe']  # New category for footwear
 }
@@ -92,11 +92,11 @@ MASK_DENSE_PARTS = {
     'lower': ['thighs', 'legs'],
     'overall': ['torso', 'thighs', 'legs', 'big arms', 'forearms'],
     'inner': ['torso'],
-    'outer': ['torso', 'big arms', 'forearms']
+    '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']
 schp_protect_parts = {
     'upper': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits'],  
@@ -110,7 +110,7 @@ schp_mask_parts = {
     'lower': ['Pants', 'Skirt', 'Dress', 'Jumpsuits', 'socks'],
     'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits', 'socks'],
     'inner': ['Upper-clothes'],
-    'outer': ['Coat',]
+    'outer': ['Coat']
 }
 
 dense_mask_parts = {
@@ -154,7 +154,6 @@ def hull_mask(mask_area: np.ndarray):
         hull = cv2.convexHull(c)
         hull_mask = cv2.fillPoly(np.zeros_like(mask_area), [hull], 255) | hull_mask
     return hull_mask
-    
 
 class AutoMasker:
     def __init__(
@@ -168,7 +167,7 @@ class AutoMasker:
         
         self.densepose_processor = DensePose(densepose_ckpt, device)
         self.schp_processor_atr = SCHP(ckpt_path=os.path.join(schp_ckpt, 'exp-schp-201908301523-atr.pth'), device=device)
-        self.schp_processor_lip = SCHP(ckpt_path=os.path.join(schp_ckpt, 'exp-schp-201908261155-lip.pth'), device=device)
+        self.schp_processor_lip = SCHP(ckpt_path =os.path.join(schp_ckpt, 'exp-schp-201908261155-lip.pth'), device=device)
         
         self.mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
 
@@ -190,45 +189,42 @@ 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', '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 
+        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)
 
-    # 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'])) | \
-        part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING)
+        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)
         
         # Mask Area
         strong_mask_area = part_mask_of(MASK_CLOTH_PARTS[part], schp_lip_mask, LIP_MAPPING) | \
@@ -237,13 +233,12 @@ class AutoMasker:
         mask_dense_area = part_mask_of(MASK_DENSE_PARTS[part], densepose_mask, DENSE_INDEX_MAP)
         mask_dense_area = cv2.resize(mask_dense_area.astype(np.uint8), None, fx=0.25, fy=0.25, interpolation=cv2.INTER_NEAREST)
         mask_dense_area = cv2.dilate(mask_dense_area, dilate_kernel, iterations=2)
-        mask_dense_area = cv2.resize(mask_dense_area.astype(np.uint8), None, fx=4, fy=4, interpolation=cv2.INTER_NEAREST)
+        mask_dense_area = cv2.resize(mask_dense_area .astype(np.uint8), None, fx=4, fy=4, interpolation=cv2.INTER_NEAREST)
 
-
-        mask_area = (np.ones_like(densepose_mask) & (~weak_protect_area) & (~background_area)) | mask_dense_area
+        mask_area = (np.ones_like(densepose_mask) & (~body_protect_area) & (~background_area)) | mask_dense_area
 
         mask_area = hull_mask(mask_area * 255) // 255  # Convex Hull to expand the mask area
-        mask_area = mask_area & (~weak_protect_area)
+        mask_area = mask_area & (~body_protect_area)
         mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
         mask_area[mask_area < 25] = 0
         mask_area[mask_area >= 25] = 1
@@ -272,6 +267,5 @@ class AutoMasker:
             'schp_atr': preprocess_results['schp_atr']
         }
 
-
 if __name__ == '__main__':
-    pass
+    pass