Update SegCloth.py

SegCloth.py

@@ -1,42 +1,40 @@
+from transformers import pipeline
 from PIL import Image
 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", "Scarf"]):
+def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
     # Segment image
     segments = segmenter(img)
 
-    # Convert image to RGBA
-    img = img.convert("RGBA")
-
     # Create list of masks
-    result_images = []
+    mask_list = []
     for s in segments:
-        if s['label'] in clothes:
-            final_mask = Image.fromarray(current_mask)
-            resized_img.putalpha(final_mask)
+        if(s['label'] in clothes):
+            mask_list.append(s['mask'])
 
+    result_images = []
 
-
-
-            # Convert the final image to base64
-            imageBase64 = encode_image_to_base64(resized_img)
-            result_images.append((s['label'], imageBase64))
-
-    return result_images  
+    # 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