Update SegCloth.py

SegCloth.py

@@ -1,3 +1,34 @@
-import os
+import spaces
+from transformers import pipeline
+from PIL import Image
+import numpy as np
 
-exec(os.environ.get('SegCloth'))
+
+# Initialize segmentation pipeline
+segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes", device="cuda")
+
+@spaces.GPU(enable_queue=True)
+def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "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'])
+
+
+    # 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 = final_mask + current_mask
+            
+    # Convert final mask from np array to PIL image
+    final_mask = Image.fromarray(final_mask)
+
+    # Apply mask to original image
+    img.putalpha(final_mask)
+
+    return img