AI挑战营（进阶）:3. onnx2rknn

waterman

AI挑战营（进阶）:3. onnx2rknn [复制链接]

在上一节中，我们获得了retinaface和facenet的onnx模型，在这一节中我们将其转换为RKNPU能够识别并运行的rknn格式。

环境安装

参考我之前写的一篇博客：https://bbs.elecfans.com/jishu_2454118_1_1.html

模型转换

retinaface

数据集准备

首先上传一张包含人脸的校准数据集，并将其路径存放至retinaface_dataset.txt文件中。

retinaface_dataset.txt内容如下：

复制

登录后复制

../../img/detection/face.jpg

模型转换

编写模型转换脚本：

复制

登录后复制

import sys

from rknn.api import RKNN

def parse_arg():
    if len(sys.argv) < 5:
        print("Usage: python3 {} [onnx_model_path] [dataset_path] [output_rknn_path] [model_type]".format(sys.argv[0]));
        exit(1)

    model_path = sys.argv[1]
    dataset_path= sys.argv[2]
    output_path = sys.argv[3]
    model_type = sys.argv[4]

    return model_path, dataset_path, output_path,model_type

if __name__ == '__main__':
    model_path, dataset_path, output_path, model_type= parse_arg()

    # Create RKNN object
    rknn = RKNN(verbose=False)

    # Pre-process config
    print('--> Config model')
    if model_type == 'Retinaface':
        rknn.config(mean_values=[[104, 117, 123]], std_values=[[1, 1, 1]], target_platform='rv1103',
                                quantized_algorithm="normal", quant_img_RGB2BGR=True,optimization_level=0)
        print("Use retinaface mode")
    else:
        rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rv1103')

    print('done')

    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=model_path)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=True, dataset=dataset_path)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn(output_path)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    # Release
    rknn.release()

执行convert.py进行模型转换：

复制

登录后复制

python convert.py ../model_zoo/export/retinaface.onnx ./dataset/retinaface_dataset.txt ../model_zoo/retinaface.rknn Retinaface

facenet

数据集准备

首先上传一张包含人脸的校准数据集，并将其路径存放至facenet_dataset.txt文件中。

facenet_dataset.txt内容如下：

复制

登录后复制

../../img/recongnition/calib.jpg

模型转换

执行convert.py进行模型转换：

复制

登录后复制

python convert.py ../model_zoo/export/facenet.onnx ./dataset/facenet_dataset.txt ../model_zoo/facenet.rknn Facenet

转换完成后，就能够分别得到两个相应的rknn模型文件于model_zoo目录下：

模型评估

数据集准备

首先准备验证数据集：

目录结构如下：

复制

登录后复制

eval
├── dataset
│   ├── facenet
│   │   ├── 1_001.jpg
│   │   ├── 1_002.jpg
│   │   └── 2_001.jpg
│   └── retinaface
│       └── j1.jpg
├── facenet.py
└── retinaface.py

retinaface

复制

登录后复制

import os
import sys
import urllib
import urllib.request
import time
import numpy as np
import cv2
from math import ceil
from itertools import product as product

from rknn.api import RKNN
DATASET_PATH = '../dataset/retinaface_dataset.txt'
DEFAULT_QUANT = True

def letterbox_resize(image, size, bg_color):
    """
    letterbox_resize the image according to the specified size
    :param image: input image, which can be a NumPy array or file path
    :param size: target size (width, height)
    :param bg_color: background filling data 
    :return: processed image
    """
    if isinstance(image, str):
        image = cv2.imread(image)

    target_width, target_height = size
    image_height, image_width, _ = image.shape

    # 计算调整后的图像尺寸
    aspect_ratio = min(target_width / image_width, target_height / image_height)
    new_width = int(image_width * aspect_ratio)
    new_height = int(image_height * aspect_ratio)

    # 使用 cv2.resize() 进行等比缩放
    image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)

    # 创建新的画布并进行填充
    result_image = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
    offset_x = (target_width - new_width) // 2
    offset_y = (target_height - new_height) // 2
    result_image[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image
    return result_image, aspect_ratio, offset_x, offset_y

def PriorBox(image_size): #image_size Support (320,320) and (640,640)
    anchors = []
    min_sizes = [[16, 32], [64, 128], [256, 512]]
    steps = [8, 16, 32]
    feature_maps = [[ceil(image_size[0] / step), ceil(image_size[1] / step)] for step in steps]
    for k, f in enumerate(feature_maps):
        min_sizes_ = min_sizes[k]
        for i, j in product(range(f[0]), range(f[1])):
            for min_size in min_sizes_:
                s_kx = min_size / image_size[1]
                s_ky = min_size / image_size[0]
                dense_cx = [x * steps[k] / image_size[1] for x in [j + 0.5]]
                dense_cy = [y * steps[k] / 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]
    output = np.array(anchors).reshape(-1, 4)
    print("image_size:",image_size," num_priors=",output.shape[0])
    return output


def box_decode(loc, priors):
    """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
    """
    variances = [0.1, 0.2]
    boxes = np.concatenate((
        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
        priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), axis=1)
    boxes[:, :2] -= boxes[:, 2:] / 2
    boxes[:, 2:] += boxes[:, :2]
    return boxes


def decode_landm(pre, priors):
    """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
    """
    variances = [0.1, 0.2]
    landmarks = 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:]
    ), axis=1)
    return landmarks


def nms(dets, thresh):
    """Pure Python NMS baseline."""
    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


if __name__ == '__main__':
    # 创建RKNN对象 
    rknn = RKNN()

    # 预处理设置
    print('--> Config model')
    rknn.config(mean_values=[[104, 117, 123]], std_values=[[1, 1, 1]], target_platform="rv1106",
                quantized_algorithm="normal", quant_img_RGB2BGR=True)  # mmse
    print('done')

    # 载入模型
    print('--> Loading model')
    ret = rknn.load_onnx(model="../../model_zoo/export/retinaface.onnx")
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # 创建模型
    print('--> Building model')
    ret = rknn.build(do_quantization=True, dataset=DATASET_PATH)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')


    # 输入图像
    img = cv2.imread('./dataset/retinaface/j1.jpg')
    img_height, img_width, _ = img.shape
    model_height, model_width = (640, 640)
    letterbox_img, aspect_ratio, offset_x, offset_y = letterbox_resize(img, (model_height,model_width), 114)  # letterbox缩放
    infer_img = letterbox_img[..., ::-1]  # BGR2RGB

    # 初始化运行时环境
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

    # 运行
    print('--> Running model')
    outputs = rknn.inference(inputs=[infer_img])
    loc, conf, landmarks = outputs #获取输出数据

    priors = PriorBox(image_size=(model_height, model_width)) # 获取先验框
    boxes = box_decode(loc.squeeze(0), priors)                # 解码输出数据

    # letterbox
    scale = np.array([model_width, model_height,
                      model_width, model_height])
    boxes = boxes * scale // 1  # face box
    boxes[...,0::2] =np.clip((boxes[...,0::2] - offset_x) / aspect_ratio, 0, img_width)  #letterbox
    boxes[...,1::2] =np.clip((boxes[...,1::2] - offset_y) / aspect_ratio, 0, img_height) #letterbox
    scores = conf.squeeze(0)[:, 1]  # 人脸检测的置信度
    landmarks = decode_landm(landmarks.squeeze(
        0), priors)  # face keypoint data
    scale_landmarks = np.array([model_width, model_height, model_width, model_height,
                                model_width, model_height, model_width, model_height,
                                model_width, model_height])
    landmarks = landmarks * scale_landmarks // 1
    landmarks[...,0::2] = np.clip((landmarks[...,0::2] - offset_x) / aspect_ratio, 0, img_width) #letterbox
    landmarks[...,1::2] = np.clip((landmarks[...,1::2] - offset_y) / aspect_ratio, 0, img_height) #letterbox

    # 丢弃置信度过低的部分
    inds = np.where(scores > 0.5)[0]
    boxes = boxes[inds]
    landmarks = landmarks[inds]
    scores = scores[inds]

    order = scores.argsort()[::-1]
    boxes = boxes[order]
    landmarks = landmarks[order]
    scores = scores[order]

    # 非极大值抑制
    dets = np.hstack((boxes, scores[:, np.newaxis])).astype(
        np.float32, copy=False)
    keep = nms(dets, 0.2)
    dets = dets[keep, :]
    landmarks = landmarks[keep]
    dets = np.concatenate((dets, landmarks), axis=1)

    # 画框标记
    for data in dets:
        if data[4] < 0.5:
            continue
        print("face @ (%d %d %d %d) %f"%(data[0], data[1], data[2], data[3], data[4]))
        text = "{:.4f}".format(data[4])
        data = list(map(int, data))
        cv2.rectangle(img, (data[0], data[1]),
                      (data[2], data[3]), (0, 0, 255), 2)
        cx = data[0]
        cy = data[1] + 12
        cv2.putText(img, text, (cx, cy),
                    cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
        # landmarks
        cv2.circle(img, (data[5], data[6]), 1, (0, 0, 255), 5)
        cv2.circle(img, (data[7], data[8]), 1, (0, 255, 255), 5)
        cv2.circle(img, (data[9], data[10]), 1, (255, 0, 255), 5)
        cv2.circle(img, (data[11], data[12]), 1, (0, 255, 0), 5)
        cv2.circle(img, (data[13], data[14]), 1, (255, 0, 0), 5)
    img_path = './result.jpg'
    cv2.imwrite(img_path, img)
    print("save image in", img_path)
    # 释放
    rknn.release()

执行结果如下：

facenet

复制

登录后复制

import numpy as np
import cv2
import os
from rknn.api import RKNN
from PIL import Image
from sklearn import preprocessing
from scipy.spatial.distance import pdist
os.environ['RKNN_DRAW_DATA_DISTRIBUTE']="1"

if __name__ == '__main__':
    BUILD_QUANT = True

    # Create RKNN object
    rknn = RKNN()
    print('--> config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rv1103')
    print('done')

    # Load model
    print('--> Lccoading model')
    ret = rknn.load_onnx(model='../../model_zoo/export/facenet.onnx')
    if ret != 0:
        print('Load facenet failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=True, dataset='../dataset/facenet_dataset.txt')
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

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

    # Set inputs and inference
    image_1 = Image.open("./dataset/facenet/1_001.jpg")
    image_1 = image_1.resize((160,160), Image.BICUBIC)
    img1 = np.asarray(image_1, np.uint8)
    outputs1 = np.array(rknn.inference(data_format='nhwc', inputs=[img1])[0])
    outputs1 = preprocessing.normalize(outputs1, norm='l2') 

    image_2 = Image.open("./dataset/facenet/1_002.jpg")
    image_2 = image_2.resize((160,160), Image.BICUBIC)
    img2 = np.asarray(image_2, np.uint8)
    outputs2 = np.array(rknn.inference(data_format='nhwc', inputs=[img2])[0])
    outputs2 = preprocessing.normalize(outputs2, norm='l2') 

    # Get distance 
    distance = np.linalg.norm(outputs2 - outputs1, axis=1)
    print("distance:", distance)


    rknn.release()

执行结果如下：

工程文件在附录中。

241205ZHUCE

这群人真酷

 

waterman

241205ZHUCE 发表于 2024-12-5 21:58 这群人真酷  

必须的

sgf201

大佬大佬