Update SegCloth.py

SegCloth.py

@@ -1,26 +1,12 @@
 from transformers import pipeline
-from PIL import Image, ImageChops
+from PIL import Image, ImageChops, ImageOps
 import numpy as np
 from io import BytesIO
 import base64
-from transparent_background import Remover
 
 # Initialisation du pipeline de segmentation
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
-# Fonction pour supprimer l'arrière-plan
-def remove_background(image):
-    remover = Remover()
-    if isinstance(image, Image.Image):
-        output = remover.process(image)
-    elif isinstance(image, np.ndarray):
-        image_pil = Image.fromarray(image)
-        output = remover.process(image_pil)
-    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")
@@ -35,16 +21,17 @@ 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 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)
+            # 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]))
 
             # Application du masque sur l'image vide
             empty_image.paste(segmented_part, mask=mask_image)
@@ -62,11 +49,9 @@ 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))
                 
-                # 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)
+                imageBase64 = encode_image_to_base64(cropped_image)
+                #result_images.append((s['label'], imageBase64))
                 result_images.append(imageBase64)
 
     return result_images