Update SegCloth.py

SegCloth.py

@@ -4,39 +4,33 @@ import numpy as np
 from io import BytesIO
 import io
 import base64
+
 # Initialize segmentation pipeline
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
 
-def encode_image_to_base64(image):
-    buffered = BytesIO()
-    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)
 
-    # Dictionary to store masks for each clothing type
-    masks = {clothing_type: None for clothing_type in clothes}
-
-    # Create individual masks for each clothing type
+    # Create list of masks
+    mask_list = []
     for s in segments:
-        if s['label'] in clothes:
-            mask = np.array(s['mask'])
-            if masks[s['label']] is None:
-                masks[s['label']] = mask
-            else:
-                masks[s['label']] += mask
-    
-    # Create and save images for each clothing type
+        if(s['label'] in clothes):
+            mask_list.append(s['mask'])
+
     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))
+
+    # Paste all masks on top of eachother 
+    final_mask = np.array(mask_list[0])
+    for mask in mask_list:
+        current_mask = np.array(mask)
+        final_mask_bis = Image.fromarray(current_mask)
+        img.putalpha(final_mask_bis)
+        imageBase64 = encode_image_to_base64(img)
+        result_images.append(('clothing_type', imageBase64))    
+ 
+
+  
 
     return result_images