Update SegCloth.py

SegCloth.py

@@ -13,26 +13,33 @@ 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", "Scarf"]):
     # Segment image
     segments = segmenter(img)
 
-    # Create list of masks
-    mask_list = []
-    for s in segments:
-        if(s['label'] in clothes):
-            mask_list.append(s['mask'])
+    # Convert image to RGBA
+    img = img.convert("RGBA")
 
+    # Create list of masks
     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))    
+    for s in segments:
+        if s['label'] in clothes:
+            # Extract mask and resize image to mask size
+            current_mask = np.array(s['mask'])
+            mask_size = current_mask.shape[::-1]  # Mask size is (width, height)
+            
+            # Resize the original image to match the mask size
+            resized_img = img.resize(mask_size)
+
+            # Apply mask to resized image
+            final_mask = Image.fromarray(current_mask)
+            resized_img.putalpha(final_mask)
+
+            # Convert the final image to base64
+            imageBase64 = encode_image_to_base64(resized_img)
+            result_images.append((s['label'], imageBase64))
+
+    return result_images