tonyassi commited on
Commit
8682ae7
1 Parent(s): bddf3b8

Upload 2 files

Browse files
Files changed (2) hide show
  1. SegCloth.py +33 -0
  2. requirements.txt +4 -0
SegCloth.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ from PIL import Image
3
+ import numpy as np
4
+
5
+
6
+ # Initialize segmentation pipeline
7
+ segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
8
+
9
+
10
+ def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
11
+ # Segment image
12
+ segments = segmenter(img)
13
+
14
+ # Create list of masks
15
+ mask_list = []
16
+ for s in segments:
17
+ if(s['label'] in clothes):
18
+ mask_list.append(s['mask'])
19
+
20
+
21
+ # Paste all masks on top of eachother
22
+ final_mask = np.array(mask_list[0])
23
+ for mask in mask_list:
24
+ current_mask = np.array(mask)
25
+ final_mask = final_mask + current_mask
26
+
27
+ # Convert final mask from np array to PIL image
28
+ final_mask = Image.fromarray(final_mask)
29
+
30
+ # Apply mask to original image
31
+ img.putalpha(final_mask)
32
+
33
+ return img
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ transformers
2
+ torch
3
+ pillow
4
+ numpy