瑞芯微:基于RK3568的Yolo-NAS部署

You Only Look Once  神经架构搜索(YOLO-NAS)是最新最先进的(SOTA)实时目标检测模型。 在 COCO 数据集上进行评估并与其前身 YOLOv6 和 YOLOv8  相比,YOLO-NAS 以更低的延迟实现了更高的 mAP 值。

YOLO-NAS 作为 Deci 维护的 super-gradient包的一部分提供。

下图展示了Deci在YOLO-NAS上的基准测试结果:

瑞芯微:基于RK3568的Yolo-NAS部署_第1张图片

瑞芯微:基于RK3568的Yolo-NAS部署_第2张图片

置信度很高呀。。。接下来我们讲yolo-nas部署到rk中去玩玩。。。

 
import cv2
import numpy as np
 
from rknn.api import RKNN
import os
 
if __name__ == '__main__':
 
    platform = 'rk3568'
    exp = 'yolo-nas-s'
    Width = 640
    Height = 640
    MODEL_PATH = './onnx_models/yolo-nas-s.onnx'
    NEED_BUILD_MODEL = True
    # NEED_BUILD_MODEL = False
    im_file = './bus.jpg'
 
    # Create RKNN object
    rknn = RKNN()
 
    OUT_DIR = "rknn_models"
    RKNN_MODEL_PATH = './{}/{}_kk.rknn'.format(OUT_DIR,exp+'-'+str(Width)+'-'+str(Height))
    if NEED_BUILD_MODEL:
        DATASET = './pose.txt'
        rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform="rk3568")
        # Load model
        print('--> Loading model')
        ret = rknn.load_onnx(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)
        if ret != 0:
            print('build model failed.')
            exit(ret)
        print('done')
 
        # Export rknn model
        if not os.path.exists(OUT_DIR):
            os.mkdir(OUT_DIR)
        print('--> Export RKNN model: {}'.format(RKNN_MODEL_PATH))
        ret = rknn.export_rknn(RKNN_MODEL_PATH)
        if ret != 0:
            print('Export rknn model failed.')
            exit(ret)
        print('done')
    else:
        ret = rknn.load_rknn(RKNN_MODEL_PATH)
 
    rknn.release()

结果:

瑞芯微:基于RK3568的Yolo-NAS部署_第3张图片

瑞芯微:基于RK3568的Yolo-NAS部署_第4张图片

你可能感兴趣的:(AI-训练+部署,YOLO,人工智能,rknn)