-
Notifications
You must be signed in to change notification settings - Fork 8
/
face_segment_yuvalcontour.py
executable file
·141 lines (117 loc) · 5.32 KB
/
face_segment_yuvalcontour.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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
#!/usr/bin/env python
import numpy as np
from PIL import Image
from PIL import ExifTags
import matplotlib.pyplot as plt
from scipy import ndimage
from scipy.spatial import ConvexHull
from skimage import draw
import os
import caffe
import argparse
from face_alignment.api import FaceAlignment, LandmarksType, NetworkSize
from util import *
def main(args):
image_paths = read_list(args.image_list)
# init
caffe.set_device(0)
caffe.set_mode_gpu()
# load net
net = caffe.Net(args.prototxt, args.caffemodel, caffe.TEST)
for path in image_paths:
# load image, switch to BGR, subtract mean, and make dims C x H x W for Caffe
print path
if path[-3:] == 'jpg' or path[-3:] == 'png':
imi = open_image(path)
# resize for memory
width, height = imi.size
if height > 800:
imi = imi.resize((int(800*width/height), 800))
width, height = imi.size
else:
continue
dir = 'images/' + path[path.rindex('/', 0, path.rindex('/'))+1:path.rindex('/')] + '/'
if not os.path.exists(dir):
os.makedirs(dir)
txtfile = dir+path[path.rindex('/')+1:-3]+'txt'
# use 2D-FAN detect landmarks
fa = FaceAlignment(LandmarksType._2D, enable_cuda=True,
flip_input=False, use_cnn_face_detector=True)
try:
landmarks = fa.get_landmarks(np.array(imi))[-1]
save_landmarks(txtfile, landmarks)
landmarks = landmarks.astype('uint16')
except:
# no face deteced
continue
if args.crop == 'middle':
imi, landmarks = crop_image_middle(landmarks, imi)
elif args.crop == 'min':
imi, landmarks = crop_image_min(landmarks, imi)
width, height = imi.size
if '300' in args.prototxt:
imi = imi.resize((300, 300))
landmarks[:,0] = landmarks[:,0] * 300 / width
landmarks[:,1] = landmarks[:,1] * 300 / height
width = height = 300
else:
imi = imi.resize((500, 500))
landmarks[:,0] = landmarks[:,0] * 500 / width
landmarks[:,1] = landmarks[:,1] * 500 / height
width = height = 500
im = np.array(imi, dtype=np.float32)
im = im[:,:,::-1]
im -= np.array((104.00698793,116.66876762,122.67891434))
im = im.transpose((2,0,1))
# shape for input (data blob is N x C x H x W), set data
net.blobs['data'].reshape(1, *im.shape)
net.blobs['data'].data[...] = im
# run net and take argmax for prediction
net.forward()
mask = net.blobs['score'].data[0].argmax(axis=0)
save = True if args.save == 'True' else False
path = path[:-1] if path[-1] == '/' else path
if '300' in args.prototxt:
image_name = path[path.rindex('/')+1:-4] + '_yuvalcontour_300_nocrf_' + args.crop + '.png'
else:
image_name = path[path.rindex('/')+1:-4] + '_yuvalcontour_nocrf_' + args.crop + '.png'
# draw landmarks
lm = np.array(imi)
for i in range(landmarks.shape[0]):
rr, cc = draw.circle(landmarks[i,1].astype('int32'), landmarks[i,0].astype('int32'), 1)
lm[rr, cc, :] = np.array((255, 0, 0))
# create mask contour
hull = ConvexHull(landmarks)
mask_contour = np.zeros((height, width))
rr, cc = draw.polygon(landmarks[hull.vertices,1], landmarks[hull.vertices,0], mask_contour.shape)
mask_contour[rr,cc] = 1
mask = np.clip(mask + mask_contour, 0, 1)
im_seg = imi * np.tile((mask!=0)[:,:,np.newaxis], (1,1,3))
show_result(lm, mask, im_seg, save=save, filename=dir+image_name)
# generate prob
#prob = np.concatenate(((1-mask)[np.newaxis,:,:]*0.9 + mask[np.newaxis,:,:]*0.1, mask[np.newaxis,:,:]*0.9+(1-mask)[np.newaxis,:,:]*0.1), axis=0)
prob = ndimage.gaussian_filter(mask*1.0, sigma=5)
prob = np.concatenate(((1-prob)[np.newaxis,:,:], prob[np.newaxis,:,:]), axis=0)
# add CRF
map = CRF(prob, np.array(imi))
if '300' in args.prototxt:
image_name = path[path.rindex('/')+1:-4] + '_yuvalcontour_300_crf_' + args.crop + '.png'
else:
image_name = path[path.rindex('/')+1:-4] + '_yuvalcontour_crf_' + args.crop + '.png'
show_result(imi, map, np.tile((map!=0)[:,:,np.newaxis], (1,1,3)) * imi, save=save, filename=dir+image_name)
if __name__=="__main__":
parser = argparse.ArgumentParser(description='Face segmentation of yuval.')
parser.add_argument('--image_list', default='input/list.txt',
type=str, help='path to image')
parser.add_argument('--prototxt',
default='model/face_seg_fcn8s_300_deploy.prototxt',
type=str, help='path to prototxt')
parser.add_argument('--caffemodel',
default='model/face_seg_fcn8s_300.caffemodel',
type=str, help='path to caffemodel')
parser.add_argument('--crop', choices=['min', 'middle', 'no'],
default='min', help='choose min/middle/no crop')
parser.add_argument('--save', choices=['True', 'False'],
default='False', help='choose if save final result')
args = parser.parse_args()
main(args)