Update SegCloth.py

SegCloth.py

@@ -4,7 +4,6 @@ import numpy as np
 from io import BytesIO
 import io
 import base64
-
 # Initialize segmentation pipeline
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
@@ -14,36 +13,30 @@ def encode_image_to_base64(image):
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-
-def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
     # Segment image
     segments = segmenter(img)
 
-    # List to hold the results
-    results = []
+    # Dictionary to store masks for each clothing type
+    masks = {clothing_type: None for clothing_type in clothes}
 
-    # Process each segment
+    # Create individual masks for each clothing type
     for s in segments:
         if s['label'] in clothes:
-            # Create a blank image with the same size as the original
-            clothing_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
-            
-            # Apply the mask to the new image
             mask = np.array(s['mask'])
-            mask_image = Image.fromarray(mask * 255)  # Convert mask to 255 range for alpha channel
-            
-            # Paste mask onto the blank image
-            clothing_image.paste(img, mask=mask_image)
-            
-            # Convert image to base64
-            image_base64 = encode_image_to_base64(clothing_image)
-            
-            # Add to results list
-            results.append({
-                "type": s['label'],
-                "image_base64": image_base64
-            })
-
-    return results
-
-
+            if masks[s['label']] is None:
+                masks[s['label']] = mask
+            else:
+                masks[s['label']] += mask
+    
+    # Create and save images for each clothing type
+    result_images = []
+    for clothing_type, mask in masks.items():
+        if mask is not None:
+            final_mask = Image.fromarray(np.uint8(mask * 255))  # Normalize mask to 0-255 range
+            result_image = img.copy()
+            result_image.putalpha(final_mask)
+            imageBase64 = encode_image_to_base64(result_image)
+            result_images.append((clothing_type, imageBase64))
+
+    return result_images