Update SegCloth.py

SegCloth.py

@@ -1,5 +1,5 @@
 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
@@ -12,7 +12,7 @@ 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"], margin=10):
     # Segmentation de l'image
     segments = segmenter(img)
     
@@ -36,8 +36,21 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
             # 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))
+            # Déterminer la bounding box du masque
+            bbox = mask_image.getbbox()
+            if bbox:
+                # Ajouter la marge autour de la bounding box
+                left, top, right, bottom = bbox
+                left = max(0, left - margin)
+                top = max(0, top - margin)
+                right = min(img.width, right + margin)
+                bottom = min(img.height, bottom + margin)
+                
+                # 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
+                imageBase64 = encode_image_to_base64(cropped_image)
+                result_images.append((s['label'], imageBase64))
 
-    return result_images
+    return result_images