瑞芯微转化人脸检测retinaface模型_python

瑞芯微转化人脸检测retinaface模型

一、运行docker
二、转换步骤
- 1.使用https://netron.app/查看模型的输入及输出
- 2.设置转换模型参数
- 3.运行文件生成rknn模型文件
- 4.调用模型测试
- 5.运行测试
总结

一、运行docker

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb rknn-toolkit:1.7.1 /bin/bash

二、转换步骤 1.使用https://netron.app/查看模型的输入及输出

2.设置转换模型参数

test.py代码如下：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
import time

ONNX_MODEL = 'det_10g.onnx'
RKNN_MODEL = 'det_10g_500.rknn'

if __name__ == '__main__':
    # Create RKNN object
    #rknn = RKNN(verbose = True)
    rknn = RKNN()

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         #inputs=['input.1'],
                         #input_size_list=[[3,608, 608]],
                         #outputs=['277','279'])#调试查看中间输出
                         outputs=['448','471','494','451','474','497','454','477','500'])
    if ret != 0:
        print('Load failed!')
        exit(ret)
    print('done')


    # pre-process config
    print('--> Config model')   
    #rknn.config(batch_size=10,target_platform=['rv1126'],mean_values=[[127.5, 127.5, 127.5]], std_values=[[128.0,128.0, 128.0]],quantized_dtype='dynamic_fixed_point-#i8',quantized_algorithm = "mmse")
    rknn.config(batch_size=10,target_platform=['rv1126'],mean_values=[[127.5, 127.5, 127.5]], std_values=[[128.0,128.0, 128.0]])
    print('done')

    # Build model
    print('--> Building model')
    t1 =time.time()
    #ret = rknn.build(do_quantization=False,pre_compile=False) #不量化测试
    ret = rknn.build(do_quantization=True,dataset='./data.txt',pre_compile=True)
    t2 = time.time() - t1
    print("cost time = ",t2)
    
    if ret != 0:
        print('Build failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export failed!')
        exit(ret)
    print('done')
    
    #性能分析，可删除
    print('--> Accuracy analysis')
    rknn.accuracy_analysis(inputs='./test.txt', target='rv1126')
    print('done')

    rknn.release()

3.运行文件生成rknn模型文件

python test.py

4.调用模型测试

run.py代码如下：

import numpy as np
import cv2
from PIL import Image
from rknn.api import RKNN
import time

GRID0 = 19
GRID1 = 38
GRID2 = 76
LISTSIZE = 85
SPAN = 3
NUM_CLS = 80
MAX_BOXES = 500
OBJ_THRESH = 0.5
NMS_THRESH = 0.5
fmc = 3
_feat_stride_fpn = [8, 16, 32]
_num_anchors = 2
use_kps = True
center_cache = {}

def distance2bbox(points, distance, max_shape=None):
    """Decode distance prediction to bounding box.

    Args:
        points (Tensor): Shape (n, 2), [x, y].
        distance (Tensor): Distance from the given point to 4
            boundaries (left, top, right, bottom).
        max_shape (tuple): Shape of the image.

    Returns:
        Tensor: Decoded bboxes.
    """
    x1 = points[:, 0] - distance[:, 0]
    y1 = points[:, 1] - distance[:, 1]
    x2 = points[:, 0] + distance[:, 2]
    y2 = points[:, 1] + distance[:, 3]
    if max_shape is not None:
        x1 = x1.clamp(min=0, max=max_shape[1])
        y1 = y1.clamp(min=0, max=max_shape[0])
        x2 = x2.clamp(min=0, max=max_shape[1])
        y2 = y2.clamp(min=0, max=max_shape[0])
    return np.stack([x1, y1, x2, y2], axis=-1)

def distance2kps(points, distance, max_shape=None):
    """Decode distance prediction to bounding box.

    Args:
        points (Tensor): Shape (n, 2), [x, y].
        distance (Tensor): Distance from the given point to 4
            boundaries (left, top, right, bottom).
        max_shape (tuple): Shape of the image.

    Returns:
        Tensor: Decoded bboxes.
    """
    preds = []
    for i in range(0, distance.shape[1], 2):
        px = points[:, i%2] + distance[:, i]
        py = points[:, i%2+1] + distance[:, i+1]
        if max_shape is not None:
            px = px.clamp(min=0, max=max_shape[1])
            py = py.clamp(min=0, max=max_shape[0])
        preds.append(px)
        preds.append(py)
    return np.stack(preds, axis=-1)

def nms(dets):
    thresh = NMS_THRESH
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep

def load_model():
        rknn = RKNN(verbose = True)
        print('-->loading model')
        rknn.load_rknn('./det_10g_500.rknn')
        print('loading model done')

        print('--> Init runtime environment')
        ret = rknn.init_runtime(target='rv1126')
        if ret != 0:
                print('Init runtime environment failed')
                exit(ret)
        print('done')
        return rknn



if __name__ == '__main__':
    rknn = load_model()
    im = cv2.imread('./data/2.jpg')
    h, w, channels = im.shape
    #print(im)
    dim = (608, 608)
    # resize image
    det_img = cv2.resize(im, dim, interpolation = cv2.INTER_AREA)
    det_img = np.asarray(det_img)   
    
    print('--> inter')
    #print(x2)
    #增加一个维度
    det_img = det_img[:, :, :, np.newaxis]
    #转换为模型需要的输入维度(1, 3, 608, 608)
    det_img = det_img.transpose([3, 2, 0, 1])
    print(det_img)
 # Inference
    
    print('--> Running model')
    t1 = time.time()
    net_outs = rknn.inference(inputs=det_img,data_format="nchw")
    t2 = time.time() -t1
    print('--> t2 = ',t2)
    #net_outs = rknn.inference(inputs=[det_img],data_format="nchw")
    print('--> inputs')
    
    print(len(net_outs[2]))
    print(net_outs[2])
    print('outputs over!!!!!!')
    
    scores_list = []
    bboxes_list = []
    kpss_list = []
    det_scale = 1
    for idx, stride in enumerate(_feat_stride_fpn):
        scores = net_outs[idx]
        bbox_preds = net_outs[idx+fmc]

        #print("scores",scores)
        # print("bbox_preds",bbox_preds,fmc)
        # print("scores",net_outs[0],"box",net_outs[3])
        for idx_i in range(len(scores)):
            score_num = float(scores[idx_i])
            if score_num > 0.5:
                kps_preds = net_outs[idx+fmc*2] * stride
                print("score_num:",idx_i,float(score_num))
        bbox_preds = bbox_preds * stride        
        kps_preds = net_outs[idx+fmc*2] * stride
        
        height = 608// stride#input_height // stride
        width = 608// stride#input_width // stride
        K = height * width
        key = (height, width, stride)
        
        anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32)
        anchor_centers = (anchor_centers * stride).reshape( (-1, 2) )
        if _num_anchors>1:
            anchor_centers = np.stack([anchor_centers]*_num_anchors, axis=1).reshape( (-1,2) )
        if len(center_cache)<100:
            center_cache[key] = anchor_centers

        pos_inds = np.where(scores>=OBJ_THRESH)[0]       
        bboxes = distance2bbox(anchor_centers, bbox_preds)
        pos_scores = scores[pos_inds]
        pos_bboxes = bboxes[pos_inds]
        scores_list.append(pos_scores)
        bboxes_list.append(pos_bboxes)      
        kpss = distance2kps(anchor_centers, kps_preds)         
        kpss = kpss.reshape( (kpss.shape[0], -1, 2) )
        pos_kpss = kpss[pos_inds]
        kpss_list.append(pos_kpss)

    max_num = 0
    scores = np.vstack(scores_list)
    scores_ravel = scores.ravel()
    order = scores_ravel.argsort()[::-1]
    bboxes = np.vstack(bboxes_list) / det_scale
    
    kpss = np.vstack(kpss_list) / det_scale
    pre_det = np.hstack((bboxes, scores)).astype(np.float32, copy=False)
    pre_det = pre_det[order, :]
    keep = nms(pre_det)
    det = pre_det[keep, :]
    
    kpss = kpss[order,:,:]
    kpss = kpss[keep,:,:]
   
    if max_num > 0 and det.shape[0] > max_num:
        area = (det[:, 2] - det[:, 0]) * (det[:, 3] -det[:, 1])
        img_center = img.shape[0] // 2, img.shape[1] // 2
        offsets = np.vstack([
            (det[:, 0] + det[:, 2]) / 2 - img_center[1],
            (det[:, 1] + det[:, 3]) / 2 - img_center[0]
        ])
        offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
        if metric=='max':
            values = area
        else:
            values = area - offset_dist_squared * 2.0  # some extra weight on the centering
        bindex = np.argsort(
            values)[::-1]  # some extra weight on the centering
        bindex = bindex[0:max_num]
        det = det[bindex, :]
        if kpss is not None:
            kpss = kpss[bindex, :]


    print('--> Begin evaluate modelperformance')
    perf_results = rknn.eval_perf(inputs=[det_img])
    print('done')
   
    print(det)
    for face_i in range(len(det)): 
        cv2.rectangle(im,(int(det[face_i][0]*w *1.0/608),int(det[face_i][1]*h *1.0/608)),(int(det[face_i][2]*w *1.0/608),int(det[face_i][3]*h *1.0/608)),(0,0,255),3)
    cv2.imwrite("out.jpg",im)

    rknn.release()