RK3588 AI 应用开发 （YOLOX-目标检测）【玩转华为云】

RK3588 AI 应用开发 （YOLOX-目标检测）

一、模型训练和转换

YOLOX是YOLO系列的优化版本，引入了解耦头、数据增强、无锚点以及标签分类等目标检测领域的优秀进展，拥有较好的精度表现，同时对工程部署友好。

模型的训练与转换教程已经开放在AI Gallery中，其中包含训练数据、训练代码、模型转换脚本。

在ModelArts的Notebook环境中训练后，再转换成对应平台的模型格式：onnx格式可以用在Windows设备上，RK系列设备上需要转换为rknn格式。

二、应用开发

1. 开发 Gradio 界面

import cv2
import json
import base64
import requests
import numpy as np
import gradio as gr
def test_image(image_path):
    try:
        image_bgr = cv2.imread(image_path)
        image_string = cv2.imencode('.jpg', image_bgr)[1].tobytes()
        image_base64 = base64.b64encode(image_string).decode('utf-8')
        params = {"image_base64": image_base64}

        response = requests.post(f'http://{ip}:{port}{url}', data=json.dumps(params),
                                headers={"Content-Type": "application/json"})
        
        if response.status_code == 200:
            image_base64 = response.json().get("image_base64")
            image_binary = base64.b64decode(image_base64)
            image_array = np.frombuffer(image_binary, dtype=np.uint8)
            image_rgb = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
        else:
            image_rgb = None

    except Exception as e:
        return None
    else:    
        return image_rgb
if __name__ == "__main__":
    port = 8000
    ip = "127.0.0.1"
    url = "/v1/fish_det"

    demo = gr.Interface(fn=test_image, inputs=gr.Image(type="filepath"), outputs=["image"], title="YOLOX 深海鱼类检测")
    demo.launch(share=False, server_port=3000)

* Running on local URL:  http://127.0.0.1:3000
* To create a public link, set `share=True` in `launch()`.

2. 编写推理代码

%%writefile YOLOX/yolox/data/datasets/voc_classes.py
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

# VOC_CLASSES = ( '__background__', # always index 0
VOC_CLASSES = (
    "fish",
)

Overwriting YOLOX/yolox/data/datasets/voc_classes.py

import sys
sys.path.append("YOLOX")

from yolox.utils import demo_postprocess, multiclass_nms, vis
from yolox.data.data_augment import preproc as preprocess
from yolox.data.datasets.voc_classes import VOC_CLASSES

import cv2
import numpy as np
import ipywidgets as widgets
from rknnlite.api import RKNNLite
from IPython.display import display

class YOLOX:
    def __init__(self, model_path):
        self.ratio = None
        self.rknn_lite = RKNNLite()
        self.rknn_lite.load_rknn(model_path)
        self.rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
    def preprocess(self, image):
        start_img, self.ratio = preprocess(image, (320, 320), swap=(0, 1, 2))
        return np.expand_dims(start_img, axis=0)
    def rknn_infer(self, data):
        outputs = self.rknn_lite.inference(inputs=[data])
        return outputs[0]
    def post_process(self, pred):
        predictions = demo_postprocess(pred.squeeze(), (320, 320))

        boxes = predictions[:, :4]
        scores = predictions[:, 4:5] * predictions[:, 5:]

        boxes_xyxy = np.ones_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
        boxes_xyxy /= self.ratio

        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.25)
        return dets
    def predict(self, image):
        # 图像预处理
        data = self.preprocess(image)
        # 模型推理
        pred = self.rknn_infer(data)
        # 模型后处理
        dets = self.post_process(pred)
        # 绘制目标检测结果
        if dets is not None:
            final_boxes = dets[:, :4]
            final_scores, final_cls_inds = dets[:, 4], dets[:, 5]
            image = vis(image, final_boxes, final_scores, final_cls_inds, conf=0.25, class_names=VOC_CLASSES)
        return image[..., ::-1]
    def img2bytes(self, image):
        """将图片转换为字节码"""
        return bytes(cv2.imencode('.jpg', image)[1])
    def infer_video(self, video_path):
        """视频推理"""
        image_widget = widgets.Image(format='jpeg', width=800, height=600)
        display(image_widget)

        cap = cv2.VideoCapture(video_path)
        while True:
            ret, img_frame = cap.read()
            if not ret:
                break
            image_pred = self.predict(img_frame)
            image_widget.value = self.img2bytes(image_pred)
        cap.release()
    def release(self):
        """释放资源"""
        self.rknn_lite.release()

3. 图像预测

4. 视频推理

5. 创建 Flask 服务

import cv2
import base64
import numpy as np
from rknnlite.api import RKNNLite
from flask import Flask, request, jsonify
from flask_cors import CORS

app = Flask(__name__)
CORS(app)

@app.route('/v1/fish_det', methods=['POST'])
def inference():
    data = request.get_json()
    image_base64 = data.get("image_base64")
    image_binary = base64.b64decode(image_base64)
    image_array = np.frombuffer(image_binary, dtype=np.uint8)
    image_bgr = cv2.imdecode(image_array, cv2.IMREAD_COLOR)

    image_rgb = model.predict(image_bgr)

    image_string = cv2.imencode('.jpg', image_rgb)[1].tobytes()
    image_base64 = base64.b64encode(image_string).decode('utf-8')

    return jsonify({
        "image_base64": image_base64
    }), 200

if __name__ == '__main__':
    model = YOLOX('model/yolox_fish.rknn')
    app.run(host='0.0.0.0', port=8000) 
    model.release()

6. 上传图片预测

三、小结

本章介绍了基于RK3588平台使用YOLOX进行目标检测的全流程，包括模型训练与转换、Gradio界面开发、推理代码编写、图像和视频预测实现，以及Flask服务部署。整体实现了高效的鱼类检测应用，适用于嵌入式设备部署与实际场景应用。