def gen_onet_data(data_dir,
anno_file,
pnet_model_file,
rnet_model_file,
prefix_path='',
use_cuda=True,
vis=False):
pnet, rnet, _ = create_mtcnn_net(p_model_path=pnet_model_file,
r_model_path=rnet_model_file,
use_cuda=use_cuda)
mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, min_face_size=12)
imagedb = ImageDB(anno_file, mode="test", prefix_path=prefix_path)
imdb = imagedb.load_imdb()
image_reader = TestImageLoader(imdb, 1, False)
all_boxes = list()
batch_idx = 0
print('size:%d' % image_reader.size)
for databatch in image_reader:
if batch_idx % 50 == 0:
print("%d images done" % batch_idx)
im = databatch
t = time.time()
# pnet detection = [x1, y1, x2, y2, score, reg]
p_boxes, p_boxes_align = mtcnn_detector.detect_pnet(im=im)
# rnet detection
boxes, boxes_align = mtcnn_detector.detect_rnet(im=im,
dets=p_boxes_align)
if boxes_align is None:
all_boxes.append(np.array([]))
batch_idx += 1
continue
if vis:
rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB)
vision.vis_two(rgb_im, boxes, boxes_align)
t1 = time.time() - t
t = time.time()
all_boxes.append(boxes_align)
batch_idx += 1
save_path = './model_store'
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time()))
with open(save_file, 'wb') as f:
cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)
gen_onet_sample_data(data_dir, anno_file, save_file, prefix_path)
def gen_rnet_data(data_dir, anno_dir, pnet_model_file, use_cuda=True):
''' Generate the train data of RNet with trained-PNet '''
# load trained pnet model
pnet, _, _ = create_mtcnn_net(p_model_path=pnet_model_file,
use_cuda=use_cuda)
mtcnn_detector = MtcnnDetector(pnet=pnet, min_face_size=12)
# load original_anno_file, length = 12880
anno_file = os.path.join(anno_dir, 'anno_store/wide_anno_train.txt'
) # TODO :: [local_wide_anno, wide_anno_train]
imagedb = ImageDB(anno_file, mode='test', prefix_path='')
imdb = imagedb.load_imdb()
image_reader = TestImageLoader(imdb, 1, False)
print('size:%d' % image_reader.size)
batch_idx, all_boxes = 0, list()
for databatch in image_reader:
if (batch_idx + 1) % 100 == 0:
print("%d images done" % (batch_idx + 1))
im = databatch
# obtain boxes and aligned boxes
boxes_align = mtcnn_detector.detect_pnet(im=im) # Time costly
if boxes_align is None:
all_boxes.append(np.array([]))
batch_idx += 1
continue
# if vis:
# rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB)
# vision.vis_two(rgb_im, boxes[:100, :], boxes_align[:100, :])
all_boxes.append(boxes_align)
batch_idx += 1
save_path = os.path.join(anno_dir, 'rnet')
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time()))
with open(save_file, 'wb') as f:
cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)
gen_rnet_sample_data(data_dir, anno_dir, save_file)
def gen_rnet_data(data_dir,
anno_file,
pnet_model_file,
prefix_path='',
use_cuda=True,
vis=False):
"""
:param data_dir: train data
:param anno_file:
:param pnet_model_file:
:param prefix_path:
:param use_cuda:
:param vis:
:return:
"""
# load trained pnet model
pnet, _, _ = create_mtcnn_net(p_model_path=pnet_model_file,
use_cuda=use_cuda)
mtcnn_detector = MtcnnDetector(pnet=pnet, min_face_size=12)
# load original_anno_file, length = 12880
imagedb = ImageDB(anno_file, mode="test", prefix_path=prefix_path)
imdb = imagedb.load_imdb()
image_reader = TestImageLoader(imdb, 1, False)
all_boxes = list()
batch_idx = 0
print('size:%d' % image_reader.size)
for databatch in image_reader:
if batch_idx % 100 == 0:
print("%d images done" % batch_idx)
im = databatch
t = time.time()
# obtain boxes and aligned boxes
boxes, boxes_align = mtcnn_detector.detect_pnet(im=im)
if boxes_align is None:
all_boxes.append(np.array([]))
batch_idx += 1
continue
if vis:
rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB)
vision.vis_two(rgb_im, boxes, boxes_align)
t1 = time.time() - t
t = time.time()
all_boxes.append(boxes_align)
batch_idx += 1
# if batch_idx == 100:
# break
# print("shape of all boxes {0}".format(all_boxes))
# time.sleep(5)
# save_path = model_store_path()
# './model_store'
save_path = './model_store'
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time()))
with open(save_file, 'wb') as f:
cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)
gen_rnet_sample_data(data_dir, anno_file, save_file, prefix_path)
from mtcnn.config import *
pnet, rnet, onet = create_mtcnn_net(p_model_path=PNET_MODEL_PATH,r_model_path=RNET_MODEL_PATH,
o_model_path=ONET_MODEL_PATH,use_cuda=False)
mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=onet, min_face_size=48)
cap = cv2.VideoCapture(0)
f = 0
stime = time.time()
while (True):
ret, frame = cap.read() # 读取一帧的图像
frame = cv2.resize(frame,(480,360))
# frame = cv2.imread('face.jpg')
boxes, boxes_align = mtcnn_detector.detect_pnet(im=frame)
rboxes, rboxes_align = mtcnn_detector.detect_rnet(im=frame, dets=boxes_align)
# oboxes,olandmark = mtcnn_detector.detect_onet(im=frame,dets=rboxes_align)
if rboxes_align is not None:
for box in rboxes_align:
cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 0, 128), 2)
if f % 20 == 0:
fps = int(20/(time.time()-stime))
f = 0
stime = time.time()
cv2.putText(frame, '{:d}fps'.format(fps), (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8,(255, 0, 255), 2)
cv2.imshow('Face Recognition', frame)
f += 1
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release() # 释放摄像头
def gen_landmark48_data(data_dir,
anno_file,
pnet_model_file,
rnet_model_file,
prefix_path='',
use_cuda=True,
vis=False):
anno_file = os.path.join(data_dir, anno_file)
pnet, rnet, _ = create_mtcnn_net(p_model_path=pnet_model_file,
r_model_path=rnet_model_file,
use_cuda=use_cuda)
mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, min_face_size=12)
imagedb = ImageDB(anno_file,
mode="test",
prefix_path=os.path.join(data_dir, 'img'))
imdb = imagedb.load_imdb()
image_reader = TestImageLoader(imdb, 1, False)
all_boxes = list()
batch_idx = 0
for databatch in image_reader:
if batch_idx % 500 == 0:
print("%d images done" % batch_idx)
im = databatch
if im.shape[0] >= 1200 or im.shape[1] >= 1200:
all_boxes.append(np.array([]))
batch_idx += 1
continue
t = time.time()
p_boxes, p_boxes_align = mtcnn_detector.detect_pnet(im=im)
boxes, boxes_align = mtcnn_detector.detect_rnet(im=im,
dets=p_boxes_align)
if boxes_align is None:
all_boxes.append(np.array([]))
batch_idx += 1
continue
if vis:
rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB)
vision.vis_two(rgb_im, boxes, boxes_align)
t1 = time.time() - t
t = time.time()
all_boxes.append(boxes_align)
batch_idx += 1
save_path = config.MODEL_STORE_DIR
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections_celeba.pkl")
with open(save_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
