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)
'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__':
#img_dim = 320
device = torch.device(args.device)
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.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
## Print model's state_dict
#print("Model's state_dict:")
#for param_tensor in net.state_dict():
# print(param_tensor, "\t", net.state_dict()[param_tensor].size())
#print(net)
cudnn.benchmark = True
img_dim = 320 # only 1024 is supported
rgb_mean = (0, 0, 0) #(104, 117, 123) # bgr order
num_classes = 2
gpu_ids = [int(item) for item in args.gpu_ids.split(',')]
num_workers = args.num_workers
batch_size = args.batch_size
momentum = args.momentum
weight_decay = args.weight_decay
initial_lr = args.lr
gamma = args.gamma
max_epoch = args.max_epoch
training_face_rect_dir = args.training_face_rect_dir
training_face_landmark_dir = args.training_face_landmark_dir
net = YuFaceDetectNet('train', img_dim)
print("Printing net...")
#print(net)
if args.resume_net is not None:
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
img_dim = 320 # only 1024 is supported
rgb_mean = (0,0,0) #(104, 117, 123) # bgr order
num_classes = 2
gpu_ids = [int(item) for item in args.gpu_ids.split(',')]
num_workers = args.num_workers
#batch_size = args.batch_size
momentum = args.momentum
weight_decay = args.weight_decay
#initial_lr = args.lr
gamma = args.gamma
max_epoch = args.max_epoch
training_face_rect_dir = args.training_face_rect_dir
training_face_landmark_dir = args.training_face_landmark_dir
net = YuFaceDetectNet('train', img_dim)
print("Printing net...")
#print(net)
if args.resume_net is not None:
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
if not filepath:
cv2.imshow("Press 'Esc' to exit.", detections)
video_writer.write(detections)
if cv2.waitKey(1) == 27:
break
video_reader.release()
video_writer.release()
cv2.destroyAllWindows()
return None
if __name__ == '__main__':
device = torch.device('cuda:0')
torch.set_grad_enabled(False)
# net and model, initialize detector
net = YuFaceDetectNet(phase='test', size=None)
net = load_model(net, 'weights/yunet_final.pth', load_to_cpu=True)
net.eval()
print('Finished loading model!')
cudnn.benchmark = True
net = net.to(device)
file_list = os.listdir('./input_image/')
for i in file_list:
print('Dealing with {}...'.format(i))
file_id = './input_image/' + i
result = deal_img(file_id, save_result=True, print_messages=False)
file_list = os.listdir('./input_video/')
for i in file_list:
print('Dealing with {}...'.format(i))