Update SegCloth.py

SegCloth.py

@@ -2,7 +2,6 @@ from transformers import pipeline
 from PIL import Image
 import numpy as np
 from io import BytesIO
-import io
 import base64
 
 # Initialize segmentation pipeline
@@ -13,28 +12,24 @@ 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)
 
-    # Create list of masks
+    # Create list of masks and their corresponding clothing types
     mask_list = []
     for s in segments:
-        if(s['label'] in clothes):
-            mask_list.append(s['mask'])
+        if s['label'] in clothes:
+            mask_list.append((s['mask'], s['label']))
 
     result_images = []
 
-    # Paste all masks on top of eachother 
-    final_mask = np.array(mask_list[0])
-    for mask in mask_list:
+    # Paste all masks on top of each other
+    for mask, clothing_type 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))    
- 
+        result_images.append((clothing_type, imageBase64))
 
-  
-
-    return result_images
+    return result_images