Create SegCloth.py

SegCloth.py

@@ -0,0 +1,63 @@
+from transformers import pipeline
+from PIL import Image
+import numpy as np
+import io
+import base64
+
+# Initialize segmentation pipeline
+segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
+
+def image_to_base64(img):
+    buffered = io.BytesIO()
+    img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
+
+def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+    # Segment image
+    segments = segmenter(img)
+
+    # List to hold the results
+    results = []
+
+    # Process each segment
+    for s in segments:
+        if s['label'] in clothes:
+            # Convert mask to array
+            mask = np.array(s['mask'])
+            
+            # Find bounding box of the mask
+            y_nonzero, x_nonzero = np.nonzero(mask)
+            if len(x_nonzero) == 0 or len(y_nonzero) == 0:
+                continue
+            min_x, max_x = np.min(x_nonzero), np.max(x_nonzero)
+            min_y, max_y = np.min(y_nonzero), np.max(y_nonzero)
+
+            # Crop mask and original image to the bounding box
+            cropped_mask = mask[min_y:max_y+1, min_x:max_x+1]
+            cropped_image = img.crop((min_x, min_y, max_x+1, max_y+1))
+            
+            # Create an image with the same size as the bounding box
+            clothing_image = Image.new("RGBA", (max_x - min_x + 1, max_y - min_y + 1), (0, 0, 0, 0))
+            
+            # Apply the mask to the new image
+            mask_image = Image.fromarray(cropped_mask * 255)  # Convert mask to 255 range for alpha channel
+            
+            # Paste cropped image onto the blank image using the mask
+            clothing_image.paste(cropped_image, mask=mask_image)
+            
+            # Convert image to base64
+            image_base64 = image_to_base64(clothing_image)
+            
+            # Add to results list
+            results.append({
+                "type": s['label'],
+                "image_base64": image_base64
+            })
+
+    return results
+
+# Example usage
+# img = Image.open("path_to_your_image.jpg")
+# results = segment_clothing(img)
+# print(results)