Update SegCloth.py

SegCloth.py

@@ -1,10 +1,10 @@
 from transformers import pipeline
-from PIL import Image
+from PIL import Image, ImageChops
 import numpy as np
 from io import BytesIO
 import base64
 
-# Initialize segmentation pipeline
+# Initialisation du pipeline de segmentation
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
 def encode_image_to_base64(image):
@@ -13,23 +13,31 @@ def encode_image_to_base64(image):
     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"]):
-    # Segment image
+    # Segmentation de l'image
     segments = segmenter(img)
+    
+    # Liste des images segmentées
+    result_images = []
 
-    # 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'], s['label']))
-
-    result_images = []
+            # Conversion du masque en tableau NumPy
+            mask_array = np.array(s['mask'])
+            
+            # Création d'une image vide avec transparence
+            empty_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
+            
+            # Conversion du masque en image PIL (niveau de gris)
+            mask_image = Image.fromarray(mask_array).convert("L")
+            
+            # Extraction de la partie de l'image correspondant au masque
+            segmented_part = ImageChops.multiply(img.convert("RGBA"), Image.merge("RGBA", [mask_image, mask_image, mask_image, mask_image]))
 
-    # 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))
+            # Application du masque sur l'image vide
+            empty_image.paste(segmented_part, mask=mask_image)
+            
+            # Encodage de l'image résultante en base64
+            imageBase64 = encode_image_to_base64(empty_image)
+            result_images.append((s['label'], imageBase64))
 
-    return result_images
+    return result_images