Update SegCloth.py

SegCloth.py

@@ -6,12 +6,13 @@ import cv2  # OpenCV for better mask processing
 # Initialize segmentation pipeline
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
+
 def segment_clothing(img, clothes):
     # Segment image
     segments = segmenter(img)
 
     # Define clothing items to expand
-    EXPAND_CLOTHING = {"Upper-clothes", "Skirt", "Pants", "Dress"}
+    EXPAND_CLOTHING = {"Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe"}
 
     # Create list of masks
     mask_list = []
@@ -35,14 +36,19 @@ def segment_clothing(img, clothes):
     for mask in mask_list:
         final_mask = np.maximum(final_mask, mask)
 
-    # Expand selected clothing masks using morphological dilation
+    # Expand selected clothing masks using closing + dilation
     for mask in expand_mask_list:
         height, width = mask.shape
         kernel_size = max(1, int(0.02 * 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)
+        # **Step 1: Fill gaps using Closing (Dilation + Erosion)**
+        closed_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)
+
+        # **Step 2: Expand using Dilation**
+        dilated_mask = cv2.dilate(closed_mask, kernel, iterations=1)
+
+        # Merge into final mask
         final_mask = np.maximum(final_mask, dilated_mask)
 
     # Optional: Use contour filling to ensure all areas within contours are filled