-
Notifications
You must be signed in to change notification settings - Fork 2
/
SegCloth.py
33 lines (23 loc) · 891 Bytes
/
SegCloth.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from transformers import pipeline
from PIL import Image
import numpy as np
# Initialize segmentation pipeline
segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
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