RKNN3588——利用推理YOLOv8推理图片

1. yolov8_test.py

import os
import cv2
import numpy as np
from class_type import CLASSES
# 设置对象置信度阈值和非极大值抑制（NMS）阈值。
OBJ_THRESH = 0.25
NMS_THRESH = 0.45 
IMG_SIZE = (640, 640)
def filter_boxes(boxes, box_confidences, box_class_probs):
    # 筛选出满足条件的框，根据置信度和类别概率筛选出有效的框。
    box_confidences = box_confidences.reshape(-1)
    # candidate, class_num = box_class_probs.shape
    class_max_score = np.max(box_class_probs, axis=-1)
    classes = np.argmax(box_class_probs, axis=-1)
    _class_pos = np.where(class_max_score * box_confidences >= OBJ_THRESH)
    scores = (class_max_score * box_confidences)[_class_pos]
    boxes = boxes[_class_pos]
    classes = classes[_class_pos]
    return boxes, classes, scores
def nms_boxes(boxes, scores):
    # 使用非极大值抑制（NMS）来消除冗余框，保留最优的检测框。
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]
    areas = w * h
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])
        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr  None:
        self.rknn_model = model_path
        self.rknn_lite = RKNNLite()
        print(f'--> Load {self.rknn_model} model')
        ret = self.rknn_lite.load_rknn(self.rknn_model)
        if ret != 0:
            print('Load RKNNLite model failed')
            exit(ret)
        print('done')
        # 初始化运行环境
        print('--> Init runtime environment')
        ret = self.rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
        if ret != 0:
            print('Init runtime environment failed')
            exit(ret)
        print('done')
    def inference(self, img_src, IMG_SIZE):
        if img_src is None:
            print('Error: image read failed')
            return None
        self.co_helper = COCO_test_helper(enable_letter_box=True)
        img = self.co_helper.letter_box(im=img_src.copy(), new_shape=(IMG_SIZE[1], IMG_SIZE[0]), pad_color=(0, 0, 0))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.expand_dims(img, 0)
        # print(f'--> Running RKNN model')
        outputs = self.rknn_lite.inference(inputs=[img])
        return outputs
    def release(self):
        self.rknn_lite.release()
    def recover_real_box(self, boxes):
        # 还原框
        boxes = self.co_helper.get_real_box(boxes)
        return boxes
    
if __name__ == '__main__':
    yolo_model_path = 'yolov8-main/study/yolov8-240617.rknn'
    yolo_model = Model(yolo_model_path)
    img_path = r"yolov8-main/study/76_269.jpg"
    img = cv2.imread(img_path)
    yolo_result = yolo_model.inference(img, IMG_SIZE=(640,640))
    boxes, classes, scores = yolov8_post_process(yolo_result)
    boxes = yolo_model.recover_real_box(boxes=boxes)
    after_images = draw(img, boxes, scores, classes)
    cv2.imwrite("1.jpg",after_images)
    # print(yolo_result)

3. study/class_type.py

CLASSES = ("building", "building2", "statue")
coco_id_list = [1, 2, 3]