RetinaFace Pytorch实现训练、测试，pytorch模型转onnx转ncnn C++推理
首先感谢大佬的开源：
RetinaFace Pytorch：https://github.com/biubug6/Pytorch_Retinaface
pytorch转onnx转ncnn C++推理：https://github.com/biubug6/Face-Detector-1MB-with-landmark
RetinaFace Pytorch
训练和评估直接看github项目里的就可以了，主要提供了RetinaFace（resnet和mobilenet）、Slim和RFB网络结构（version-slim (network backbone simplification,slightly faster) and version-RFB (with the modified RFB module, higher precision)。基本思路类似于yolov3，在多个尺度的特征图上生成box。
下面给出单张图片和视频的推理代码，1050显卡推理速度在30ms左右。
import os
import sys
import os
import argparse
import torch
import torch.backends.cudnn as cudnn
import numpy as np
from data.config import cfg_mnet, cfg_re50
from layers.functions.prior_box import PriorBox
from utils.nms.py_cpu_nms import py_cpu_nms
import cv2
from models.retinaface import RetinaFace
from utils.box_utils import decode
import glob
import time
force_cpu = False
if force_cpu:
device = torch.device('cpu')
device = torch.cuda.current_device()
parser = argparse.ArgumentParser(description='Retinaface')
parser.add_argument('-m', '--trained_model', default='./weights/Resnet50_Final.pth',
type=str, help='Trained state_dict file path to open')
parser.add_argument('--origin_size', default=True, type=str,
help='Whether use origin image size to evaluate')
parser.add_argument('--img_folder', default='./images/',
type=str, help='dataset path')
parser.add_argument('--confidence_threshold', default=0.02,
type=float, help='confidence_threshold')
parser.add_argument('--top_k', default=5000, type=int, help='top_k')
parser.add_argument('--nms_threshold', default=0.3,
type=float, help='nms_threshold')
parser.add_argument('--keep_top_k', default=750, type=int, help='keep_top_k')
parser.add_argument('-s', '--show_image', action="store_true",
default=True, help='show detection results')
parser.add_argument('--vis_thres', default=0.3, type=float,
help='visualization_threshold')
args = parser.parse_args()
def check_keys(model, pretrained_state_dict):
ckpt_keys = set(pretrained_state_dict.keys())
model_keys = set(model.state_dict().keys())
used_pretrained_keys = model_keys & ckpt_keys
unused_pretrained_keys = ckpt_keys - model_keys
missing_keys = model_keys - ckpt_keys
print('Missing keys:{}'.format(len(missing_keys)))
print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))
print('Used keys:{}'.format(len(used_pretrained_keys)))
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
return True
def remove_prefix(state_dict, prefix):
''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''
print('remove prefix \'{}\''.format(prefix))
def f(x): return x.split(prefix, 1)[-1] if x.startswith(prefix) else x
return {f(key): value for key, value in state_dict.items()}
def load_model(model, pretrained_path):
print('Loading pretrained model from {}'.format(pretrained_path))
pretrained_dict = torch.load(pretrained_path,
map_location=lambda storage, loc: storage if force_cpu else 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
def detect_vis(net,img_raw):
img = np.float32(img_raw)
# testing scale
target_size = 1600
max_size = 2150
im_shape = img.shape
im_size_min = np.min(im_shape[0:2])
im_size_max = np.max(im_shape[0:2])
resize = float(target_size) / float(im_size_min)
# prevent bigger axis from being more than max_size:
if np.round(resize * im_size_max) > max_size:
resize = float(max_size) / float(im_size_max)
if args.origin_size:
resize = 1
if resize != 1:
img = cv2.resize(img, None, None, fx=resize,
fy=resize, interpolation=cv2.INTER_LINEAR)
im_height, im_width, _ = img.shape
scale = torch.Tensor(
[img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
img -= (104, 117, 123)
img = img.transpose(2, 0, 1)
img = torch.from_numpy(img).unsqueeze(0)
img = img.to(device)
scale = scale.to(device)
tic = time.time()
loc, conf, landms = net(img) # forward pass
print('net forward time: {}'.format(time.time() - tic))
priorbox = PriorBox(cfg_re50, image_size=(im_height, im_width))
priors = priorbox.forward()
priors = priors.to(device)
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg_re50['variance'])
boxes = boxes * scale / resize
boxes = boxes.cpu().numpy()
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > args.confidence_threshold)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:args.top_k]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(
np.float32, copy=False)
keep = py_cpu_nms(dets, args.nms_threshold)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:args.keep_top_k, :]
# show image
if args.show_image:
for b in dets:
if b[4] < args.vis_thres:
continue
text = "{:.4f}".format(b[4])
b = list(map(int, b))
cv2.rectangle(img_raw, (b[0], b[1]),
(b[2], b[3]), (0, 0, 255), 2)
cx = b[0]
cy = b[1] + 12
cv2.putText(img_raw, text, (cx, cy),
cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
cv2.imshow("res", img_raw)
def test_video(net):
cap = cv2.VideoCapture(0)
while cap.isOpened():
ret,img = cap.read()
detect_vis(net,img)
k = cv2.waitKey(1)
if k == ord('a') or k == ord('A'):
cv2.imwrite('test.jpg', img)
if k == ord('q') or k == ord('Q'):
cap.release()
def test_pic(net):
img_folder = args.img_folder
all_imgs = glob.glob(os.path.join(img_folder, '*.jpg'))
# testing begin
for i, img_f in enumerate(all_imgs):
img_raw = cv2.imread(img_f, cv2.IMREAD_COLOR)
img = np.float32(img_raw)
# testing scale
target_size = 1600
max_size = 2150
im_shape = img.shape
im_size_min = np.min(im_shape[0:2])
im_size_max = np.max(im_shape[0:2])
resize = float(target_size) / float(im_size_min)
# prevent bigger axis from being more than max_size:
if np.round(resize * im_size_max) > max_size:
resize = float(max_size) / float(im_size_max)
if args.origin_size:
resize = 1
if resize != 1:
img = cv2.resize(img, None, None, fx=resize,
fy=resize, interpolation=cv2.INTER_LINEAR)
im_height, im_width, _ = img.shape
scale = torch.Tensor(
[img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
img -= (104, 117, 123)
img = img.transpose(2, 0, 1)
img = torch.from_numpy(img).unsqueeze(0)
img = img.to(device)
scale = scale.to(device)
print('input tensor shape: {}'.format(img.size()))
tic = time.time()
loc, conf, landms = net(img) # forward pass
print('net forward time: {}'.format(time.time() - tic))
priorbox = PriorBox(cfg_re50, image_size=(im_height, im_width))
priors = priorbox.forward()
priors = priors.to(device)
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg_re50['variance'])
boxes = boxes * scale / resize
boxes = boxes.cpu().numpy()
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > args.confidence_threshold)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:args.top_k]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(
np.float32, copy=False)
keep = py_cpu_nms(dets, args.nms_threshold)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:args.keep_top_k, :]
# show image
if args.show_image:
for b in dets:
if b[4] < args.vis_thres:
continue
text = "{:.4f}".format(b[4])
b = list(map(int, b))
cv2.rectangle(img_raw, (b[0], b[1]),
(b[2], b[3]), (0, 0, 255), 2)
cx = b[0]
cy = b[1] + 12
cv2.putText(img_raw, text, (cx, cy),
cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
print(img_f.split("/")[-1][7:])
cv2.imwrite("./result/"+img_f.split("/")[-1][7:],img_raw)
if __name__ == '__main__':
torch.set_grad_enabled(False)
# net and model
net = RetinaFace(cfg=cfg_re50, phase='test')
net = load_model(net, args.trained_model)
net.eval()
print('Finished loading model!')
cudnn.benchmark = True
net = net.to(device)
test_video(net)
# test_pic(net)
pytorch转onnx转ncnn C++推理
环境：ubuntu16.04或win10；pytorch1.2（虽然作者说1.1.0以上就可以了，但是1.1.0版本onnx转换会有bug）；protobuf；opencv；ncnn。
ubuntu：
ubuntu下安装protobu、opencvf和ncnn还是比较简单的，收集了一些比较靠谱的博客。
protobuf安装：https://blog.csdn.net/u010918487/article/details/82947157
opencv安装：https://www.jianshu.com/p/f646448da265
ncnn安装：https://yyingbiu.github.io/2019/08/21/linux-xia-bian-yi-an-zhuang-ncnn/
执行以下指令，将pytorch模型转换成onnx格式
python convert_to_onnx.py --trained_model weight_file --network mobile0.25 or slim or RFB
pytorch生成onnx转换模型时有一些冗余，我们用工具简化一下onnx模型执行下面命令，转换就是正常的了，这个face_sim.onnx 就是最终的onnx文件。
pip install onnx-simplifier
python -m onnxsim face.onnx face_sim.onnx
后面按照github里面的流程来就可以了
win10
cmake、opencv和VS2015我之前安装过了，这里就不介绍了
推荐一个安装protobuf、ncnn比较靠谱的博客
https://blog.csdn.net/heiheiya/article/details/100519584
有以下几点要注意的：
1.protobuf的版本为3.4.0，我之前先安装了一个3.6.1版本的，在ncnn编译的时候onnx2caffe等工具的时候出了问题。
2.我是用命令行编译protobuf和ncnn的，这里要注意不能用普通的cmd窗口，要用Visual C++ 2015 x64 Native Build Tools Command Prompt，否则环境不行导致编译失败。
3.在编译ncnn之前要配置一下opencv，cmake可能找不到opencv的路径，在对应的CMakeLists.txt里面添加
set(OpenCV_DIR D:/mysoftware/opencv/build)
include_directories(${OpenCV_DIR}/include)
4.可以提前修改CMakeLists.txt里面相关路径，也可以直接执行以下指令后build里面有个vs2015工程.sln在vs里面配置include和lib路径，注意把ncnn的包含路径换成自己编译的，不要用github作者否则有重定义的bug！
mkdir build
cd build
cmake ..