def __init__(self,
model_path,
gpu_id=0,
minsize=40,
factor=0.5,
threshold=[0.8, 0.8, 0.9]):
rnet_model_path = os.path.join(model_path, "rnet/rnet-3000000")
pnet_model_path = os.path.join(model_path, "pnet/pnet-3000000")
onet_model_path = os.path.join(model_path, "onet/onet-500000")
if not os.path.exists(model_path) or \
not os.path.exists(os.path.join(model_path,"rnet")) or \
not os.path.exists(os.path.join(model_path,"pnet")) or \
not os.path.exists(os.path.join(model_path,"onet")):
raise Exception("Error when loading {}".format(model_path))
# default detection parameters
self.minsize = minsize
self.factor = factor
self.threshold = threshold
# load models
with tf.device('/gpu:{}'.format(gpu_id)):
with tf.Graph().as_default() as p:
config = tf.ConfigProto(allow_soft_placement=True)
self.sess = tf.Session(config=config)
self.pnet_input = tf.placeholder(tf.float32,
[None, None, None, 3])
self.pnet = PNet({'data': self.pnet_input}, mode='test')
self.pnet_output = self.pnet.get_all_output()
self.rnet_input = tf.placeholder(tf.float32, [None, 24, 24, 3])
self.rnet = RNet({'data': self.rnet_input}, mode='test')
self.rnet_output = self.rnet.get_all_output()
self.onet_input = tf.placeholder(tf.float32, [None, 48, 48, 3])
self.onet = ONet({'data': self.onet_input}, mode='test')
self.onet_output = self.onet.get_all_output()
saver_pnet = tf.train.Saver([
v for v in tf.global_variables() if v.name[0:5] == "pnet/"
])
saver_rnet = tf.train.Saver([
v for v in tf.global_variables() if v.name[0:5] == "rnet/"
])
saver_onet = tf.train.Saver([
v for v in tf.global_variables() if v.name[0:5] == "onet/"
])
saver_pnet.restore(self.sess, pnet_model_path)
self.pnet_func = lambda img: self.sess.run(
self.pnet_output, feed_dict={self.pnet_input: img})
saver_rnet.restore(self.sess, rnet_model_path)
self.rnet_func = lambda img: self.sess.run(
self.rnet_output, feed_dict={self.rnet_input: img})
saver_onet.restore(self.sess, onet_model_path)
self.onet_func = lambda img: self.sess.run(
self.onet_output, feed_dict={self.onet_input: img})
def __init__(
self, image_size=160, margin=0, min_face_size=20,
thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
select_largest=True, keep_all=False, device=None
):
super().__init__()
self.image_size = image_size
self.margin = margin
self.min_face_size = min_face_size
self.thresholds = thresholds
self.factor = factor
self.post_process = post_process
self.select_largest = select_largest
self.keep_all = keep_all
self.pnet = PNet()
self.rnet = RNet()
self.onet = ONet()
self.device = torch.device('cpu')
if device is not None:
self.device = device
self.to(device)
# Copyright 2019 smarsu. All Rights Reserved.
# --------------------------------------------------------
import os.path as osp
import cv2
import numpy as np
from mtcnn import PNet, RNet, ONet
from square import square_boxes
from broad import broad_boxes
from crop import crop
pnet = PNet(scale_factor=0.89,
conf_thrs=0.8,
nms_thrs=0.5,
min_face=60,
nms_topk=32)
pnet.sess.restore(
osp.join(pnet.model_root, '3.2153504_cycle_7_0.01_pnet_v2.npz'))
rnet = RNet(conf_thrs=0.5)
rnet.sess.restore(osp.join(rnet.model_root, '0.022445953_53_0.001_rnet.npz'))
onet = ONet(conf_thrs=0.5)
onet.sess.restore(osp.join(onet.model_root, '0.012311436_69_0.01_onet.npz'))
def detect(img, top_k=-1):
"""Do face detection with the input img.
def main(args):
image_size = 24
save_dir = str(image_size)
anno_file = 'wider_face_train.txt'
im_dir = 'WIDER_train/images/'
neg_save_dir = save_dir + '/negative'
pos_save_dir = save_dir + '/positive'
part_save_dir = save_dir + '/part'
if not os.path.exists(save_dir):
os.mkdir(save_dir)
if not os.path.exists(pos_save_dir):
os.mkdir(pos_save_dir)
if not os.path.exists(part_save_dir):
os.mkdir(part_save_dir)
if not os.path.exists(neg_save_dir):
os.mkdir(neg_save_dir)
f1 = open(save_dir + '/pos_24.txt', 'w')
f2 = open(save_dir + '/neg_24.txt', 'w')
f3 = open(save_dir + '/part_24.txt', 'w')
threshold = 0.6
with open(anno_file, 'r') as f:
annotations = f.readlines()
num = len(annotations)
print('%d pics in total' % num)
p_idx = 0 # positive
n_idx = 0 # negative
d_idx = 0 # dont care
image_idx = 0
with tf.device('/gpu:0'):
minsize = 20
factor = 0.709
model_file = args.pnet_model
with tf.Graph().as_default():
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.per_process_gpu_memory_fraction = 0.5
with tf.Session(config=config) as sess:
image = tf.placeholder(tf.float32, [None, None, None, 3])
pnet = PNet({'data': image}, mode='test')
out_tensor = pnet.get_all_output()
init_op = tf.global_variables_initializer()
sess.run(init_op)
saver = tf.train.Saver()
saver.restore(sess, model_file)
def pnet_fun(img):
return sess.run(out_tensor, feed_dict={image: img})
for annotation in annotations:
annotation = annotation.strip().split(' ')
bbox = list(map(float, annotation[1:]))
gts = np.array(bbox, dtype=np.float32).reshape(-1, 4)
img_path = im_dir + annotation[0] + '.jpg'
img = cv2.imread(img_path)
rectangles = detect_face_12net(img, minsize, pnet_fun,
threshold, factor)
image_idx += 1
view_bar(image_idx, num)
for box in rectangles:
lis = box.astype(np.int32)
mask = lis < 0
lis[mask] = 0
x_left, y_top, x_right, y_bottom, _ = lis
crop_w = x_right - x_left + 1
crop_h = y_bottom - y_top + 1
# ignore box that is too small or beyond image border
if crop_w < image_size or crop_h < image_size:
continue
Iou = IoU(box, gts)
cropped_im = img[y_top:y_bottom + 1,
x_left:x_right + 1]
resized_im = cv2.resize(cropped_im,
(image_size, image_size),
interpolation=cv2.INTER_LINEAR)
# save negative images and write label
if np.max(Iou) < 0.3:
# Iou with all gts must below 0.3
save_file = os.path.join(neg_save_dir,
'%s.jpg' % n_idx)
f2.write('%s/negative/%s' % (save_dir, n_idx) +
' 0\n')
cv2.imwrite(save_file, resized_im)
n_idx += 1
else:
# find gt_box with the highest iou
idx = np.argmax(Iou)
assigned_gt = gts[idx]
x1, y1, x2, y2 = assigned_gt
# compute bbox reg label
offset_x1 = (x1 - x_left) / float(crop_w)
offset_y1 = (y1 - y_top) / float(crop_h)
offset_x2 = (x2 - x_right) / float(crop_w)
offset_y2 = (y2 - y_bottom) / float(crop_h)
if np.max(Iou) >= 0.65:
save_file = os.path.join(
pos_save_dir, '%s.jpg' % p_idx)
f1.write('%s/positive/%s' % (save_dir, p_idx) +
' 1 %.2f %.2f %.2f %.2f\n' %
(offset_x1, offset_y1, offset_x2,
offset_y2))
cv2.imwrite(save_file, resized_im)
p_idx += 1
elif np.max(Iou) >= 0.4:
save_file = os.path.join(
part_save_dir, '%s.jpg' % d_idx)
f3.write('%s/part/%s' % (save_dir, d_idx) +
' -1 %.2f %.2f %.2f %.2f\n' %
(offset_x1, offset_y1, offset_x2,
offset_y2))
cv2.imwrite(save_file, resized_im)
d_idx += 1
f1.close()
f2.close()
f3.close()
layer = model.get_layer(layer_name)
if "conv" in layer_name:
layer.set_weights([
weights_dict[layer_name]["weights"],
weights_dict[layer_name]["biases"]
])
else:
prelu_weight = weights_dict[layer_name]['alpha']
try:
layer.set_weights([prelu_weight])
except:
layer.set_weights([prelu_weight[np.newaxis, np.newaxis, :]])
return True
pnet, rnet, onet = PNet(), RNet(), ONet()
pnet(tf.ones(shape=[1, 12, 12, 3]))
rnet(tf.ones(shape=[1, 24, 24, 3]))
onet(tf.ones(shape=[1, 48, 48, 3]))
load_weights(pnet, "./det1.npy"), load_weights(rnet,
"./det2.npy"), load_weights(
onet, "./det3.npy")
pnet.predict(tf.ones(shape=[1, 12, 12, 3]))
pnet_converter = tf.lite.TFLiteConverter.from_keras_model(pnet)
pnet_converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
with open("pnet.tflite", "wb") as f:
pnet_tflite_model = pnet_converter.convert()
f.write(pnet_tflite_model)
rnet.predict(tf.ones(shape=[1, 24, 24, 3]))
layer = model.get_layer(layer_name)
if "conv" in layer_name:
layer.set_weights([
weights_dict[layer_name]["weights"],
weights_dict[layer_name]["biases"]
])
else:
prelu_weight = weights_dict[layer_name]['alpha']
try:
layer.set_weights([prelu_weight])
except:
layer.set_weights([prelu_weight[np.newaxis, np.newaxis, :]])
return True
pnet, rnet, onet = PNet(), RNet(), ONet()
pnet(tf.ones(shape=[1, 12, 12, 3]))
rnet(tf.ones(shape=[1, 24, 24, 3]))
onet(tf.ones(shape=[1, 48, 48, 3]))
load_weights(pnet, "./det1.npy"), load_weights(rnet,
"./det2.npy"), load_weights(
onet, "./det3.npy")
image = cv2.cvtColor(cv2.imread("./multiface.jpg"), cv2.COLOR_BGR2RGB)
total_boxes, points = detect_face(image, 20, pnet, rnet, onet, [0.6, 0.7, 0.7],
0.709)
for bounding_box, keypoints in zip(total_boxes, points.T):
bounding_boxes = {
'box': [
int(bounding_box[0]),
dataset_root = '/datasets/wider'
if __name__ == '__main__':
parse = sys.argv[1] # 'train'
net = sys.argv[2]
#net = 'pnet'
if parse == 'train':
widerface = WiderFace(
'/datasets/wider/images',
'/datasets/wider/wider_face_split/wider_face_train_bbx_gt.txt')
if net == 'pnet':
pnet = PNet(batch_size=32,
no_mask=False,
rd_size=False,
min_face=60,
scale_factor=0.89)
# 0.8780253_1_0.1_pnet.npz the last size of step 1 for lr 0.1.
pnet.sess.restore(
osp.join(pnet.model_root, '3.2153504_cycle_7_0.01_pnet.npz'))
pnet.train(widerface.train_datas, 100, lr=0.001)
# pnet.train(widerface.train_datas_debug, 100, lr=0.1)
#conf, box = pnet.test(widerface.train_datas_debug(1)[0][0][0])
#print(conf)
#print(box)
#conf, box = pnet._check_model_capable(widerface.train_datas_debug)
#print(conf)
#print(box)
elif net == 'rnet':
widerfacepnet = WiderFace(
class MTCNN_detector:
def __init__(self,
model_path,
gpu_id=0,
minsize=40,
factor=0.5,
threshold=[0.8, 0.8, 0.9]):
rnet_model_path = os.path.join(model_path, "rnet/rnet-3000000")
pnet_model_path = os.path.join(model_path, "pnet/pnet-3000000")
onet_model_path = os.path.join(model_path, "onet/onet-500000")
if not os.path.exists(model_path) or \
not os.path.exists(os.path.join(model_path,"rnet")) or \
not os.path.exists(os.path.join(model_path,"pnet")) or \
not os.path.exists(os.path.join(model_path,"onet")):
raise Exception("Error when loading {}".format(model_path))
# default detection parameters
self.minsize = minsize
self.factor = factor
self.threshold = threshold
# load models
with tf.device('/gpu:{}'.format(gpu_id)):
with tf.Graph().as_default() as p:
config = tf.ConfigProto(allow_soft_placement=True)
self.sess = tf.Session(config=config)
self.pnet_input = tf.placeholder(tf.float32,
[None, None, None, 3])
self.pnet = PNet({'data': self.pnet_input}, mode='test')
self.pnet_output = self.pnet.get_all_output()
self.rnet_input = tf.placeholder(tf.float32, [None, 24, 24, 3])
self.rnet = RNet({'data': self.rnet_input}, mode='test')
self.rnet_output = self.rnet.get_all_output()
self.onet_input = tf.placeholder(tf.float32, [None, 48, 48, 3])
self.onet = ONet({'data': self.onet_input}, mode='test')
self.onet_output = self.onet.get_all_output()
saver_pnet = tf.train.Saver([
v for v in tf.global_variables() if v.name[0:5] == "pnet/"
])
saver_rnet = tf.train.Saver([
v for v in tf.global_variables() if v.name[0:5] == "rnet/"
])
saver_onet = tf.train.Saver([
v for v in tf.global_variables() if v.name[0:5] == "onet/"
])
saver_pnet.restore(self.sess, pnet_model_path)
self.pnet_func = lambda img: self.sess.run(
self.pnet_output, feed_dict={self.pnet_input: img})
saver_rnet.restore(self.sess, rnet_model_path)
self.rnet_func = lambda img: self.sess.run(
self.rnet_output, feed_dict={self.rnet_input: img})
saver_onet.restore(self.sess, onet_model_path)
self.onet_func = lambda img: self.sess.run(
self.onet_output, feed_dict={self.onet_input: img})
def destroy(self):
self.sess.close()
# Returns:
# rects: a numpy array of shape [num_face, 5].
# Denote of each row:
# [left_top_x, left_top_y, right_bottom_x, right_bottom_y, confidence]
def calc_det_result(self, image):
rects, shapes = self.calc_landmark_result(image)
return rects
# Returns:
# rectangles: a numpy array of shape [num_face, 5].
# Denote of each row:
# [left_top_x, left_top_y, right_bottom_x, right_bottom_y, confidence]
# points: a numpy array of shape [num_face, 10],
# Denote of each row:
# [left_eye_x, left_eye_y, right_eye_x, right_eye_y,
# nose_x, nose_y,
# left_mouthcorner_x,left_mouthcorner_y,right_mouthcorner_x,right_mouthcorner_y,]
def calc_landmark_result(self, image):
# TODO time test
start = cv2.getTickCount()
rectangles, shapes = tools.detect_face(image, self.minsize,
self.pnet_func, self.rnet_func,
self.onet_func, self.threshold,
self.factor)
shapes = np.transpose(shapes)
# TODO time test
usetime = (cv2.getTickCount() - start) / cv2.getTickFrequency()
print "Use time {}s.".format(usetime)
return rectangles, shapes
# SAME as calc_landmark_result but return shape [num_face, 4]
# Denote of each row:
# [left_eye_x, left_eye_y, right_eye_x, right_eye_y]
def extract_eye_result(self, shapes):
assert (shapes is not None)
assert (shapes.shape[0] > 0 and shapes.shape[1] == 10)
return shapes[:, 0:4]
# show the detection results
def show_result(self, image_path, rectangles, shapes):
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
if rectangles.shape[0] != shapes.shape[0]:
print "Error in show results {} != {}.".format(
rectangles.shape[0], shapes.shape[0])
for rect in rectangles:
cv2.rectangle(image, (int(round(rect[0])),int(round(rect[1]))), \
(int(round(rect[2])),int(round(rect[3]))), (255,255,0), 2)
for shape in shapes:
shape_num = len(shape) / 2
for i in xrange(shape_num):
pt = (int(round(shape[2 * i])), int(round(shape[2 * i + 1])))
cv2.circle(image, pt, 2, (0, 0, 255), 2)
cv2.imwrite("show.jpg", image)
def bounding_box_face(img, gpu_mem):
### Model parameters
model_dir = 'assets/save_model/'
minsize = 20
# factor = 0.7
# threshold = [0.8, 0.8, 0.8]
factor = 0.709
threshold = [0.6, 0.7, 0.7]
file_paths = get_model_filenames(model_dir)
with tf.device('/gpu:0'):
with tf.Graph().as_default():
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.per_process_gpu_memory_fraction = gpu_mem
with tf.Session(config=config) as sess:
if len(file_paths) == 3:
image_pnet = tf.placeholder(tf.float32,
[None, None, None, 3])
pnet = PNet({'data': image_pnet}, mode='test')
out_tensor_pnet = pnet.get_all_output()
image_rnet = tf.placeholder(tf.float32, [None, 24, 24, 3])
rnet = RNet({'data': image_rnet}, mode='test')
out_tensor_rnet = rnet.get_all_output()
image_onet = tf.placeholder(tf.float32, [None, 48, 48, 3])
onet = ONet({'data': image_onet}, mode='test')
out_tensor_onet = onet.get_all_output()
saver_pnet = tf.train.Saver([
v for v in tf.global_variables()
if v.name[0:5] == "pnet/"
])
saver_rnet = tf.train.Saver([
v for v in tf.global_variables()
if v.name[0:5] == "rnet/"
])
saver_onet = tf.train.Saver([
v for v in tf.global_variables()
if v.name[0:5] == "onet/"
])
saver_pnet.restore(sess, file_paths[0])
def pnet_fun(img):
return sess.run(out_tensor_pnet,
feed_dict={image_pnet: img})
saver_rnet.restore(sess, file_paths[1])
def rnet_fun(img):
return sess.run(out_tensor_rnet,
feed_dict={image_rnet: img})
saver_onet.restore(sess, file_paths[2])
def onet_fun(img):
return sess.run(out_tensor_onet,
feed_dict={image_onet: img})
else:
saver = tf.train.import_meta_graph(file_paths[0])
saver.restore(sess, file_paths[1])
def pnet_fun(img):
return sess.run(
('softmax/Reshape_1:0', 'pnet/conv4-2/BiasAdd:0'),
feed_dict={'Placeholder:0': img})
def rnet_fun(img):
return sess.run(('softmax_1/softmax:0',
'rnet/conv5-2/rnet/conv5-2:0'),
feed_dict={'Placeholder_1:0': img})
def onet_fun(img):
return sess.run(('softmax_2/softmax:0',
'onet/conv6-2/onet/conv6-2:0',
'onet/conv6-3/onet/conv6-3:0'),
feed_dict={'Placeholder_2:0': img})
rectangles, points = detect_face(img, minsize, pnet_fun,
rnet_fun, onet_fun, threshold,
factor)
tf.reset_default_graph()
return rectangles, points
def main(args):
img = cv2.imread(args.image_path)
file_paths = get_model_filenames(args.model_dir)
with tf.device('/gpu:0'):
with tf.Graph().as_default():
config = tf.ConfigProto(allow_soft_placement=True)
with tf.Session(config=config) as sess:
if len(file_paths) == 3:
image_pnet = tf.placeholder(tf.float32,
[None, None, None, 3])
pnet = PNet({'data': image_pnet}, mode='test')
out_tensor_pnet = pnet.get_all_output()
image_rnet = tf.placeholder(tf.float32, [None, 24, 24, 3])
rnet = RNet({'data': image_rnet}, mode='test')
out_tensor_rnet = rnet.get_all_output()
image_onet = tf.placeholder(tf.float32, [None, 48, 48, 3])
onet = ONet({'data': image_onet}, mode='test')
out_tensor_onet = onet.get_all_output()
saver_pnet = tf.train.Saver([
v for v in tf.global_variables()
if v.name[0:5] == "pnet/"
])
saver_rnet = tf.train.Saver([
v for v in tf.global_variables()
if v.name[0:5] == "rnet/"
])
saver_onet = tf.train.Saver([
v for v in tf.global_variables()
if v.name[0:5] == "onet/"
])
saver_pnet.restore(sess, file_paths[0])
def pnet_fun(img):
return sess.run(out_tensor_pnet,
feed_dict={image_pnet: img})
saver_rnet.restore(sess, file_paths[1])
def rnet_fun(img):
return sess.run(out_tensor_rnet,
feed_dict={image_rnet: img})
saver_onet.restore(sess, file_paths[2])
def onet_fun(img):
return sess.run(out_tensor_onet,
feed_dict={image_onet: img})
else:
saver = tf.train.import_meta_graph(file_paths[0])
saver.restore(sess, file_paths[1])
def pnet_fun(img):
return sess.run(
('softmax/Reshape_1:0', 'pnet/conv4-2/BiasAdd:0'),
feed_dict={'Placeholder:0': img})
def rnet_fun(img):
return sess.run(('softmax_1/softmax:0',
'rnet/conv5-2/rnet/conv5-2:0'),
feed_dict={'Placeholder_1:0': img})
def onet_fun(img):
return sess.run(('softmax_2/softmax:0',
'onet/conv6-2/onet/conv6-2:0',
'onet/conv6-3/onet/conv6-3:0'),
feed_dict={'Placeholder_2:0': img})
start_time = time.time()
rectangles, points = detect_face(img, args.minsize, pnet_fun,
rnet_fun, onet_fun,
args.threshold, args.factor)
duration = time.time() - start_time
print(duration)
print('rectangles->', rectangles)
print('pts->', points)
points = np.transpose(points)
for rectangle in rectangles:
cv2.putText(img, str(rectangle[4]),
(int(rectangle[0]), int(rectangle[1])),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
cv2.rectangle(img, (int(rectangle[0]), int(rectangle[1])),
(int(rectangle[2]), int(rectangle[3])),
(255, 0, 0), 1)
for point in points:
for i in range(0, 10, 2):
cv2.circle(img, (int(point[i]), int(point[i + 1])), 2,
(0, 255, 0))
cv2.imshow("test", img)
if args.save_image:
cv2.imwrite(args.save_name, img)
if cv2.waitKey(0) & 0xFF == ord('q'):
cv2.destroyAllWindows()