Update SegCloth.py

SegCloth.py

@@ -6,7 +6,7 @@ import cv2  # OpenCV for better mask processing
 # Initialize segmentation pipeline
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
-def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+def segment_clothing(img, clothes):
     # Segment image
     segments = segmenter(img)
 
@@ -16,6 +16,9 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
         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 not mask_list:
+        return img  # Return original image if no clothes found
+
     # Initialize final mask with zeros
     final_mask = np.zeros_like(mask_list[0], dtype=np.uint8)
 
@@ -23,24 +26,22 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     for mask in mask_list:
         final_mask = np.maximum(final_mask, mask)
 
-    # Optional: Dilate the mask to ensure coverage at edges
-    #kernel = np.ones((5, 5), np.uint8)
-    #final_mask = cv2.dilate(final_mask, kernel, iterations=2)
-
-    # Optional: Erode to slightly smoothen the mask
-    #final_mask = cv2.erode(final_mask, kernel, iterations=1)
+    # 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)
 
     # Optional: Use contour filling to ensure all areas within contours are filled
     contours, _ = cv2.findContours(final_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     cv2.drawContours(final_mask, contours, -1, (255), thickness=cv2.FILLED)
 
-    # Apply Gaussian blur to smooth edges and reduce noise
-    #final_mask = cv2.GaussianBlur(final_mask, (7, 7), 0)
-
     # Convert mask to binary (0 or 255) if needed for alpha channel
     _, final_mask = cv2.threshold(final_mask, 127, 255, cv2.THRESH_BINARY)
 
-    # Convert final mask from np array to PIL image
+    # Convert final mask from numpy array to PIL image
     final_mask = Image.fromarray(final_mask)
 
     # Apply mask to original image (convert to RGBA first)
@@ -48,3 +49,4 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     img.putalpha(final_mask)
 
     return img
+