def gen_onet_data(data_dir,
anno_file,
pnet_model_file,
rnet_model_file,
prefix_path='',
use_cuda=True,
vis=False):
mtcnn_detector = MtcnnDetector(p_model_path=pnet_model_file,
r_model_path=rnet_model_file,
o_model_path=None,
min_face_size=12,
use_cuda=True)
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()
# detect an image by pnet and rnet
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:
vision.vis_face(im, boxes_align)
t1 = time.time() - t
print('time cost for image ', batch_idx, '/', image_reader.size, ': ',
t1)
all_boxes.append(boxes_align)
batch_idx += 1
save_path = config.TRAIN_DATA_DIR
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path,
"pnet_rnet_detections_%d.pkl" % int(time.time()))
with open(save_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
# save_file = '/home/liujing/Codes/MTCNN/data/pnet_detections_1532582821.pkl'
get_onet_sample_data(data_dir, anno_file, save_file, prefix_path)
def gen_rnet_data(data_dir,
anno_file,
pnet_model_file,
prefix_path='',
use_cuda=True,
vis=False):
# load the pnet and pnet_detector ,利用刚刚训练的pnet网络生成rnet的训练数据
mtcnn_detector = MtcnnDetector(
p_model_path=pnet_model_file, # pnet_model_file自行设置
r_model_path=None,
o_model_path=None,
min_face_size=12,
use_cuda=True)
device = mtcnn_detector.device
# 生成rnet网络的训练集
imagedb = ImageDB(anno_file, mode="test", prefix_path=prefix_path)
imdb = imagedb.load_imdb()
image_reader = TestImageLoader(imdb, 1, False)
all_boxes = []
batch_idx = 0
for databatch in image_reader:
if batch_idx % 100 == 0:
print("%d images done" % batch_idx)
im = databatch
t = time.time()
boxes, boxes_align = mtcnn_detector.detect_pnet(im)
if boxes_align is None:
all_boxes.append(np.array([]))
continue
if vis:
vision.vis_face(im, boxes_align)
t1 = time.time() - t
print('time cost for image {} / {} : {:.4f}'.format(
batch_idx, image_reader.size, t1))
all_boxes.append(boxes_align)
batch_idx += 1
save_path = config.TRAIN_DATA_DIR
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path,
"pnet_detections_%d.pkl" % int(time.time()))
with open(save_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
# save_file = '/home/liujing/Codes/MTCNN/data/pnet_detections_1532530263.pkl'
get_rnet_sample_data(data_dir, anno_file, save_file, prefix_path)
if attrs != []:
for a in attrs:
txt = plate_color[a[0]] + ' ' + plate_layer[
a[1]] + ' ' + plate_type[a[2]]
cv2.putText(img, txt, (int(bboxs[i, 0]), int(bboxs[i, 1])),
cv2.FONT_HERSHEY_COMPLEX, 2, (0, 0, 255), 5)
i += 1
return img
if __name__ == '__main__':
mtcnn_detector = MtcnnDetector(
p_model_path=
"./plate_annotations/pnet/log_bs512_lr0.010_0930/check_point/model_050.pth",
r_model_path=
"./plate_annotations/rnet/log_bs512_lr0.001_0930/check_point/model_050.pth",
o_model_path=
"./plate_annotations/onet/log_bs512_lr0.001_0930/check_point/model_050.pth",
min_face_size=60,
use_cuda=False) # 加载模型参数，构造检测器
logger.info("Init the MtcnnDetector.")
project_root = pathlib.Path()
inputPath = project_root / "data" / "test_new_energy"
outputPath = project_root / "data" / "test_new_energy/results"
outputPath.mkdir(exist_ok=True)
start = time.time()
corr = 0
count = 0
for num, input_img_filename in enumerate(inputPath.iterdir()):
logger.info("Start to process No.{} image.".format(num))
cv2.rectangle(img, (int(bboxs[i, 0]), int(bboxs[i, 1])),
(int(bboxs[i, 2]), int(bboxs[i, 3])), (0, 255, 0), 3)
for p in landmarks:
for i in range(5):
cv2.circle(img, (int(p[2 * i]), int(p[2 * i + 1])), 3, (0, 0, 255),
-1)
return img
if __name__ == '__main__':
# set pnet, rnet, onet model path as your training model path
mtcnn_detector = MtcnnDetector(
p_model_path=
"./results/pnet/log_bs512_lr0.010_072402/check_point/model_050.pth",
r_model_path=
"./results/rnet/log_bs512_lr0.001_072502/check_point/model_050.pth",
o_model_path=
"./results/onet/log_bs512_lr0.001_072602/check_point/model_050.pth",
min_face_size=24,
use_cuda=True) # 加载模型参数，构造检测器
logger.info("Init the MtcnnDetector.")
project_root = pathlib.Path()
inputPath = project_root / "data" / "test_images"
outputPath = project_root / "data" / "you_result"
outputPath.mkdir(exist_ok=True)
with torch.no_grad():
start = time.time()
for num, input_img_filename in enumerate(inputPath.iterdir()):
logger.info("Start to predict No.{} image.".format(num))
img_name = input_img_filename.name
file_name = 'FDDB-folds/FDDB-folds-%02d.txt' % (n + 1)
file_name = os.path.join(data_dir, file_name)
fid = open(file_name, 'r')
image_names = []
for im_name in fid.readlines():
image_names.append(im_name.strip('\n'))
imdb.append(image_names)
return imdb
if __name__ == "__main__":
mtcnn_detector = MtcnnDetector(
p_model_path=
'./results/pnet/log_bs512_lr0.010_072402/check_point/model_050.pth',
r_model_path=
'./results/rnet/log_bs512_lr0.001_072502/check_point/model_050.pth',
o_model_path=
'./results/pnet/log_bs512_lr0.001_0726402/check_point/model_050.pth',
min_face_size=12,
use_cuda=False)
# logger.info("Init the MtcnnDetector.")
imdb = get_imdb_fddb(data_dir)
nfold = len(imdb)
for i in range(nfold):
image_names = imdb[i]
dets_file_name = os.path.join(out_dir,
'FDDB-det-fold-%02d.txt' % (i + 1))
fid = open(dets_file_name, 'w')
# image_names_abs = [os.path.join(data_dir, 'originalPics', image_name + '.jpg') for image_name in image_names]