Update SegCloth.py

SegCloth.py

@@ -1,12 +1,11 @@
 from transformers import pipeline
-from PIL import Image
+from PIL import Image, ImageFilter
 import numpy as np
-
+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"]):
     # Segment image
     segments = segmenter(img)
@@ -15,19 +14,30 @@ def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "D
     mask_list = []
     for s in segments:
         if(s['label'] in clothes):
-            mask_list.append(s['mask'])
+            mask_list.append(np.array(s['mask'], dtype=np.uint8))  # Convert to numpy array and ensure it's uint8
 
+    # Initialize final mask with zeros
+    final_mask = np.zeros_like(mask_list[0], dtype=np.uint8)
 
-    # Paste all masks on top of eachother 
-    final_mask = np.array(mask_list[0])
+    # Combine masks into one
     for mask in mask_list:
-        current_mask = np.array(mask)
-        final_mask = final_mask + current_mask
-            
+        final_mask = np.maximum(final_mask, mask)
+
+    # Optional: Smooth the mask to reduce rough edges (using Gaussian blur)
+    final_mask = cv2.GaussianBlur(final_mask, (7, 7), 0)
+
+    # Optional: Dilate the mask to ensure coverage at edges
+    kernel = np.ones((5,5), np.uint8)
+    final_mask = cv2.dilate(final_mask, kernel, iterations=1)
+
+    # 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
     final_mask = Image.fromarray(final_mask)
 
-    # Apply mask to original image
+    # Apply mask to original image (convert to RGBA first)
+    img = img.convert("RGBA")
     img.putalpha(final_mask)
 
-    return img
+    return img