retinaface onnx

tensorrt代码：

https://github.com/azhe198827/retinaface_tensorRT

 

https://github.com/lesurp/retinanet_cpp/blob/master/examples/single_image.cpp

https://github.com/wuchaodzxx/tensorrt_retinaface

废话不说，上代码

retina face，

import os
import time
from math import ceil

import onnxruntime
import numpy as np
import cv2
import argparse
import argparse
import numpy as np
from data import cfg_mnet, cfg_peleenet
from utils.nms.py_cpu_nms import py_cpu_nms
from math import ceil
from itertools import product as product


#sigmoid函数
def sigmoid(x):
    s = 1 / (1 + np.exp(-1*x))
    return s
def softmax(x, axis=1):
    # 计算每行的最大值
    row_max = x.max(axis=axis)

    # 每行元素都需要减去对应的最大值，否则求exp(x)会溢出，导致inf情况
    row_max = row_max.reshape(-1, 1)
    x = x - row_max

    x_exp = np.exp(x)
    x_sum = np.sum(x_exp, axis=axis, keepdims=True)
    s = x_exp / x_sum
    return s

def decode_landm(pre, priors, variances):
    """Decode landm from predictions using priors to undo
    the encoding we did for offset regression at train time.
    Args:
        pre (tensor): landm predictions for loc layers,
            Shape: [num_priors,10]
        priors (tensor): Prior boxes in center-offset form.
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        decoded landm predictions
    """
    landms = np.concatenate((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
                        ), 1)
    return landms

def decode(loc, priors, variances):
    """Decode locations from predictions using priors to undo
    the encoding we did for offset regression at train time.
    Args:
        loc (tensor): location predictions for loc layers,
            Shape: [num_priors,4]
        priors (tensor): Prior boxes in center-offset form.
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        decoded bounding box predictions
    """

    boxes = np.concatenate((
        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
        priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1)
    boxes[:, :2] -= boxes[:, 2:] / 2
    boxes[:, 2:] += boxes[:, :2]
    return boxes

class PriorBox(object):
    def __init__(self, cfg, image_size=None, phase='train'):
        super(PriorBox, self).__init__()
        self.min_sizes = cfg['min_sizes']
        self.steps = cfg['steps']
        self.image_size = image_size
        self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]

    def forward(self):
        anchors = []
        for k, f in enumerate(self.feature_maps):
            min_sizes = self.min_sizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / self.image_size[1]
                    s_ky = min_size / self.image_size[0]
                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]

        # back to torch land
        output = np.array(anchors)
        output = output.reshape(-1, 4)
        return output
def sigmoid(x):
    # TODO: Implement sigmoid function
    return 1/(1 + np.exp(-x))
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_shape', help="caffe's caffemodel file path",default=(480,360))
    parser.add_argument('--img_path', help="test image path", default="d:/cat.jpg")
    parser.add_argument('--onnx_path', help="onnx model file path",  default="mobileretina.onnx")
    # parser.add_argument('--onnx_path', help="onnx model file path",  default=r"pelee_detector.onnx")
    # parser.add_argument('--onnx_path', help="onnx model file path",  default="yolov3.onnx")
    parser.add_argument('--confidence_threshold', help="confidence threshold", default=0.9, type=float)
    parser.add_argument('--nms_thres', help="nms threshold",  default=0.6, type=float)
    parser.add_argument('--top_k',help="to choice anchors", default=20, type=int)
    parser.add_argument('--nms_threshold',help="to choice anchors", default=0.4, type=float)
    parser.add_argument('--show_image',help="to choice anchors", default=True, type=bool)
    args = parser.parse_args()

    cfg = cfg_peleenet
    onnx_path = args.onnx_path
    session = onnxruntime.InferenceSession(onnx_path)
    input_shape = args.input_shape #模型输入尺寸
    nms_threshold = args.nms_thres
    img_path = args.img_path

    print("image path:",img_path)
    print("onnx model path:",onnx_path)

    # list_path = r"D:\project\face\face_mask\2020\0/"
    list_path = r"D:\input\faces/"

    g = os.walk(list_path)
    files = ['%s\\%s' % (i[0], j) for i in g for j in i[-1] if
             j.endswith('jpg')]
    width=input_shape[0]
    height=input_shape[1]
    scale = np.array([width, height, width, height])

    scale1 = np.array([width, height, width, height,
                           width, height, width, height,
                           width, height])

    resize_level=1
    count = 0
    ok_count = 0
    priorbox = PriorBox(cfg, image_size=(height, width))
    priors = priorbox.forward()

    # vc = cv2.VideoCapture(r"D:\project\face\Face-Track-Detect-Extract\videos\2_Obama.mp4")  # 读入视频文件
    vc = cv2.VideoCapture(0)  # 读入视频文件

    while True:  # 循环读取视频帧
        rval, img_raw = vc.read()

    # for file in files:
    #     file=r"d:/7_Cheering_Cheering_7_426.png"
    #     img_raw = cv2.imread(file)
        if img_raw is None:
            # print(file)
            continue
        start = time.time()

        img_raw = cv2.resize(img_raw, input_shape)
        img=cv2.cvtColor(img_raw,cv2.COLOR_BGR2RGB)
        img = np.float32(img)
        img -= (104, 117, 123)
        image = img[:, :, ::-1].transpose((2, 0, 1))
        TestData = image[np.newaxis, :, :, :]

        start2=time.time()
        inname = [input.name for input in session.get_inputs()][0]
        outname = [output.name for output in session.get_outputs()]

        loc, conf,landmarks = session.run(outname, {inname:TestData})
        print('net time', time.time() - start2)
        start1=time.time()

        boxes = decode(np.squeeze(loc, axis=0), priors, cfg['variance'])
        boxes = boxes * scale / resize_level
        scores = np.squeeze(conf, axis=0)[:,1]

        landmarks = decode_landm(np.squeeze(landmarks, axis=0), priors, cfg['variance'])
        landmarks = landmarks * scale1 / resize_level
        # ignore low scores
        inds = np.where(scores > args.confidence_threshold)[0]
        boxes = boxes[inds]
        landmarks = landmarks[inds]
        scores = scores[inds]

        # keep top-K before NMS
        order = scores.argsort()[::-1][:args.top_k]
        boxes = boxes[order]
        landmarks = landmarks[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, :]
        landmarks = landmarks[keep, :]
        print('time',time.time()-start,start1-start)
        dets = np.concatenate((dets, landmarks), axis=1)
        if args.show_image:
            for box in dets:
                if (box[3] < - box[1]) or (box[2] < - box[0]):
                    continue
                elif box[0] < 0 or box[1] < 0:

                    box[0] = max(0, box[0])
                    box[1] = max(0, box[1])

                if (box[3] + box[1] > 2 * img_raw.shape[0]) or (box[0] + box[2] > 2 * img_raw.shape[1]):
                    continue
                elif box[3] > img_raw.shape[0] or box[2] > img_raw.shape[1]:
                    box[3] = min(img_raw.shape[0], box[3])
                    box[2] = max(img_raw.shape[1], box[2])

                if (box[2] - box[0]) > 4 * (box[3] - box[1]) or (box[2] - box[0]) * 4 < (box[3] - box[1]):
                    continue
                # if box[3]*resize_level > img_raw.shape[0] + 5 or box[2]*resize_level > img_raw.shape[1] + 5:
                #     # print('out_show', img_raw.shape, int(box[2]*resize_level),int(box[3]*resize_level))
                #     continue
                text = "{:.2f}".format(box[4])
                box = list(map(int, box))
                cv2.rectangle(img_raw, (box[0] * resize_level, box[1] * resize_level),
                              (box[2] * resize_level, box[3] * resize_level), (0, 0, 255), 1)

                cv2.circle(img_raw, (box[5], box[6]), 1, (0, 0, 255), 4)
                cv2.circle(img_raw, (box[7], box[8]), 1, (0, 255, 255), 4)
                cv2.circle(img_raw, (box[9], box[10]), 1, (255, 0, 255), 4)
                cv2.circle(img_raw, (box[11], box[12]), 1, (0, 255, 0), 4)
                cv2.circle(img_raw, (box[13], box[14]), 1, (255, 0, 0), 4)
                cx = box[0] * resize_level + 18
                cy = box[1] * resize_level + 18
                # cv2.putText(img_raw, text, (cx, cy),                           cv2.FONT_HERSHEY_DUPLEX, 0.3, (0, 255, 0))

                # landms
                # cv2.circle(img_raw, (b[5], b[6]), 1, (0, 0, 255), 4)
                # cv2.circle(img_raw, (b[7], b[8]), 1, (0, 255, 255), 4)
                # cv2.circle(img_raw, (b[9], b[10]), 1, (255, 0, 255), 4)
                # cv2.circle(img_raw, (b[11], b[12]), 1, (0, 255, 0), 4)
                # cv2.circle(img_raw, (b[13], b[14]), 1, (255, 0, 0), 4)
            # save image

            if img_raw.shape[0] > 1080:
                fy = 1070 / img_raw.shape[0]
                img_raw = cv2.resize(img_raw, (0, 0), fx=fy, fy=fy, interpolation=cv2.INTER_NEAREST)
            cv2.imshow("sdf", img_raw)
            cv2.waitKey(1)

        # print(time.time()-start,"inputs name:", inname, "outputs name:", outname,prediction)
    # drawBox(boxes,img,img_shape)


if __name__ == '__main__':
    main()