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
for databatch in image_reader:
if batch_idx % 100 == 0:
print("%d images done" % batch_idx)
im = databatch
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_%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 train_net(annotation_file,
model_store_path,
end_epoch=16,
frequent=200,
lr=0.01,
batch_size=128,
use_cuda=False):
imagedb = ImageDB(annotation_file)
gt_imdb = imagedb.load_imdb()
gt_imdb = imagedb.append_flipped_images(gt_imdb)
train_pnet(model_store_path=model_store_path,
end_epoch=end_epoch,
imdb=gt_imdb,
batch_size=batch_size,
frequent=frequent,
base_lr=lr,
use_cuda=use_cuda)