def main(args):
torch.set_grad_enabled(False)
device = torch.device(args.device)
# Initialize the net and load the model
print('Loading pretrained model from {}'.format(args.trained_model))
net = YuFaceDetectNet(phase='test', size=None)
net = load_model(net, args.trained_model)
net.eval()
if device.type == 'cuda':
cudnn.benchmark = True
net = net.to(device)
print('Finished loading model!')
# init data loader for WIDER Face
print('Loading data for {}...'.format(args.widerface_split))
widerface = WIDERFace(args.widerface_root, split=args.widerface_split)
print('Finished loading data!')
# start testing
scales = []
if args.multi_scale:
scales = [0.25, 0.50, 0.75, 1.25, 1.50, 1.75, 2.0]
print('Performing testing with scales: 1. {}, conf_threshold: {}'.format(
str(scales), args.confidence_threshold))
priors_dict = {}
for idx in tqdm(range(len(widerface))):
img, event, name = widerface[idx] # img_subpath = '0--Parade/XXX.jpg'
if img.shape in priors_dict:
priors = priors_dict[img.shape]
else:
height, width, _ = img.shape
priors = PriorBox(cfg, image_size=(height, width)).forward()
priors_dict[img.shape] = priors
dets = detect_face(net, img, priors, device)
available_scales = get_available_scales(img.shape[0], img.shape[1],
scales)
for available_scale in available_scales:
det = detect_face(net, img, None, device, scale=available_scale)
if det.shape[0] != 0:
dets = np.row_stack((dets, det))
# nms
dets = nms_opencv(dets,
score_thresh=args.confidence_threshold,
nms_thresh=args.nms_threshold,
top_k=args.top_k,
keep_top_k=args.keep_top_k)
save_res(dets, event, name)
# widerface_eval
print('Evaluating:')
evaluation(args.res_dir,
os.path.join(args.widerface_root, 'eval_tools/ground_truth'))
def main(args):
torch.set_grad_enabled(False)
device = torch.device(args.device)
# Initialize the net and load the model
print('Loading pretrained model from {}'.format(args.trained_model))
net = YuFaceDetectNet(phase='test', size=None)
net = load_model(net, args.trained_model)
net.eval()
if device.type == 'cuda':
cudnn.benchmark = True
net = net.to(device)
print('Finished loading model!')
# init data loader for WIDER Face
print('Loading data for {}...'.format(args.widerface_split))
widerface = WIDERFace(args.widerface_root, split=args.widerface_split)
print('Finished loading data!')
# start testing
scales = [1.]
if args.multi_scale:
scales = [0.25, 0.50, 0.75, 1.25, 1.50, 1.75, 2.0]
print('Performing testing with scales: {}, conf_threshold: {}'.format(
str(scales), args.confidence_threshold))
for idx in tqdm(range(len(widerface))):
img, event, name = widerface[idx] # img_subpath = '0--Parade/XXX.jpg'
dets = detect_face(net,
img,
device,
conf_thresh=args.confidence_threshold)
available_scales = get_available_scales(img.shape[0], img.shape[1],
scales)
for available_scale in available_scales:
det = detect_face(net,
img,
device,
scale=available_scale,
conf_thresh=args.confidence_threshold)
if det is not None: dets = np.row_stack((dets, det))
# nms
dets = simple_nms(dets)
# dets = bbox_vote(dets)
save_res(dets, event, name)
return {f(key): value for key, value in state_dict.items()}
def load_model(model, pretrained_path, load_to_cpu):
print('Loading pretrained model from {}'.format(pretrained_path))
if load_to_cpu:
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
else:
device = torch.cuda.current_device()
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
else:
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
check_keys(model, pretrained_dict)
model.load_state_dict(pretrained_dict, strict=False)
return model
if __name__ == '__main__':
torch.set_grad_enabled(False)
# net and model
net = YuFaceDetectNet(phase='test', size=None ) # initialize detector
net = load_model(net, args.trained_model, True)
net.eval()
print('Finished loading model!')
net.export_cpp(args.output)