def train_net(args):
imagedb = ImageDB(args.anno_file)
train_imdb = imagedb.load_imdb()
# train_imdb = imagedb.append_flipped_images(train_imdb)
imagedb = ImageDB(args.eval_file)
eval_imdb = imagedb.load_imdb()
print('train : %d\teval : %d' % (len(train_imdb), len(eval_imdb)))
print(args)
train.train_onet(args, train_imdb, eval_imdb)
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 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.train_onet(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)
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 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_data = TrainImageReader(gt_imdb, 24, 16, shuffle=False)
accuracy_list = []
cls_loss_list = []
bbox_loss_list = []
for batch_idx, (image, (gt_label, gt_bbox,
gt_landmark)) in enumerate(train_data):
print(len(image))
print(gt_bbox)
break
def train_net(annotation_file,
model_store_path,
end_epoch=16,
frequent=200,
lr=0.01,
lr_epoch_decay=[9],
batch_size=128,
use_cuda=False,
load=''):
imagedb = ImageDB(annotation_file)
gt_imdb = imagedb.load_imdb()
print('DATASIZE', len(gt_imdb))
gt_imdb = imagedb.append_flipped_images(gt_imdb)
print('FLIP DATASIZE', len(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,
lr_epoch_decay=lr_epoch_decay,
use_cuda=use_cuda,
load=load)
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)
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)
def test_Onet_without_PRnet(self, annotation, outputDir, test_moudel, xxyy,
savePic):
imagedb = ImageDB(annotation)
gt_imdb = imagedb.load_imdb()
pnet, rnet, onet_jiang = create_mtcnn_net(
p_model_path="./original_model/pnet_epoch.pt",
r_model_path="./original_model/rnet_epoch.pt",
o_model_path="./original_model/" + test_moudel + ".pt",
use_cuda=False)
mtcnn_detector = MtcnnDetector(pnet=pnet,
rnet=rnet,
onet=onet_jiang,
min_face_size=24)
test_data = TrainImageReader(gt_imdb,
48,
batch_size=100,
shuffle=False) # 读入1个batch的数据
# train_data.reset()
total_errors = 0
cnt = 0
for i, (images, (gt_labels, gt_bboxes,
gt_landmarks)) in enumerate(test_data): # 取1个batch
list_imgs = [images[i, :, :, :]
for i in range(images.shape[0])] # 100张图片
list_bboxes = [gt_bboxes[i, :] for i in range(gt_bboxes.shape[0])]
list_gt_landmarks = [
gt_landmarks[i, :] for i in range(gt_landmarks.shape[0])
]
mix = list(zip(list_imgs, list_bboxes, list_gt_landmarks))
batch_errors = []
for img, gt_bbox, gt_landmark in mix: # 取1个图片
if xxyy:
bboxs, landmarks = mtcnn_detector.detect_onet_xxyy(
img, gt_bbox) # 原始的图片用原始网络检测,xxyy
else:
bboxs, landmarks = mtcnn_detector.detect_onet(
img, gt_bbox) # 原始的图片用原始网络检测,xxyy
if landmarks.size:
cnt += 1
bboxs = bboxs[:1] # 多个检测框保留第一个
landmarks = landmarks[:1]
if savePic:
vis_face(img, bboxs, landmarks,
self.output_dir + str(cnt) + ".jpg") # 保存图片
gt_landmark = np.array(gt_landmark).reshape(5, 2)
landmarks = np.array(landmarks).reshape(5, 2)
normDist = np.linalg.norm(gt_landmark[1] -
gt_landmark[0]) # 左右眼距离
error = np.mean(
np.sqrt(np.sum(
(landmarks - gt_landmark)**2, axis=1))) / normDist
batch_errors.append(error)
batch_errors = np.array(batch_errors).sum()
total_errors += batch_errors
print("%s: %s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
if cnt > 999:
print("%s:%s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
f = open("landmark_test.txt", "a+")
f.write("%s, moudel_name:%s.pt, %s pics mean error is %s\n" %
(datetime.datetime.now(), test_moudel, cnt,
np.array(total_errors).reshape(1, -1).sum() / cnt))
f.close()
return
print("%s:%s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
class LandmarkTestor(object):
def __init__(self, annotation, output_dir):
self.output_dir = output_dir
self.imagedb = ImageDB(annotation)
self.gt_imdb = self.imagedb.load_imdb()
def test_face_alignment(self, test_moudel, savePic):
pnet, rnet, onet_jiang = create_mtcnn_net(
p_model_path="./original_model/pnet_epoch.pt",
r_model_path="./original_model/rnet_epoch.pt",
o_model_path="./original_model/" + test_moudel + ".pt",
use_cuda=False)
mtcnn_detector = MtcnnDetector(pnet=pnet,
rnet=rnet,
onet=onet_jiang,
min_face_size=24)
train_data = TrainImageReader(self.gt_imdb,
48,
batch_size=100,
shuffle=False) # 读入1个batch的数据
# train_data.reset()
total_errors = 0
cnt = 0
for i, (images, (gt_labels, gt_bboxes,
gt_landmarks)) in enumerate(train_data): #取1个batch
# if i == 12:
# print(i)
list_imgs = [images[i, :, :, :]
for i in range(images.shape[0])] # 100张图片
list_bboxes = [gt_bboxes[i, :] for i in range(gt_bboxes.shape[0])]
list_gt_landmarks = [
gt_landmarks[i, :] for i in range(gt_landmarks.shape[0])
]
mix = list(zip(list_imgs, list_bboxes, list_gt_landmarks))
batch_errors = []
for img, gt_bbox, gt_landmark in mix: # 取1个图片
bboxs, landmarks = mtcnn_detector.detect_face(
img) # 原始的图片用原始网络检测
if landmarks.size:
cnt += 1
bboxs = bboxs[:1] # 多个检测框保留第一个
landmarks = landmarks[:1]
if savePic:
vis_face(img, bboxs, landmarks,
self.output_dir + str(cnt) + ".jpg") # 保存图片
gt_landmark = np.array(gt_landmark).reshape(5, 2)
landmarks = np.array(landmarks).reshape(5, 2)
normDist = np.linalg.norm(gt_landmark[1] -
gt_landmark[0]) # 左右眼距离
error = np.mean(
np.sqrt(np.sum(
(landmarks - gt_landmark)**2, axis=1))) / normDist
# print("the %sth pic error is : %s"%(cnt, error))
batch_errors.append(error)
batch_errors = np.array(batch_errors).sum()
total_errors += batch_errors
print("%s: %s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
f = open("landmark_test.txt", "a+")
f.write("%s, moudel_name:%s.pt, %s pics mean error is %s\n" %
(datetime.datetime.now(), test_moudel, cnt,
np.array(total_errors).reshape(1, -1).sum() / cnt))
f.close()
print("%s:%s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
def test_Onet_without_PRnet(self, annotation, outputDir, test_moudel, xxyy,
savePic):
imagedb = ImageDB(annotation)
gt_imdb = imagedb.load_imdb()
pnet, rnet, onet_jiang = create_mtcnn_net(
p_model_path="./original_model/pnet_epoch.pt",
r_model_path="./original_model/rnet_epoch.pt",
o_model_path="./original_model/" + test_moudel + ".pt",
use_cuda=False)
mtcnn_detector = MtcnnDetector(pnet=pnet,
rnet=rnet,
onet=onet_jiang,
min_face_size=24)
test_data = TrainImageReader(gt_imdb,
48,
batch_size=100,
shuffle=False) # 读入1个batch的数据
# train_data.reset()
total_errors = 0
cnt = 0
for i, (images, (gt_labels, gt_bboxes,
gt_landmarks)) in enumerate(test_data): # 取1个batch
list_imgs = [images[i, :, :, :]
for i in range(images.shape[0])] # 100张图片
list_bboxes = [gt_bboxes[i, :] for i in range(gt_bboxes.shape[0])]
list_gt_landmarks = [
gt_landmarks[i, :] for i in range(gt_landmarks.shape[0])
]
mix = list(zip(list_imgs, list_bboxes, list_gt_landmarks))
batch_errors = []
for img, gt_bbox, gt_landmark in mix: # 取1个图片
if xxyy:
bboxs, landmarks = mtcnn_detector.detect_onet_xxyy(
img, gt_bbox) # 原始的图片用原始网络检测,xxyy
else:
bboxs, landmarks = mtcnn_detector.detect_onet(
img, gt_bbox) # 原始的图片用原始网络检测,xxyy
if landmarks.size:
cnt += 1
bboxs = bboxs[:1] # 多个检测框保留第一个
landmarks = landmarks[:1]
if savePic:
vis_face(img, bboxs, landmarks,
self.output_dir + str(cnt) + ".jpg") # 保存图片
gt_landmark = np.array(gt_landmark).reshape(5, 2)
landmarks = np.array(landmarks).reshape(5, 2)
normDist = np.linalg.norm(gt_landmark[1] -
gt_landmark[0]) # 左右眼距离
error = np.mean(
np.sqrt(np.sum(
(landmarks - gt_landmark)**2, axis=1))) / normDist
batch_errors.append(error)
batch_errors = np.array(batch_errors).sum()
total_errors += batch_errors
print("%s: %s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
if cnt > 999:
print("%s:%s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
f = open("landmark_test.txt", "a+")
f.write("%s, moudel_name:%s.pt, %s pics mean error is %s\n" %
(datetime.datetime.now(), test_moudel, cnt,
np.array(total_errors).reshape(1, -1).sum() / cnt))
f.close()
return
print("%s:%s pics mean error is %s" %
(datetime.datetime.now(), cnt, total_errors / cnt))
def test_Onet_without_PRnet1(inputDir, outputDir, model):
pnet, rnet, onet_jiang = create_mtcnn_net(
p_model_path="./original_model/pnet_epoch.pt",
r_model_path="./original_model/rnet_epoch.pt",
o_model_path="./original_model/" + model + ".pt",
use_cuda=False)
mtcnn_detector = MtcnnDetector(pnet=pnet,
rnet=rnet,
onet=onet_jiang,
min_face_size=24)
files = os.listdir(inputDir)
i = 0
for image in files:
i += 1
image = os.path.join(inputDir, image)
img = cv2.imread(image)
img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
landmarks2_jiang = mtcnn_detector.detect_onet_raw(img)
vis_face_test(img_bg, landmarks2_jiang,
outputDir + model + "-" + str(i) + ".jpg")
if i == 50:
break
def test_img_and_label_48(self,
label_path,
store=True,
store_path="./landmark_show/"):
if store and not os.path.exists(store_path):
os.makedirs(store_path)
with open(label_path, 'r') as f:
annotations = f.readlines()
im_idx_list = list()
num_of_images = len(annotations)
print("processing %d images in total" % num_of_images)
i = 0
for annotation in annotations:
i += 1
annotation = annotation.strip().split(' ')
im_idx = annotation[0]
landmark = list(map(float, annotation[6:]))
landmark = np.array(landmark, dtype=np.float)
img_test = cv2.imread(im_idx)
img_test = cv2.cvtColor(img_test, cv2.COLOR_BGR2RGB)
vis_face_landmark_img_label(img_test, landmark,
"./landmark_show/" + str(i) + ".jpg")