Update SegCloth.py

SegCloth.py

@@ -1,5 +1,5 @@
 from transformers import pipeline
-from PIL import Image, ImageChops, ImageOps
+from PIL import Image, ImageChops
 import numpy as np
 from io import BytesIO
 import base64
@@ -8,7 +8,7 @@ from transparent_background import Remover
 # Initialisation du pipeline de segmentation
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
-#@spaces.GPU
+# Fonction pour supprimer l'arrière-plan
 def remove_background(image):
     remover = Remover()
     if isinstance(image, Image.Image):
@@ -19,7 +19,8 @@ def remove_background(image):
     else:
         raise TypeError("Unsupported image type")
     return output
-    
+
+# Fonction pour encoder une image en base64
 def encode_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format="PNG")
@@ -34,17 +35,16 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
 
     for s in segments:
         if s['label'] in clothes:
-            # 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]))
+            # Conversion du masque en tableau NumPy et ajustement de l'image
+            mask_array = np.array(s['mask'])
+            mask_image = Image.fromarray(mask_array).convert("L")  # Convertir le masque en niveau de gris
+
+            # Appliquer le masque à l'image d'origine (en RGBA)
+            mask_rgba = Image.merge("RGBA", [mask_image, mask_image, mask_image, mask_image])
+            segmented_part = ImageChops.multiply(img.convert("RGBA"), mask_rgba)
 
             # Application du masque sur l'image vide
             empty_image.paste(segmented_part, mask=mask_image)
@@ -62,10 +62,11 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
                 # Recadrer l'image à la taille du masque avec la marge
                 cropped_image = empty_image.crop((left, top, right, bottom))
                 
-                # Encodage de l'image recadrée en base64
+                # Supprimer l'arrière-plan
                 image_rm_background = remove_background(cropped_image)
+                
+                # Encodage de l'image recadrée en base64
                 imageBase64 = encode_image_to_base64(image_rm_background)
-                #result_images.append((s['label'], imageBase64))
                 result_images.append(imageBase64)
 
     return result_images