Update SegCloth.py

SegCloth.py

@@ -10,29 +10,40 @@ def segment_clothing(img, clothes):
     # Segment image
     segments = segmenter(img)
 
+    # Define clothing items to expand
+    EXPAND_CLOTHING = {"Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"}
+
     # Create list of masks
     mask_list = []
+    expand_mask_list = []  # Separate list for clothes that need expansion
+
     for s in segments:
+        mask = np.array(s['mask'], dtype=np.uint8)  # Convert mask to numpy array
         if s['label'] in clothes:
-            mask_list.append(np.array(s['mask'], dtype=np.uint8))  # Convert to numpy array and ensure it's uint8
+            if s['label'] in EXPAND_CLOTHING:
+                expand_mask_list.append(mask)  # Store separately for expansion
+            else:
+                mask_list.append(mask)  # Keep others as they are
 
-    if not mask_list:
-        return img  # Return original image if no clothes found
+    if not mask_list and not expand_mask_list:
+        return img  # Return original image if no relevant items found
 
     # Initialize final mask with zeros
-    final_mask = np.zeros_like(mask_list[0], dtype=np.uint8)
+    final_mask = np.zeros_like(mask_list[0] if mask_list else expand_mask_list[0], dtype=np.uint8)
 
-    # Combine masks into one
+    # Combine normal masks into one
     for mask in mask_list:
         final_mask = np.maximum(final_mask, mask)
 
-    # Expand clothing boundaries using morphological dilation
-    height, width = final_mask.shape
-    kernel_size = max(1, int(0.05 * min(height, width)))  # Ensure at least 1 pixel
-    kernel = np.ones((kernel_size, kernel_size), np.uint8)
-    
-    # Dilate mask to expand clothing area
-    final_mask = cv2.dilate(final_mask, kernel, iterations=1)
+    # Expand selected clothing masks using morphological dilation
+    for mask in expand_mask_list:
+        height, width = mask.shape
+        kernel_size = max(1, int(0.05 * min(height, width)))  # 5% expansion
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
+
+        # Dilate mask and merge into final mask
+        dilated_mask = cv2.dilate(mask, kernel, iterations=1)
+        final_mask = np.maximum(final_mask, dilated_mask)
 
     # Optional: Use contour filling to ensure all areas within contours are filled
     contours, _ = cv2.findContours(final_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)