├── README.md ├── yolov8obb_horizon └── mapper │ ├── 01_check.sh │ ├── 02_preprocess.sh │ ├── 03_build.sh │ ├── 04_inference.sh │ ├── 05_evaluate.sh │ ├── cal_data │ └── test.jpg │ ├── data_preprocess.py │ ├── inference_image_demo.py │ ├── model │ └── yolov8n-obb.onnx │ ├── model_output │ ├── yolov8obb.bin │ └── yolov8obb_quantized_model.onnx │ ├── preprocess.py │ ├── src_data │ └── test.jpg │ ├── test.jpg │ ├── test_onnx_result.jpg │ └── yolov8obb_config.yaml ├── yolov8obb_onnx ├── test.jpg ├── test_onnx_result.jpg ├── yolov8n-obb.onnx └── yolov8obb_onnx_demo.py ├── yolov8obb_rknn ├── images_list.txt ├── images_list │ └── test.jpg ├── onnx2rknn_demo_ZQ.py ├── test.jpg ├── test_onnx_result.jpg ├── yolov8n-obb.onnx └── yyolov8n-obb.rknn └── yolov8obb_tensorRT ├── onnx2trt_rt7.py ├── tensorRT_inferenc_demo.py ├── test.jpg ├── test_onnx_result.jpg ├── yolov8n-obb.onnx └── yolov8n-obb.trt /README.md: -------------------------------------------------------------------------------- 1 | # yolov8-obb_onnx_rknn_horizon_tensonRT 2 | 3 | yolov8-obb 旋转目标检测部署，瑞芯微RKNN芯片部署、地平线Horizon芯片部署、TensorRT部署 4 | 5 | 导出 onnx 参考链接 [【yolov8-obb 旋转目标检测 瑞芯微RKNN芯片部署、地平线Horizon芯片部署、TensorRT部署】](https://blog.csdn.net/zhangqian_1/article/details/139437315) 6 | 7 | **特别说明：本示例提供的代码只适用按照对应参考导出的onnx，其它方式导出onnx自行写后处理。** 8 | 9 | # 文件夹结构说明 10 | 11 | yolov8obb_onnx：onnx模型、测试图像、测试结果、测试demo脚本 12 | 13 | yolov8obb_TensorRT：TensorRT版本模型、测试图像、测试结果、测试demo脚本、onnx模型、onnx2tensorRT脚本(tensorRT-7.2.3.4) 14 | 15 | yolov8obb_rknn：rknn模型、测试（量化）图像、测试结果、onnx2rknn转换测试脚本 16 | 17 | yolov8obb_horizon：地平线模型、测试（量化）图像、测试结果、转换测试脚本、测试量化后onnx模型脚本 18 | 19 | # 测试结果 20 | 21 | python 测试结果 22 | 23 | ![test](https://github.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/assets/22290931/fc44788c-5736-4ab5-baf8-6fd04d906b82) 24 | 25 | 26 | onnx 测试结果 27 | 28 | ![image](https://github.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/blob/main/yolov8obb_onnx/test_onnx_result.jpg) 29 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/01_check.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env sh 2 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 3 | # 4 | # The material in this file is confidential and contains trade secrets 5 | # of Horizon Robotics Inc. This is proprietary information owned by 6 | # Horizon Robotics Inc. No part of this work may be disclosed, 7 | # reproduced, copied, transmitted, or used in any way for any purpose, 8 | # without the express written permission of Horizon Robotics Inc. 9 | 10 | set -e -v 11 | cd $(dirname $0) || exit 12 | 13 | model_type="onnx" 14 | onnx_model="./model/yolov8n-obb.onnx" 15 | output="./yolov8obb_checker.log" 16 | march="bernoulli2" 17 | 18 | hb_mapper checker --model-type ${model_type} \ 19 | --model ${onnx_model} \ 20 | --output ${output} --march ${march} 21 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/02_preprocess.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 3 | # 4 | # The material in this file is confidential and contains trade secrets 5 | # of Horizon Robotics Inc. This is proprietary information owned by 6 | # Horizon Robotics Inc. No part of this work may be disclosed, 7 | # reproduced, copied, transmitted, or used in any way for any purpose, 8 | # without the express written permission of Horizon Robotics Inc. 9 | 10 | set -e -v 11 | cd $(dirname $0) || exit 12 | 13 | python3 data_preprocess.py \ 14 | --src_dir ./src_data \ 15 | --dst_dir ./cal_data \ 16 | --pic_ext .rgb \ 17 | --read_mode opencv 18 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/03_build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 3 | # 4 | # The material in this file is confidential and contains trade secrets 5 | # of Horizon Robotics Inc. This is proprietary information owned by 6 | # Horizon Robotics Inc. No part of this work may be disclosed, 7 | # reproduced, copied, transmitted, or used in any way for any purpose, 8 | # without the express written permission of Horizon Robotics Inc. 9 | 10 | set -e -v 11 | cd $(dirname $0) 12 | config_file="./yolov8obb_config.yaml" 13 | model_type="onnx" 14 | # build model 15 | hb_mapper makertbin --config ${config_file} \ 16 | --model-type ${model_type} 17 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/04_inference.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 3 | # 4 | # The material in this file is confidential and contains trade secrets 5 | # of Horizon Robotics Inc. This is proprietary information owned by 6 | # Horizon Robotics Inc. No part of this work may be disclosed, 7 | # reproduced, copied, transmitted, or used in any way for any purpose, 8 | # without the express written permission of Horizon Robotics Inc. 9 | 10 | python3 -u inference_image_demo.py 11 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/05_evaluate.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 3 | # 4 | # The material in this file is confidential and contains trade secrets 5 | # of Horizon Robotics Inc. This is proprietary information owned by 6 | # Horizon Robotics Inc. No part of this work may be disclosed, 7 | # reproduced, copied, transmitted, or used in any way for any purpose, 8 | # without the express written permission of Horizon Robotics Inc. 9 | 10 | set -v -e 11 | cd $(dirname $0) || exit 12 | 13 | #for converted quanti model evaluation 14 | quanti_model_file="./model_output/yolov5_672x672_nv12_quantized_model.onnx" 15 | quanti_input_layout="NHWC" 16 | 17 | original_model_file="./model_output/yolov5_672x672_nv12_original_float_model.onnx" 18 | original_input_layout="NCHW" 19 | 20 | if [[ $1 =~ "origin" ]]; then 21 | model=$original_model_file 22 | layout=$original_input_layout 23 | input_offset=128 24 | else 25 | model=$quanti_model_file 26 | layout=$quanti_input_layout 27 | input_offset=128 28 | fi 29 | 30 | image_path="../../../01_common/data/coco/coco_val2017/images/" 31 | anno_path="../../../01_common/data/coco/coco_val2017/annotations/instances_val2017.json" 32 | 33 | 34 | if [ -z $2 ]; then 35 | total_image_number=5000 36 | else 37 | total_image_number=$2 38 | fi 39 | 40 | # ------------------------------------------------------------------------------------------------------------- 41 | # shell command "sh 05_evaluate.sh" runs quanti full evaluation by default 42 | # If quanti model eval is intended, please run the shell via command "sh 05_evaluate.sh quanti" 43 | # If float model eval is intended, please run the shell via command "sh 05_evaluate.sh origin"# 44 | # If quanti model quick eval test is intended, please run the shell via command "sh 05_evaluate.sh quanti 20" 45 | # If float model quick eval test is intended, please run the shell via command "sh 05_evaluate.sh origin 20" 46 | # ------------------------------------------------------------------------------------------------------------- 47 | # quanti model eval 48 | python3 -u ../../det_evaluate.py \ 49 | --model=${model} \ 50 | --image_path=${image_path} \ 51 | --annotation_path=${anno_path} \ 52 | --input_layout=${layout} \ 53 | --total_image_number=${total_image_number} \ 54 | --input_offset ${input_offset} 55 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/cal_data/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/cal_data/test.jpg -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/data_preprocess.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 2 | # 3 | # The material in this file is confidential and contains trade secrets 4 | # of Horizon Robotics Inc. This is proprietary information owned by 5 | # Horizon Robotics Inc. No part of this work may be disclosed, 6 | # reproduced, copied, transmitted, or used in any way for any purpose, 7 | # without the express written permission of Horizon Robotics Inc. 8 | 9 | import os 10 | import sys 11 | sys.path.append('.') 12 | 13 | import click 14 | import numpy as np 15 | from preprocess import calibration_transformers 16 | import skimage.io 17 | import cv2 18 | 19 | transformers = calibration_transformers() 20 | 21 | sys.path.append("../../../01_common/python/data/") 22 | from dataloader import DataLoader 23 | from dataset import CifarDataset 24 | 25 | regular_process_list = [ 26 | ".rgb", 27 | ".rgbp", 28 | ".bgr", 29 | ".bgrp", 30 | ".yuv", 31 | ".feature", 32 | ".cali", 33 | ] 34 | 35 | 36 | def read_image(src_file, read_mode): 37 | if read_mode == "skimage": 38 | image = skimage.img_as_float(skimage.io.imread(src_file)).astype( 39 | np.float32) 40 | elif read_mode == "opencv": 41 | image = cv2.imread(src_file) 42 | else: 43 | raise ValueError(f"Invalid read mode {read_mode}") 44 | if image.ndim != 3: # expend gray scale image to three channels 45 | image = image[..., np.newaxis] 46 | image = np.concatenate([image, image, image], axis=-1) 47 | return image 48 | 49 | 50 | def regular_preprocess(src_file, transformers, dst_dir, pic_ext, read_mode): 51 | image = [read_image(src_file, read_mode)] 52 | for trans in transformers: 53 | image = trans(image) 54 | 55 | filename = os.path.basename(src_file) 56 | short_name, ext = os.path.splitext(filename) 57 | pic_name = os.path.join(dst_dir, short_name + pic_ext) 58 | print("write:%s" % pic_name) 59 | dtype = np.float32 if dst_dir.endswith("_f32") else np.uint8 60 | image[0].astype(dtype).tofile(pic_name) 61 | 62 | 63 | def cifar_preprocess(src_file, data_loader, dst_dir, pic_ext, cal_img_num): 64 | for i in range(cal_img_num): 65 | image, label = next(data_loader) 66 | filename = os.path.basename(src_file) 67 | pic_name = os.path.join(dst_dir + '/' + str(i) + pic_ext) 68 | print("write:%s" % pic_name) 69 | image[0].astype(np.uint8).tofile(pic_name) 70 | 71 | 72 | @click.command(help=''' 73 | A Tool used to generate preprocess pics for calibration. 74 | ''') 75 | @click.option('--src_dir', type=str, help='calibration source file') 76 | @click.option('--dst_dir', type=str, help='generated calibration file') 77 | @click.option('--pic_ext', 78 | type=str, 79 | default=".cali", 80 | help='picture extension.') 81 | @click.option('--read_mode', 82 | type=click.Choice(["skimage", "opencv"]), 83 | default="opencv", 84 | help='picture extension.') 85 | @click.option('--cal_img_num', type=int, default=100, help='cali picture num.') 86 | def main(src_dir, dst_dir, pic_ext, read_mode, cal_img_num): 87 | '''A Tool used to generate preprocess pics for calibration.''' 88 | pic_num = 0 89 | os.makedirs(dst_dir, exist_ok=True) 90 | if pic_ext.strip().split('_')[0] in regular_process_list: 91 | print("regular preprocess") 92 | for src_name in sorted(os.listdir(src_dir)): 93 | pic_num += 1 94 | if pic_num > cal_img_num: 95 | break 96 | src_file = os.path.join(src_dir, src_name) 97 | regular_preprocess(src_file, transformers, dst_dir, pic_ext, 98 | read_mode) 99 | elif pic_ext.strip().split('_')[0] == ".cifar": 100 | print("cifar preprocess") 101 | data_loader = DataLoader(CifarDataset(src_dir), transformers, 1) 102 | cifar_preprocess(src_dir, data_loader, dst_dir, pic_ext, cal_img_num) 103 | else: 104 | raise ValueError(f"invalid pic_ext {pic_ext}") 105 | 106 | 107 | if __name__ == '__main__': 108 | main() 109 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/inference_image_demo.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from horizon_tc_ui import HB_ONNXRuntime 3 | from horizon_tc_ui.utils.tool_utils import init_root_logger 4 | import math 5 | import cv2 6 | import numpy as np 7 | 8 | 9 | CLASSES = ['plane', 'ship', 'storage tank', 'baseball diamond', 'tennis court', 'basketball court', 10 | 'ground track field', 'harbor', 'bridge', 'large vehicle', 'small vehicle', 'helicopter', 'roundabout', 11 | 'soccer ball field', 'swimming pool'] 12 | 13 | meshgrid = [] 14 | 15 | class_num = len(CLASSES) 16 | head_num = 3 17 | strides = [8, 16, 32] 18 | map_size = [[80, 80], [40, 40], [20, 20]] 19 | reg_num = 16 20 | nms_thresh = 0.5 21 | object_thresh = 0.25 22 | 23 | input_height = 640 24 | input_width = 640 25 | 26 | 27 | 28 | class CSXYWHR: 29 | def __init__(self, classId, score, x, y, w, h, angle): 30 | self.classId = classId 31 | self.score = score 32 | self.x = x 33 | self.y = y 34 | self.w = w 35 | self.h = h 36 | self.angle = angle 37 | 38 | 39 | class DetectBox: 40 | def __init__(self, classId, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle): 41 | self.classId = classId 42 | self.score = score 43 | self.pt1x = pt1x 44 | self.pt1y = pt1y 45 | self.pt2x = pt2x 46 | self.pt2y = pt2y 47 | self.pt3x = pt3x 48 | self.pt3y = pt3y 49 | self.pt4x = pt4x 50 | self.pt4y = pt4y 51 | self.angle = angle 52 | 53 | 54 | def GenerateMeshgrid(): 55 | for index in range(head_num): 56 | for i in range(map_size[index][0]): 57 | for j in range(map_size[index][1]): 58 | meshgrid.append(j + 0.5) 59 | meshgrid.append(i + 0.5) 60 | 61 | 62 | def get_covariance_matrix(boxes): 63 | a, b, c = boxes.w, boxes.h, boxes.angle 64 | cos = math.cos(c) 65 | sin = math.sin(c) 66 | cos2 = math.pow(cos, 2) 67 | sin2 = math.pow(sin, 2) 68 | return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin 69 | 70 | 71 | def probiou(obb1, obb2, eps=1e-7): 72 | x1, y1 = obb1.x, obb1.y 73 | x2, y2 = obb2.x, obb2.y 74 | a1, b1, c1 = get_covariance_matrix(obb1) 75 | a2, b2, c2 = get_covariance_matrix(obb2) 76 | 77 | t1 = (((a1 + a2) * math.pow((y1 - y2), 2) + (b1 + b2) * math.pow((x1 - x2), 2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.25 78 | t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.5 79 | 80 | temp1 = (a1 * b1 - math.pow(c1, 2)) if (a1 * b1 - math.pow(c1, 2)) > 0 else 0 81 | temp2 = (a2 * b2 - math.pow(c2, 2)) if (a2 * b2 - math.pow(c2, 2)) > 0 else 0 82 | t3 = math.log((((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2)) / (4 * math.sqrt((temp1 * temp2)) + eps)+ eps)) * 0.5 83 | 84 | if (t1 + t2 + t3) > 100: 85 | bd = 100 86 | elif (t1 + t2 + t3) < eps: 87 | bd = eps 88 | else: 89 | bd = t1 + t2 + t3 90 | hd = math.sqrt((1.0 - math.exp(-bd) + eps)) 91 | return 1 - hd 92 | 93 | 94 | def nms_rotated(boxes, nms_thresh): 95 | pred_boxes = [] 96 | sort_boxes = sorted(boxes, key=lambda x: x.score, reverse=True) 97 | for i in range(len(sort_boxes)): 98 | if sort_boxes[i].classId != -1: 99 | pred_boxes.append(sort_boxes[i]) 100 | for j in range(i + 1, len(sort_boxes), 1): 101 | ious = probiou(sort_boxes[i], sort_boxes[j]) 102 | if ious > nms_thresh: 103 | sort_boxes[j].classId = -1 104 | return pred_boxes 105 | 106 | 107 | def sigmoid(x): 108 | return 1 / (1 + math.exp(-x)) 109 | 110 | 111 | def xywhr2xyxyxyxy(x, y, w, h, angle): 112 | cos_value = math.cos(angle) 113 | sin_value = math.sin(angle) 114 | 115 | vec1x= w / 2 * cos_value 116 | vec1y = w / 2 * sin_value 117 | vec2x = -h / 2 * sin_value 118 | vec2y = h / 2 * cos_value 119 | 120 | pt1x = x + vec1x + vec2x 121 | pt1y = y + vec1y + vec2y 122 | 123 | pt2x = x + vec1x - vec2x 124 | pt2y = y + vec1y - vec2y 125 | 126 | pt3x = x - vec1x - vec2x 127 | pt3y = y - vec1y - vec2y 128 | 129 | pt4x = x - vec1x + vec2x 130 | pt4y = y - vec1y + vec2y 131 | return pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y 132 | 133 | 134 | def postprocess(out): 135 | print('postprocess ... ') 136 | 137 | detect_result = [] 138 | output = [] 139 | for i in range(len(out)): 140 | output.append(out[i].reshape((-1))) 141 | 142 | gridIndex = -2 143 | cls_index = 0 144 | cls_max = 0 145 | 146 | for index in range(head_num): 147 | reg = output[index * 2 + 0] 148 | cls = output[index * 2 + 1] 149 | ang = output[head_num * 2 + index] 150 | 151 | for h in range(map_size[index][0]): 152 | for w in range(map_size[index][1]): 153 | gridIndex += 2 154 | 155 | if 1 == class_num: 156 | cls_max = sigmoid(cls[0 * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 157 | cls_index = 0 158 | else: 159 | for cl in range(class_num): 160 | cls_val = cls[cl * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] 161 | if 0 == cl: 162 | cls_max = cls_val 163 | cls_index = cl 164 | else: 165 | if cls_val > cls_max: 166 | cls_max = cls_val 167 | cls_index = cl 168 | cls_max = sigmoid(cls_max) 169 | 170 | if cls_max > object_thresh: 171 | regdfl = [] 172 | for lc in range(4): 173 | sfsum = 0 174 | locval = 0 175 | for df in range(reg_num): 176 | temp = math.exp(reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 177 | reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] = temp 178 | sfsum += temp 179 | 180 | for df in range(reg_num): 181 | sfval = reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] / sfsum 182 | locval += sfval * df 183 | regdfl.append(locval) 184 | 185 | angle = (sigmoid(ang[h * map_size[index][1] + w]) - 0.25) * math.pi 186 | 187 | left, top, right, bottom = regdfl[0], regdfl[1], regdfl[2], regdfl[3] 188 | cos, sin = math.cos(angle), math.sin(angle) 189 | fx = (right - left) / 2 190 | fy = (bottom - top) / 2 191 | 192 | cx = ((fx * cos - fy * sin) + meshgrid[gridIndex + 0]) * strides[index] 193 | cy = ((fx * sin + fy * cos) + meshgrid[gridIndex + 1])* strides[index] 194 | cw = (left + right) * strides[index] 195 | ch = (top + bottom) * strides[index] 196 | 197 | box = CSXYWHR(cls_index, cls_max, cx, cy, cw, ch, angle) 198 | 199 | detect_result.append(box) 200 | # NMS 201 | print('before nms num is:', len(detect_result)) 202 | pred_boxes = nms_rotated(detect_result, nms_thresh) 203 | 204 | print('after nms num is:', len(pred_boxes)) 205 | 206 | resutl = [] 207 | for i in range(len(pred_boxes)): 208 | classid = pred_boxes[i].classId 209 | score = pred_boxes[i].score 210 | cx = pred_boxes[i].x 211 | cy = pred_boxes[i].y 212 | cw = pred_boxes[i].w 213 | ch = pred_boxes[i].h 214 | angle = pred_boxes[i].angle 215 | 216 | bw_ = cw if cw > ch else ch 217 | bh_ = ch if cw > ch else cw 218 | bt = angle % math.pi if cw > ch else (angle + math.pi / 2) % math.pi 219 | 220 | pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y = xywhr2xyxyxyxy(cx, cy, bw_, bh_, bt) 221 | 222 | bbox = DetectBox(classid, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle) 223 | resutl.append(bbox) 224 | 225 | return resutl 226 | 227 | 228 | def preprocess(src_image, input_width, input_height): 229 | src_image = cv2.cvtColor(src_image, cv2.COLOR_BGR2RGB) 230 | img = cv2.resize(src_image, (input_height, input_width)) 231 | return img 232 | 233 | 234 | def inference(model_path, image_path, input_layout, input_offset): 235 | # init_root_logger("inference.log", console_level=logging.INFO, file_level=logging.DEBUG) 236 | 237 | sess = HB_ONNXRuntime(model_file=model_path) 238 | sess.set_dim_param(0, 0, '?') 239 | 240 | if input_layout is None: 241 | logging.warning(f"input_layout not provided. Using {sess.layout[0]}") 242 | input_layout = sess.layout[0] 243 | 244 | origin_image = cv2.imread(image_path) 245 | image_h, image_w = origin_image.shape[:2] 246 | image = preprocess(origin_image, input_width, input_height) 247 | image_data = np.expand_dims(image, axis=0) 248 | 249 | input_name = sess.input_names[0] 250 | output_name = sess.output_names 251 | output = sess.run(output_name, {input_name: image_data}, input_offset=input_offset) 252 | 253 | print('inference finished, output len is:', len(output)) 254 | 255 | out = [] 256 | for i in range(len(output)): 257 | out.append(output[i]) 258 | 259 | pred_boxes = postprocess(out) 260 | 261 | print('obj num is :', len(pred_boxes)) 262 | 263 | for i in range(len(pred_boxes)): 264 | classId = pred_boxes[i].classId 265 | score = pred_boxes[i].score 266 | pt1x = int(pred_boxes[i].pt1x / input_width * image_w) 267 | pt1y = int(pred_boxes[i].pt1y / input_width * image_h) 268 | pt2x = int(pred_boxes[i].pt2x / input_width * image_w) 269 | pt2y = int(pred_boxes[i].pt2y / input_width * image_h) 270 | pt3x = int(pred_boxes[i].pt3x / input_width * image_w) 271 | pt3y = int(pred_boxes[i].pt3y / input_width * image_h) 272 | pt4x = int(pred_boxes[i].pt4x / input_width * image_w) 273 | pt4y = int(pred_boxes[i].pt4y / input_width * image_h) 274 | angle = pred_boxes[i].angle 275 | 276 | cv2.line(origin_image, (pt1x, pt1y), (pt2x, pt2y), (0, 255, 0), 2) 277 | cv2.line(origin_image, (pt2x, pt2y), (pt3x, pt3y), (0, 255, 0), 2) 278 | cv2.line(origin_image, (pt3x, pt3y), (pt4x, pt4y), (0, 255, 0), 2) 279 | cv2.line(origin_image, (pt4x, pt4y), (pt1x, pt1y), (0, 255, 0), 2) 280 | 281 | title = CLASSES[classId] + "%.2f" % score 282 | cv2.putText(origin_image, title, (pt1x, pt1y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2, cv2.LINE_AA) 283 | 284 | cv2.imwrite('./test_onnx_result.jpg', origin_image) 285 | 286 | 287 | if __name__ == '__main__': 288 | print('This main ... ') 289 | GenerateMeshgrid() 290 | 291 | model_path = './model_output/yolov8obb_quantized_model.onnx' 292 | image_path = './test.jpg' 293 | input_layout = 'NHWC' 294 | input_offset = 128 295 | 296 | inference(model_path, image_path, input_layout, input_offset) 297 | 298 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/model/yolov8n-obb.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/model/yolov8n-obb.onnx -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/model_output/yolov8obb.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/model_output/yolov8obb.bin -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/model_output/yolov8obb_quantized_model.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/model_output/yolov8obb_quantized_model.onnx -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/preprocess.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021 Horizon Robotics.All Rights Reserved. 2 | # 3 | # The material in this file is confidential and contains trade secrets 4 | # of Horizon Robotics Inc. This is proprietary information owned by 5 | # Horizon Robotics Inc. No part of this work may be disclosed, 6 | # reproduced, copied, transmitted, or used in any way for any purpose, 7 | # without the express written permission of Horizon Robotics Inc. 8 | 9 | import sys 10 | sys.path.append("../../../01_common/python/data/") 11 | from transformer import * 12 | from dataloader import * 13 | 14 | 15 | def calibration_transformers(): 16 | """ 17 | step： 18 | 1、pad resize to 672 * 672 19 | 2、NHWC to NCHW 20 | 3、bgr to rgb 21 | """ 22 | transformers = [ 23 | PadResizeTransformer(target_size=(672, 672)), 24 | HWC2CHWTransformer(), 25 | BGR2RGBTransformer(data_format="CHW"), 26 | ] 27 | return transformers 28 | 29 | 30 | def infer_transformers(input_shape, input_layout="NHWC"): 31 | """ 32 | step： 33 | 1、pad resize to target_size(input_shape) 34 | 2、bgr to rgb 35 | 3、rgb to nv12 36 | 3、nv12 to yuv444 37 | :param input_shape: input shape(target size) 38 | :param input_layout: NCHW / NHWC 39 | """ 40 | transformers = [ 41 | PadResizeTransformer(target_size=input_shape), 42 | BGR2RGBTransformer(data_format="HWC"), 43 | RGB2NV12Transformer(data_format="HWC"), 44 | NV12ToYUV444Transformer(target_size=input_shape, 45 | yuv444_output_layout=input_layout[1:]), 46 | ] 47 | return transformers 48 | 49 | 50 | def infer_image_preprocess(image_file, input_layout, input_shape): 51 | """ 52 | image for single image inference 53 | note: imread_mode [skimage / opencv] 54 | opencv read image as 8-bit unsigned integers BGR in range [0, 255] 55 | skimage read image as float32 RGB in range [0, 1] 56 | make sure to use the same imread_mode as the model training 57 | :param image_file: image file 58 | :param input_layout: NCHW / NHWC 59 | :param input_shape: input shape（target size） 60 | :return: origin image, processed image (uint8, 0-255) 61 | """ 62 | transformers = infer_transformers(input_shape, input_layout) 63 | origin_image, processed_image = SingleImageDataLoaderWithOrigin( 64 | transformers, image_file, imread_mode="opencv") 65 | return origin_image, processed_image 66 | 67 | 68 | def eval_image_preprocess(image_path, annotation_path, input_shape, 69 | input_layout): 70 | """ 71 | image for full scale evaluation 72 | note: imread_mode [skimage / opencv] 73 | opencv read image as 8-bit unsigned integers BGR in range [0, 255] 74 | skimage read image as float32 RGB in range [0, 1] 75 | make sure to use the same imread_mode as the model training 76 | :param image_path: image path 77 | :param annotation_path: annotation path 78 | :param input_shape: input shape（target size） 79 | :param input_layout: input layout 80 | :return: data loader 81 | """ 82 | transformers = infer_transformers(input_shape, input_layout) 83 | data_loader = COCODataLoader(transformers, 84 | image_path, 85 | annotation_path, 86 | imread_mode='opencv') 87 | 88 | return data_loader 89 | -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/src_data/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/src_data/test.jpg -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/test.jpg -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/test_onnx_result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_horizon/mapper/test_onnx_result.jpg -------------------------------------------------------------------------------- /yolov8obb_horizon/mapper/yolov8obb_config.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020 Horizon Robotics.All Rights Reserved. 2 | # 3 | # The material in this file is confidential and contains trade secrets 4 | # of Horizon Robotics Inc. This is proprietary information owned by 5 | # Horizon Robotics Inc. No part of this work may be disclosed, 6 | # reproduced, copied, transmitted, or used in any way for any purpose, 7 | # without the express written permission of Horizon Robotics Inc. 8 | 9 | # 模型转化相关的参数 10 | # ------------------------------------ 11 | # model conversion related parameters 12 | model_parameters: 13 | # Onnx浮点网络数据模型文件 14 | # ----------------------------------------------------------- 15 | # the model file of floating-point ONNX neural network data 16 | onnx_model: './model/yolov8n-obb.onnx' 17 | 18 | # 适用BPU架构 19 | # -------------------------------- 20 | # the applicable BPU architecture 21 | march: "bernoulli2" 22 | 23 | # 指定模型转换过程中是否输出各层的中间结果，如果为True，则输出所有层的中间输出结果， 24 | # -------------------------------------------------------------------------------------- 25 | # specifies whether or not to dump the intermediate results of all layers in conversion 26 | # if set to True, then the intermediate results of all layers shall be dumped 27 | layer_out_dump: False 28 | 29 | # 日志文件的输出控制参数， 30 | # debug输出模型转换的详细信息 31 | # info只输出关键信息 32 | # warn输出警告和错误级别以上的信息 33 | # ---------------------------------------------------------------------------------------- 34 | # output control parameter of log file(s), 35 | # if set to 'debug', then details of model conversion will be dumped 36 | # if set to 'info', then only important imformation will be dumped 37 | # if set to 'warn', then information ranked higher than 'warn' and 'error' will be dumped 38 | log_level: 'debug' 39 | 40 | # 模型转换输出的结果的存放目录 41 | # ----------------------------------------------------------- 42 | # the directory in which model conversion results are stored 43 | working_dir: 'model_output' 44 | 45 | # 模型转换输出的用于上板执行的模型文件的名称前缀 46 | # ----------------------------------------------------------------------------------------- 47 | # model conversion generated name prefix of those model files used for dev board execution 48 | output_model_file_prefix: 'yolov8obb' 49 | 50 | # 模型输入相关参数, 若输入多个节点, 则应使用';'进行分隔, 使用默认缺省设置则写None 51 | # -------------------------------------------------------------------------- 52 | # model input related parameters, 53 | # please use ";" to seperate when inputting multiple nodes, 54 | # please use None for default setting 55 | input_parameters: 56 | 57 | # (选填) 模型输入的节点名称, 此名称应与模型文件中的名称一致, 否则会报错, 不填则会使用模型文件中的节点名称 58 | # -------------------------------------------------------------------------------------------------------- 59 | # (Optional) node name of model input, 60 | # it shall be the same as the name of model file, otherwise an error will be reported, 61 | # the node name of model file will be used when left blank 62 | input_name: "" 63 | 64 | # 网络实际执行时，输入给网络的数据格式，包括 nv12/rgb/bgr/yuv444/gray/featuremap, 65 | # ------------------------------------------------------------------------------------------ 66 | # the data formats to be passed into neural network when actually performing neural network 67 | # available options: nv12/rgb/bgr/yuv444/gray/featuremap, 68 | input_type_rt: 'rgb' 69 | 70 | # 网络实际执行时输入的数据排布, 可选值为 NHWC/NCHW 71 | # 若input_type_rt配置为nv12，则此处参数不需要配置 72 | # ------------------------------------------------------------------ 73 | # the data layout formats to be passed into neural network when actually performing neural network, available options: NHWC/NCHW 74 | # If input_type_rt is configured as nv12, then this parameter does not need to be configured 75 | input_layout_rt: 'NCHW' 76 | 77 | # 网络训练时输入的数据格式，可选的值为rgb/bgr/gray/featuremap/yuv444 78 | # -------------------------------------------------------------------- 79 | # the data formats in network training 80 | # available options: rgb/bgr/gray/featuremap/yuv444 81 | input_type_train: 'rgb' 82 | 83 | # 网络训练时输入的数据排布, 可选值为 NHWC/NCHW 84 | # ------------------------------------------------------------------ 85 | # the data layout in network training, available options: NHWC/NCHW 86 | input_layout_train: 'NCHW' 87 | 88 | # (选填) 模型网络的输入大小, 以'x'分隔, 不填则会使用模型文件中的网络输入大小，否则会覆盖模型文件中输入大小 89 | # ------------------------------------------------------------------------------------------- 90 | # (Optional)the input size of model network, seperated by 'x' 91 | # note that the network input size of model file will be used if left blank 92 | # otherwise it will overwrite the input size of model file 93 | input_shape: '' 94 | 95 | # 网络实际执行时，输入给网络的batch_size, 默认值为1 96 | # --------------------------------------------------------------------- 97 | # the data batch_size to be passed into neural network when actually performing neural network, default value: 1 98 | #input_batch: 1 99 | 100 | # 网络输入的预处理方法，主要有以下几种： 101 | # no_preprocess 不做任何操作 102 | # data_mean 减去通道均值mean_value 103 | # data_scale 对图像像素乘以data_scale系数 104 | # data_mean_and_scale 减去通道均值后再乘以scale系数 105 | # ------------------------------------------------------------------------------------------- 106 | # preprocessing methods of network input, available options: 107 | # 'no_preprocess' indicates that no preprocess will be made 108 | # 'data_mean' indicates that to minus the channel mean, i.e. mean_value 109 | # 'data_scale' indicates that image pixels to multiply data_scale ratio 110 | # 'data_mean_and_scale' indicates that to multiply scale ratio after channel mean is minused 111 | norm_type: 'data_scale' 112 | 113 | # 图像减去的均值, 如果是通道均值，value之间必须用空格分隔 114 | # -------------------------------------------------------------------------- 115 | # the mean value minused by image 116 | # note that values must be seperated by space if channel mean value is used 117 | mean_value: '' 118 | 119 | # 图像预处理缩放比例，如果是通道缩放比例，value之间必须用空格分隔 120 | # --------------------------------------------------------------------------- 121 | # scale value of image preprocess 122 | # note that values must be seperated by space if channel scale value is used 123 | scale_value: 0.003921568627451 124 | 125 | # 模型量化相关参数 126 | # ----------------------------- 127 | # model calibration parameters 128 | calibration_parameters: 129 | 130 | # 模型量化的参考图像的存放目录，图片格式支持Jpeg、Bmp等格式，输入的图片 131 | # 应该是使用的典型场景，一般是从测试集中选择20~100张图片，另外输入 132 | # 的图片要覆盖典型场景，不要是偏僻场景，如过曝光、饱和、模糊、纯黑、纯白等图片 133 | # 若有多个输入节点, 则应使用';'进行分隔 134 | # ------------------------------------------------------------------------------------------------- 135 | # the directory where reference images of model quantization are stored 136 | # image formats include JPEG, BMP etc. 137 | # should be classic application scenarios, usually 20~100 images are picked out from test datasets 138 | # in addition, note that input images should cover typical scenarios 139 | # and try to avoid those overexposed, oversaturated, vague, 140 | # pure blank or pure white images 141 | # use ';' to seperate when there are multiple input nodes 142 | cal_data_dir: './cal_data' 143 | 144 | # 如果输入的图片文件尺寸和模型训练的尺寸不一致时，并且preprocess_on为true， 145 | # 则将采用默认预处理方法(skimage resize)， 146 | # 将输入图片缩放或者裁减到指定尺寸，否则，需要用户提前把图片处理为训练时的尺寸 147 | # --------------------------------------------------------------------------------- 148 | # In case the size of input image file is different from that of in model training 149 | # and that preprocess_on is set to True, 150 | # shall the default preprocess method(skimage resize) be used 151 | # i.e., to resize or crop input image into specified size 152 | # otherwise user must keep image size as that of in training in advance 153 | preprocess_on: True 154 | 155 | # 模型量化的算法类型，支持kl、max、default、load，通常采用default即可满足要求, 若为QAT导出的模型, 则应选择load 156 | # ---------------------------------------------------------------------------------- 157 | # types of model quantization algorithms, usually default will meet the need 158 | # available options:kl, max, default and load 159 | # if converted model is quanti model exported from QAT , then choose `load` 160 | calibration_type: 'default' 161 | 162 | # 编译器相关参数 163 | # ---------------------------- 164 | # compiler related parameters 165 | compiler_parameters: 166 | 167 | # 编译策略，支持bandwidth和latency两种优化模式; 168 | # bandwidth以优化ddr的访问带宽为目标； 169 | # latency以优化推理时间为目标 170 | # ------------------------------------------------------------------------------------------- 171 | # compilation strategy, there are 2 available optimization modes: 'bandwidth' and 'lantency' 172 | # the 'bandwidth' mode aims to optimize ddr access bandwidth 173 | # while the 'lantency' mode aims to optimize inference duration 174 | compile_mode: 'latency' 175 | 176 | # 设置debug为True将打开编译器的debug模式，能够输出性能仿真的相关信息，如帧率、DDR带宽占用等 177 | # ----------------------------------------------------------------------------------- 178 | # the compiler's debug mode will be enabled by setting to True 179 | # this will dump performance simulation related information 180 | # such as: frame rate, DDR bandwidth usage etc. 181 | debug: False 182 | 183 | # 编译模型指定核数，不指定默认编译单核模型, 若编译双核模型，将下边注释打开即可 184 | # ------------------------------------------------------------------------------------- 185 | # specifies number of cores to be used in model compilation 186 | # as default, single core is used as this value left blank 187 | # please delete the "# " below to enable dual-core mode when compiling dual-core model 188 | # core_num: 2 189 | 190 | # 优化等级可选范围为O0~O3 191 | # O0不做任何优化, 编译速度最快，优化程度最低, 192 | # O1-O3随着优化等级提高，预期编译后的模型的执行速度会更快，但是所需编译时间也会变长。 193 | # 推荐用O2做最快验证 194 | # ---------------------------------------------------------------------------------------------------------- 195 | # optimization level ranges between O0~O3 196 | # O0 indicates that no optimization will be made 197 | # the faster the compilation, the lower optimization level will be 198 | # O1-O3: as optimization levels increase gradually, model execution, after compilation, shall become faster 199 | # while compilation will be prolonged 200 | # it is recommended to use O2 for fastest verification 201 | optimize_level: 'O3' 202 | -------------------------------------------------------------------------------- /yolov8obb_onnx/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_onnx/test.jpg -------------------------------------------------------------------------------- /yolov8obb_onnx/test_onnx_result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_onnx/test_onnx_result.jpg -------------------------------------------------------------------------------- /yolov8obb_onnx/yolov8n-obb.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_onnx/yolov8n-obb.onnx -------------------------------------------------------------------------------- /yolov8obb_onnx/yolov8obb_onnx_demo.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import argparse 4 | import os 5 | import sys 6 | import cv2 7 | import numpy as np 8 | import math 9 | import onnxruntime as ort 10 | 11 | ROOT = os.getcwd() 12 | if str(ROOT) not in sys.path: 13 | sys.path.append(str(ROOT)) 14 | 15 | 16 | CLASSES = ['plane', 'ship', 'storage tank', 'baseball diamond', 'tennis court', 'basketball court', 17 | 'ground track field', 'harbor', 'bridge', 'large vehicle', 'small vehicle', 'helicopter', 'roundabout', 18 | 'soccer ball field', 'swimming pool'] 19 | 20 | meshgrid = [] 21 | 22 | class_num = len(CLASSES) 23 | head_num = 3 24 | strides = [8, 16, 32] 25 | map_size = [[80, 80], [40, 40], [20, 20]] 26 | reg_num = 16 27 | nms_thresh = 0.5 28 | object_thresh = 0.25 29 | 30 | input_height = 640 31 | input_width = 640 32 | 33 | 34 | 35 | class CSXYWHR: 36 | def __init__(self, classId, score, x, y, w, h, angle): 37 | self.classId = classId 38 | self.score = score 39 | self.x = x 40 | self.y = y 41 | self.w = w 42 | self.h = h 43 | self.angle = angle 44 | 45 | 46 | class DetectBox: 47 | def __init__(self, classId, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle): 48 | self.classId = classId 49 | self.score = score 50 | self.pt1x = pt1x 51 | self.pt1y = pt1y 52 | self.pt2x = pt2x 53 | self.pt2y = pt2y 54 | self.pt3x = pt3x 55 | self.pt3y = pt3y 56 | self.pt4x = pt4x 57 | self.pt4y = pt4y 58 | self.angle = angle 59 | 60 | 61 | def GenerateMeshgrid(): 62 | for index in range(head_num): 63 | for i in range(map_size[index][0]): 64 | for j in range(map_size[index][1]): 65 | meshgrid.append(j + 0.5) 66 | meshgrid.append(i + 0.5) 67 | 68 | 69 | def get_covariance_matrix(boxes): 70 | a, b, c = boxes.w, boxes.h, boxes.angle 71 | cos = math.cos(c) 72 | sin = math.sin(c) 73 | cos2 = math.pow(cos, 2) 74 | sin2 = math.pow(sin, 2) 75 | return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin 76 | 77 | 78 | def probiou(obb1, obb2, eps=1e-7): 79 | x1, y1 = obb1.x, obb1.y 80 | x2, y2 = obb2.x, obb2.y 81 | a1, b1, c1 = get_covariance_matrix(obb1) 82 | a2, b2, c2 = get_covariance_matrix(obb2) 83 | 84 | t1 = (((a1 + a2) * math.pow((y1 - y2), 2) + (b1 + b2) * math.pow((x1 - x2), 2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.25 85 | t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.5 86 | 87 | temp1 = (a1 * b1 - math.pow(c1, 2)) if (a1 * b1 - math.pow(c1, 2)) > 0 else 0 88 | temp2 = (a2 * b2 - math.pow(c2, 2)) if (a2 * b2 - math.pow(c2, 2)) > 0 else 0 89 | t3 = math.log((((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2)) / (4 * math.sqrt((temp1 * temp2)) + eps)+ eps)) * 0.5 90 | 91 | if (t1 + t2 + t3) > 100: 92 | bd = 100 93 | elif (t1 + t2 + t3) < eps: 94 | bd = eps 95 | else: 96 | bd = t1 + t2 + t3 97 | hd = math.sqrt((1.0 - math.exp(-bd) + eps)) 98 | return 1 - hd 99 | 100 | 101 | def nms_rotated(boxes, nms_thresh): 102 | pred_boxes = [] 103 | sort_boxes = sorted(boxes, key=lambda x: x.score, reverse=True) 104 | for i in range(len(sort_boxes)): 105 | if sort_boxes[i].classId != -1: 106 | pred_boxes.append(sort_boxes[i]) 107 | for j in range(i + 1, len(sort_boxes), 1): 108 | ious = probiou(sort_boxes[i], sort_boxes[j]) 109 | if ious > nms_thresh: 110 | sort_boxes[j].classId = -1 111 | return pred_boxes 112 | 113 | 114 | def sigmoid(x): 115 | return 1 / (1 + math.exp(-x)) 116 | 117 | 118 | def xywhr2xyxyxyxy(x, y, w, h, angle): 119 | cos_value = math.cos(angle) 120 | sin_value = math.sin(angle) 121 | 122 | vec1x= w / 2 * cos_value 123 | vec1y = w / 2 * sin_value 124 | vec2x = -h / 2 * sin_value 125 | vec2y = h / 2 * cos_value 126 | 127 | pt1x = x + vec1x + vec2x 128 | pt1y = y + vec1y + vec2y 129 | 130 | pt2x = x + vec1x - vec2x 131 | pt2y = y + vec1y - vec2y 132 | 133 | pt3x = x - vec1x - vec2x 134 | pt3y = y - vec1y - vec2y 135 | 136 | pt4x = x - vec1x + vec2x 137 | pt4y = y - vec1y + vec2y 138 | return pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y 139 | 140 | 141 | def postprocess(out): 142 | print('postprocess ... ') 143 | 144 | detect_result = [] 145 | output = [] 146 | for i in range(len(out)): 147 | output.append(out[i].reshape((-1))) 148 | 149 | gridIndex = -2 150 | cls_index = 0 151 | cls_max = 0 152 | 153 | for index in range(head_num): 154 | reg = output[index * 2 + 0] 155 | cls = output[index * 2 + 1] 156 | ang = output[head_num * 2 + index] 157 | 158 | for h in range(map_size[index][0]): 159 | for w in range(map_size[index][1]): 160 | gridIndex += 2 161 | 162 | if 1 == class_num: 163 | cls_max = sigmoid(cls[0 * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 164 | cls_index = 0 165 | else: 166 | for cl in range(class_num): 167 | cls_val = cls[cl * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] 168 | if 0 == cl: 169 | cls_max = cls_val 170 | cls_index = cl 171 | else: 172 | if cls_val > cls_max: 173 | cls_max = cls_val 174 | cls_index = cl 175 | cls_max = sigmoid(cls_max) 176 | 177 | if cls_max > object_thresh: 178 | regdfl = [] 179 | for lc in range(4): 180 | sfsum = 0 181 | locval = 0 182 | for df in range(reg_num): 183 | temp = math.exp(reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 184 | reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] = temp 185 | sfsum += temp 186 | 187 | for df in range(reg_num): 188 | sfval = reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] / sfsum 189 | locval += sfval * df 190 | regdfl.append(locval) 191 | 192 | angle = (sigmoid(ang[h * map_size[index][1] + w]) - 0.25) * math.pi 193 | 194 | left, top, right, bottom = regdfl[0], regdfl[1], regdfl[2], regdfl[3] 195 | cos, sin = math.cos(angle), math.sin(angle) 196 | fx = (right - left) / 2 197 | fy = (bottom - top) / 2 198 | 199 | cx = ((fx * cos - fy * sin) + meshgrid[gridIndex + 0]) * strides[index] 200 | cy = ((fx * sin + fy * cos) + meshgrid[gridIndex + 1])* strides[index] 201 | cw = (left + right) * strides[index] 202 | ch = (top + bottom) * strides[index] 203 | 204 | box = CSXYWHR(cls_index, cls_max, cx, cy, cw, ch, angle) 205 | 206 | detect_result.append(box) 207 | # NMS 208 | print('before nms num is:', len(detect_result)) 209 | pred_boxes = nms_rotated(detect_result, nms_thresh) 210 | 211 | print('after nms num is:', len(pred_boxes)) 212 | 213 | resutl = [] 214 | for i in range(len(pred_boxes)): 215 | classid = pred_boxes[i].classId 216 | score = pred_boxes[i].score 217 | cx = pred_boxes[i].x 218 | cy = pred_boxes[i].y 219 | cw = pred_boxes[i].w 220 | ch = pred_boxes[i].h 221 | angle = pred_boxes[i].angle 222 | 223 | bw_ = cw if cw > ch else ch 224 | bh_ = ch if cw > ch else cw 225 | bt = angle % math.pi if cw > ch else (angle + math.pi / 2) % math.pi 226 | 227 | pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y = xywhr2xyxyxyxy(cx, cy, bw_, bh_, bt) 228 | 229 | bbox = DetectBox(classid, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle) 230 | resutl.append(bbox) 231 | 232 | return resutl 233 | 234 | 235 | def precess_image(img_src, resize_w, resize_h): 236 | image = cv2.resize(img_src, (resize_w, resize_h), interpolation=cv2.INTER_LINEAR) 237 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 238 | image = image.astype(np.float32) 239 | image /= 255.0 240 | 241 | return image 242 | 243 | 244 | def detect(image_path): 245 | 246 | origin_image = cv2.imread(image_path) 247 | image_h, image_w = origin_image.shape[:2] 248 | image = precess_image(origin_image, input_width, input_height) 249 | 250 | image = image.transpose((2, 0, 1)) 251 | image = np.expand_dims(image, axis=0) 252 | 253 | ort_session = ort.InferenceSession('./yolov8n-obb.onnx') 254 | pred_results = (ort_session.run(None, {'data': image})) 255 | 256 | out = [] 257 | for i in range(len(pred_results)): 258 | out.append(pred_results[i]) 259 | pred_boxes = postprocess(out) 260 | 261 | print('obj num is :', len(pred_boxes)) 262 | 263 | for i in range(len(pred_boxes)): 264 | classId = pred_boxes[i].classId 265 | score = pred_boxes[i].score 266 | pt1x = int(pred_boxes[i].pt1x / input_width * image_w) 267 | pt1y = int(pred_boxes[i].pt1y / input_width * image_h) 268 | pt2x = int(pred_boxes[i].pt2x / input_width * image_w) 269 | pt2y = int(pred_boxes[i].pt2y / input_width * image_h) 270 | pt3x = int(pred_boxes[i].pt3x / input_width * image_w) 271 | pt3y = int(pred_boxes[i].pt3y / input_width * image_h) 272 | pt4x = int(pred_boxes[i].pt4x / input_width * image_w) 273 | pt4y = int(pred_boxes[i].pt4y / input_width * image_h) 274 | angle = pred_boxes[i].angle 275 | 276 | cv2.line(origin_image, (pt1x, pt1y), (pt2x, pt2y), (0, 255, 0), 2) 277 | cv2.line(origin_image, (pt2x, pt2y), (pt3x, pt3y), (0, 255, 0), 2) 278 | cv2.line(origin_image, (pt3x, pt3y), (pt4x, pt4y), (0, 255, 0), 2) 279 | cv2.line(origin_image, (pt4x, pt4y), (pt1x, pt1y), (0, 255, 0), 2) 280 | 281 | title = CLASSES[classId] + "%.2f" % score 282 | cv2.putText(origin_image, title, (pt1x, pt1y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2, cv2.LINE_AA) 283 | 284 | cv2.imwrite('./test_onnx_result.jpg', origin_image) 285 | 286 | 287 | if __name__ == '__main__': 288 | print('This is main ....') 289 | GenerateMeshgrid() 290 | image_path = './test.jpg' 291 | detect(image_path) -------------------------------------------------------------------------------- /yolov8obb_rknn/images_list.txt: -------------------------------------------------------------------------------- 1 | ./images_list/test.jpg 2 | -------------------------------------------------------------------------------- /yolov8obb_rknn/images_list/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_rknn/images_list/test.jpg -------------------------------------------------------------------------------- /yolov8obb_rknn/onnx2rknn_demo_ZQ.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib 3 | import traceback 4 | import time 5 | import sys 6 | import numpy as np 7 | import cv2 8 | from rknn.api import RKNN 9 | import math 10 | 11 | ONNX_MODEL = './yolov8n-obb.onnx' 12 | RKNN_MODEL = './yyolov8n-obb.rknn' 13 | DATASET = './images_list.txt' 14 | 15 | QUANTIZE_ON = True 16 | 17 | 18 | CLASSES = ['plane', 'ship', 'storage tank', 'baseball diamond', 'tennis court', 'basketball court', 19 | 'ground track field', 'harbor', 'bridge', 'large vehicle', 'small vehicle', 'helicopter', 'roundabout', 20 | 'soccer ball field', 'swimming pool'] 21 | 22 | meshgrid = [] 23 | 24 | class_num = len(CLASSES) 25 | head_num = 3 26 | strides = [8, 16, 32] 27 | map_size = [[80, 80], [40, 40], [20, 20]] 28 | reg_num = 16 29 | nms_thresh = 0.5 30 | object_thresh = 0.25 31 | 32 | input_height = 640 33 | input_width = 640 34 | 35 | 36 | 37 | class CSXYWHR: 38 | def __init__(self, classId, score, x, y, w, h, angle): 39 | self.classId = classId 40 | self.score = score 41 | self.x = x 42 | self.y = y 43 | self.w = w 44 | self.h = h 45 | self.angle = angle 46 | 47 | 48 | class DetectBox: 49 | def __init__(self, classId, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle): 50 | self.classId = classId 51 | self.score = score 52 | self.pt1x = pt1x 53 | self.pt1y = pt1y 54 | self.pt2x = pt2x 55 | self.pt2y = pt2y 56 | self.pt3x = pt3x 57 | self.pt3y = pt3y 58 | self.pt4x = pt4x 59 | self.pt4y = pt4y 60 | self.angle = angle 61 | 62 | 63 | def GenerateMeshgrid(): 64 | for index in range(head_num): 65 | for i in range(map_size[index][0]): 66 | for j in range(map_size[index][1]): 67 | meshgrid.append(j + 0.5) 68 | meshgrid.append(i + 0.5) 69 | 70 | 71 | def get_covariance_matrix(boxes): 72 | a, b, c = boxes.w, boxes.h, boxes.angle 73 | cos = math.cos(c) 74 | sin = math.sin(c) 75 | cos2 = math.pow(cos, 2) 76 | sin2 = math.pow(sin, 2) 77 | return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin 78 | 79 | 80 | def probiou(obb1, obb2, eps=1e-7): 81 | x1, y1 = obb1.x, obb1.y 82 | x2, y2 = obb2.x, obb2.y 83 | a1, b1, c1 = get_covariance_matrix(obb1) 84 | a2, b2, c2 = get_covariance_matrix(obb2) 85 | 86 | t1 = (((a1 + a2) * math.pow((y1 - y2), 2) + (b1 + b2) * math.pow((x1 - x2), 2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.25 87 | t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.5 88 | 89 | temp1 = (a1 * b1 - math.pow(c1, 2)) if (a1 * b1 - math.pow(c1, 2)) > 0 else 0 90 | temp2 = (a2 * b2 - math.pow(c2, 2)) if (a2 * b2 - math.pow(c2, 2)) > 0 else 0 91 | t3 = math.log((((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2)) / (4 * math.sqrt((temp1 * temp2)) + eps)+ eps)) * 0.5 92 | 93 | if (t1 + t2 + t3) > 100: 94 | bd = 100 95 | elif (t1 + t2 + t3) < eps: 96 | bd = eps 97 | else: 98 | bd = t1 + t2 + t3 99 | hd = math.sqrt((1.0 - math.exp(-bd) + eps)) 100 | return 1 - hd 101 | 102 | 103 | def nms_rotated(boxes, nms_thresh): 104 | pred_boxes = [] 105 | sort_boxes = sorted(boxes, key=lambda x: x.score, reverse=True) 106 | for i in range(len(sort_boxes)): 107 | if sort_boxes[i].classId != -1: 108 | pred_boxes.append(sort_boxes[i]) 109 | for j in range(i + 1, len(sort_boxes), 1): 110 | ious = probiou(sort_boxes[i], sort_boxes[j]) 111 | if ious > nms_thresh: 112 | sort_boxes[j].classId = -1 113 | return pred_boxes 114 | 115 | 116 | def sigmoid(x): 117 | return 1 / (1 + math.exp(-x)) 118 | 119 | 120 | def xywhr2xyxyxyxy(x, y, w, h, angle): 121 | cos_value = math.cos(angle) 122 | sin_value = math.sin(angle) 123 | 124 | vec1x= w / 2 * cos_value 125 | vec1y = w / 2 * sin_value 126 | vec2x = -h / 2 * sin_value 127 | vec2y = h / 2 * cos_value 128 | 129 | pt1x = x + vec1x + vec2x 130 | pt1y = y + vec1y + vec2y 131 | 132 | pt2x = x + vec1x - vec2x 133 | pt2y = y + vec1y - vec2y 134 | 135 | pt3x = x - vec1x - vec2x 136 | pt3y = y - vec1y - vec2y 137 | 138 | pt4x = x - vec1x + vec2x 139 | pt4y = y - vec1y + vec2y 140 | return pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y 141 | 142 | 143 | def postprocess(out): 144 | print('postprocess ... ') 145 | 146 | detect_result = [] 147 | output = [] 148 | for i in range(len(out)): 149 | output.append(out[i].reshape((-1))) 150 | 151 | gridIndex = -2 152 | cls_index = 0 153 | cls_max = 0 154 | 155 | for index in range(head_num): 156 | reg = output[index * 2 + 0] 157 | cls = output[index * 2 + 1] 158 | ang = output[head_num * 2 + index] 159 | 160 | for h in range(map_size[index][0]): 161 | for w in range(map_size[index][1]): 162 | gridIndex += 2 163 | 164 | if 1 == class_num: 165 | cls_max = sigmoid(cls[0 * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 166 | cls_index = 0 167 | else: 168 | for cl in range(class_num): 169 | cls_val = cls[cl * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] 170 | if 0 == cl: 171 | cls_max = cls_val 172 | cls_index = cl 173 | else: 174 | if cls_val > cls_max: 175 | cls_max = cls_val 176 | cls_index = cl 177 | cls_max = sigmoid(cls_max) 178 | 179 | if cls_max > object_thresh: 180 | regdfl = [] 181 | for lc in range(4): 182 | sfsum = 0 183 | locval = 0 184 | for df in range(reg_num): 185 | temp = math.exp(reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 186 | reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] = temp 187 | sfsum += temp 188 | 189 | for df in range(reg_num): 190 | sfval = reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] / sfsum 191 | locval += sfval * df 192 | regdfl.append(locval) 193 | 194 | angle = (sigmoid(ang[h * map_size[index][1] + w]) - 0.25) * math.pi 195 | 196 | left, top, right, bottom = regdfl[0], regdfl[1], regdfl[2], regdfl[3] 197 | cos, sin = math.cos(angle), math.sin(angle) 198 | fx = (right - left) / 2 199 | fy = (bottom - top) / 2 200 | 201 | cx = ((fx * cos - fy * sin) + meshgrid[gridIndex + 0]) * strides[index] 202 | cy = ((fx * sin + fy * cos) + meshgrid[gridIndex + 1])* strides[index] 203 | cw = (left + right) * strides[index] 204 | ch = (top + bottom) * strides[index] 205 | 206 | box = CSXYWHR(cls_index, cls_max, cx, cy, cw, ch, angle) 207 | 208 | detect_result.append(box) 209 | # NMS 210 | print('before nms num is:', len(detect_result)) 211 | pred_boxes = nms_rotated(detect_result, nms_thresh) 212 | 213 | print('after nms num is:', len(pred_boxes)) 214 | 215 | resutl = [] 216 | for i in range(len(pred_boxes)): 217 | classid = pred_boxes[i].classId 218 | score = pred_boxes[i].score 219 | cx = pred_boxes[i].x 220 | cy = pred_boxes[i].y 221 | cw = pred_boxes[i].w 222 | ch = pred_boxes[i].h 223 | angle = pred_boxes[i].angle 224 | 225 | bw_ = cw if cw > ch else ch 226 | bh_ = ch if cw > ch else cw 227 | bt = angle % math.pi if cw > ch else (angle + math.pi / 2) % math.pi 228 | 229 | pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y = xywhr2xyxyxyxy(cx, cy, bw_, bh_, bt) 230 | 231 | bbox = DetectBox(classid, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle) 232 | resutl.append(bbox) 233 | 234 | return resutl 235 | 236 | 237 | 238 | 239 | def export_rknn_inference(img): 240 | # Create RKNN object 241 | rknn = RKNN(verbose=False) 242 | 243 | # pre-process config 244 | print('--> Config model') 245 | rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588') 246 | print('done') 247 | 248 | # Load ONNX model 249 | print('--> Loading model') 250 | ret = rknn.load_onnx(model=ONNX_MODEL, outputs=["reg1", "cls2", "reg3", "cls4", "reg5", "cls6", "ang1", "ang2", "ang3"]) 251 | if ret != 0: 252 | print('Load model failed!') 253 | exit(ret) 254 | print('done') 255 | 256 | # Build model 257 | print('--> Building model') 258 | ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET, rknn_batch_size=1) 259 | if ret != 0: 260 | print('Build model failed!') 261 | exit(ret) 262 | print('done') 263 | 264 | # Export RKNN model 265 | print('--> Export rknn model') 266 | ret = rknn.export_rknn(RKNN_MODEL) 267 | if ret != 0: 268 | print('Export rknn model failed!') 269 | exit(ret) 270 | print('done') 271 | 272 | # Init runtime environment 273 | print('--> Init runtime environment') 274 | ret = rknn.init_runtime() 275 | # ret = rknn.init_runtime(target='rk3566') 276 | if ret != 0: 277 | print('Init runtime environment failed!') 278 | exit(ret) 279 | print('done') 280 | 281 | # Inference 282 | print('--> Running model') 283 | outputs = rknn.inference(inputs=[img]) 284 | rknn.release() 285 | print('done') 286 | 287 | return outputs 288 | 289 | 290 | if __name__ == '__main__': 291 | print('This is main ...') 292 | GenerateMeshgrid() 293 | image_path = './test.jpg' 294 | origin_image = cv2.imread(image_path) 295 | image_h, image_w = origin_image.shape[:2] 296 | 297 | image = cv2.resize(origin_image, (input_width, input_height), interpolation=cv2.INTER_LINEAR) 298 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 299 | 300 | image = np.expand_dims(image, 0) 301 | outputs = export_rknn_inference(image) 302 | 303 | out = [] 304 | for i in range(len(outputs)): 305 | out.append(outputs[i]) 306 | pred_boxes = postprocess(out) 307 | 308 | print('obj num is :', len(pred_boxes)) 309 | 310 | for i in range(len(pred_boxes)): 311 | classId = pred_boxes[i].classId 312 | score = pred_boxes[i].score 313 | pt1x = int(pred_boxes[i].pt1x / input_width * image_w) 314 | pt1y = int(pred_boxes[i].pt1y / input_width * image_h) 315 | pt2x = int(pred_boxes[i].pt2x / input_width * image_w) 316 | pt2y = int(pred_boxes[i].pt2y / input_width * image_h) 317 | pt3x = int(pred_boxes[i].pt3x / input_width * image_w) 318 | pt3y = int(pred_boxes[i].pt3y / input_width * image_h) 319 | pt4x = int(pred_boxes[i].pt4x / input_width * image_w) 320 | pt4y = int(pred_boxes[i].pt4y / input_width * image_h) 321 | angle = pred_boxes[i].angle 322 | 323 | cv2.line(origin_image, (pt1x, pt1y), (pt2x, pt2y), (0, 255, 0), 2) 324 | cv2.line(origin_image, (pt2x, pt2y), (pt3x, pt3y), (0, 255, 0), 2) 325 | cv2.line(origin_image, (pt3x, pt3y), (pt4x, pt4y), (0, 255, 0), 2) 326 | cv2.line(origin_image, (pt4x, pt4y), (pt1x, pt1y), (0, 255, 0), 2) 327 | 328 | title = CLASSES[classId] + "%.2f" % score 329 | cv2.putText(origin_image, title, (pt1x, pt1y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2, cv2.LINE_AA) 330 | 331 | cv2.imwrite('./test_onnx_result.jpg', origin_image) 332 | 333 | -------------------------------------------------------------------------------- /yolov8obb_rknn/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_rknn/test.jpg -------------------------------------------------------------------------------- /yolov8obb_rknn/test_onnx_result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_rknn/test_onnx_result.jpg -------------------------------------------------------------------------------- /yolov8obb_rknn/yolov8n-obb.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_rknn/yolov8n-obb.onnx -------------------------------------------------------------------------------- /yolov8obb_rknn/yyolov8n-obb.rknn: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_rknn/yyolov8n-obb.rknn -------------------------------------------------------------------------------- /yolov8obb_tensorRT/onnx2trt_rt7.py: -------------------------------------------------------------------------------- 1 | import tensorrt as trt 2 | 3 | G_LOGGER = trt.Logger() 4 | 5 | batch_size = 1 6 | imput_h = 640 7 | imput_w = 640 8 | 9 | 10 | def get_engine(onnx_model_name, trt_model_name): 11 | explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) 12 | with trt.Builder(G_LOGGER) as builder, builder.create_network(explicit_batch) as network, trt.OnnxParser(network, 13 | G_LOGGER) as parser: 14 | builder.max_batch_size = batch_size 15 | builder.max_workspace_size = 2 << 30 16 | print('Loading ONNX file from path {}...'.format(onnx_model_name)) 17 | with open(onnx_model_name, 'rb') as model: 18 | print('Beginning ONNX file parsing') 19 | if not parser.parse(model.read()): 20 | for error in range(parser.num_errors): 21 | print(parser.get_error(error)) 22 | 23 | print('Completed parsing of ONNX file') 24 | print('Building an engine from file {}; this may take a while...'.format(onnx_model_name)) 25 | 26 | #### 27 | # builder.int8_mode = True 28 | # builder.int8_calibrator = calib 29 | builder.fp16_mode = True 30 | #### 31 | 32 | print("num layers:", network.num_layers) 33 | # last_layer = network.get_layer(network.num_layers - 1) 34 | # if not last_layer.get_output(0): 35 | # network.mark_output(network.get_layer(network.num_layers - 1).get_output(0))//有的模型需要，有的模型在转onnx的之后已经指定了，就不需要这行 36 | 37 | network.get_input(0).shape = [batch_size, 3, imput_h, imput_w] 38 | engine = builder.build_cuda_engine(network) 39 | print("engine:", engine) 40 | print("Completed creating Engine") 41 | with open(trt_model_name, "wb") as f: 42 | f.write(engine.serialize()) 43 | return engine 44 | 45 | 46 | def main(): 47 | onnx_file_path = './yolov8n-obb.onnx' 48 | engine_file_path = './yolov8n-obb.trt' 49 | 50 | engine = get_engine(onnx_file_path, engine_file_path) 51 | 52 | 53 | if __name__ == '__main__': 54 | print("This is main ...") 55 | main() 56 | -------------------------------------------------------------------------------- /yolov8obb_tensorRT/tensorRT_inferenc_demo.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import tensorrt as trt 4 | import pycuda.driver as cuda 5 | import pycuda.autoinit 6 | import math 7 | import time 8 | 9 | TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) 10 | 11 | 12 | CLASSES = ['plane', 'ship', 'storage tank', 'baseball diamond', 'tennis court', 'basketball court', 13 | 'ground track field', 'harbor', 'bridge', 'large vehicle', 'small vehicle', 'helicopter', 'roundabout', 14 | 'soccer ball field', 'swimming pool'] 15 | 16 | meshgrid = [] 17 | 18 | class_num = len(CLASSES) 19 | head_num = 3 20 | strides = [8, 16, 32] 21 | map_size = [[80, 80], [40, 40], [20, 20]] 22 | reg_num = 16 23 | nms_thresh = 0.5 24 | object_thresh = 0.25 25 | 26 | input_height = 640 27 | input_width = 640 28 | 29 | 30 | # Simple helper data class that's a little nicer to use than a 2-tuple. 31 | class HostDeviceMem(object): 32 | def __init__(self, host_mem, device_mem): 33 | self.host = host_mem 34 | self.device = device_mem 35 | 36 | def __str__(self): 37 | return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) 38 | 39 | def __repr__(self): 40 | return self.__str__() 41 | 42 | 43 | def allocate_buffers(engine): 44 | inputs = [] 45 | outputs = [] 46 | bindings = [] 47 | stream = cuda.Stream() 48 | for binding in engine: 49 | size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size 50 | dtype = trt.nptype(engine.get_binding_dtype(binding)) 51 | # Allocate host and device buffers 52 | host_mem = cuda.pagelocked_empty(size, dtype) 53 | device_mem = cuda.mem_alloc(host_mem.nbytes) 54 | # Append the device buffer to device bindings. 55 | bindings.append(int(device_mem)) 56 | # Append to the appropriate list. 57 | if engine.binding_is_input(binding): 58 | inputs.append(HostDeviceMem(host_mem, device_mem)) 59 | else: 60 | outputs.append(HostDeviceMem(host_mem, device_mem)) 61 | return inputs, outputs, bindings, stream 62 | 63 | 64 | def get_engine_from_bin(engine_file_path): 65 | print('Reading engine from file {}'.format(engine_file_path)) 66 | with open(engine_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime: 67 | return runtime.deserialize_cuda_engine(f.read()) 68 | 69 | 70 | # This function is generalized for multiple inputs/outputs. 71 | # inputs and outputs are expected to be lists of HostDeviceMem objects. 72 | def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): 73 | # Transfer input data to the GPU. 74 | [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] 75 | # Run inference. 76 | context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) 77 | # Transfer predictions back from the GPU. 78 | [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] 79 | # Synchronize the stream 80 | stream.synchronize() 81 | # Return only the host outputs. 82 | return [out.host for out in outputs] 83 | 84 | 85 | 86 | class CSXYWHR: 87 | def __init__(self, classId, score, x, y, w, h, angle): 88 | self.classId = classId 89 | self.score = score 90 | self.x = x 91 | self.y = y 92 | self.w = w 93 | self.h = h 94 | self.angle = angle 95 | 96 | 97 | class DetectBox: 98 | def __init__(self, classId, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle): 99 | self.classId = classId 100 | self.score = score 101 | self.pt1x = pt1x 102 | self.pt1y = pt1y 103 | self.pt2x = pt2x 104 | self.pt2y = pt2y 105 | self.pt3x = pt3x 106 | self.pt3y = pt3y 107 | self.pt4x = pt4x 108 | self.pt4y = pt4y 109 | self.angle = angle 110 | 111 | 112 | def GenerateMeshgrid(): 113 | for index in range(head_num): 114 | for i in range(map_size[index][0]): 115 | for j in range(map_size[index][1]): 116 | meshgrid.append(j + 0.5) 117 | meshgrid.append(i + 0.5) 118 | 119 | 120 | def get_covariance_matrix(boxes): 121 | a, b, c = boxes.w, boxes.h, boxes.angle 122 | cos = math.cos(c) 123 | sin = math.sin(c) 124 | cos2 = math.pow(cos, 2) 125 | sin2 = math.pow(sin, 2) 126 | return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin 127 | 128 | 129 | def probiou(obb1, obb2, eps=1e-7): 130 | x1, y1 = obb1.x, obb1.y 131 | x2, y2 = obb2.x, obb2.y 132 | a1, b1, c1 = get_covariance_matrix(obb1) 133 | a2, b2, c2 = get_covariance_matrix(obb2) 134 | 135 | t1 = (((a1 + a2) * math.pow((y1 - y2), 2) + (b1 + b2) * math.pow((x1 - x2), 2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.25 136 | t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / ((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2) + eps)) * 0.5 137 | 138 | temp1 = (a1 * b1 - math.pow(c1, 2)) if (a1 * b1 - math.pow(c1, 2)) > 0 else 0 139 | temp2 = (a2 * b2 - math.pow(c2, 2)) if (a2 * b2 - math.pow(c2, 2)) > 0 else 0 140 | t3 = math.log((((a1 + a2) * (b1 + b2) - math.pow((c1 + c2), 2)) / (4 * math.sqrt((temp1 * temp2)) + eps)+ eps)) * 0.5 141 | 142 | if (t1 + t2 + t3) > 100: 143 | bd = 100 144 | elif (t1 + t2 + t3) < eps: 145 | bd = eps 146 | else: 147 | bd = t1 + t2 + t3 148 | hd = math.sqrt((1.0 - math.exp(-bd) + eps)) 149 | return 1 - hd 150 | 151 | 152 | def nms_rotated(boxes, nms_thresh): 153 | pred_boxes = [] 154 | sort_boxes = sorted(boxes, key=lambda x: x.score, reverse=True) 155 | for i in range(len(sort_boxes)): 156 | if sort_boxes[i].classId != -1: 157 | pred_boxes.append(sort_boxes[i]) 158 | for j in range(i + 1, len(sort_boxes), 1): 159 | ious = probiou(sort_boxes[i], sort_boxes[j]) 160 | if ious > nms_thresh: 161 | sort_boxes[j].classId = -1 162 | return pred_boxes 163 | 164 | 165 | def sigmoid(x): 166 | return 1 / (1 + math.exp(-x)) 167 | 168 | 169 | def xywhr2xyxyxyxy(x, y, w, h, angle): 170 | cos_value = math.cos(angle) 171 | sin_value = math.sin(angle) 172 | 173 | vec1x= w / 2 * cos_value 174 | vec1y = w / 2 * sin_value 175 | vec2x = -h / 2 * sin_value 176 | vec2y = h / 2 * cos_value 177 | 178 | pt1x = x + vec1x + vec2x 179 | pt1y = y + vec1y + vec2y 180 | 181 | pt2x = x + vec1x - vec2x 182 | pt2y = y + vec1y - vec2y 183 | 184 | pt3x = x - vec1x - vec2x 185 | pt3y = y - vec1y - vec2y 186 | 187 | pt4x = x - vec1x + vec2x 188 | pt4y = y - vec1y + vec2y 189 | return pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y 190 | 191 | 192 | def postprocess(out): 193 | print('postprocess ... ') 194 | 195 | detect_result = [] 196 | output = [] 197 | for i in range(len(out)): 198 | output.append(out[i].reshape((-1))) 199 | 200 | gridIndex = -2 201 | cls_index = 0 202 | cls_max = 0 203 | 204 | for index in range(head_num): 205 | reg = output[head_num + index * 2 + 0] 206 | cls = output[head_num + index * 2 + 1] 207 | ang = output[index] 208 | 209 | for h in range(map_size[index][0]): 210 | for w in range(map_size[index][1]): 211 | gridIndex += 2 212 | 213 | if 1 == class_num: 214 | cls_max = sigmoid(cls[0 * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 215 | cls_index = 0 216 | else: 217 | for cl in range(class_num): 218 | cls_val = cls[cl * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] 219 | if 0 == cl: 220 | cls_max = cls_val 221 | cls_index = cl 222 | else: 223 | if cls_val > cls_max: 224 | cls_max = cls_val 225 | cls_index = cl 226 | cls_max = sigmoid(cls_max) 227 | 228 | if cls_max > object_thresh: 229 | regdfl = [] 230 | for lc in range(4): 231 | sfsum = 0 232 | locval = 0 233 | for df in range(reg_num): 234 | temp = math.exp(reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w]) 235 | reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] = temp 236 | sfsum += temp 237 | 238 | for df in range(reg_num): 239 | sfval = reg[((lc * reg_num) + df) * map_size[index][0] * map_size[index][1] + h * map_size[index][1] + w] / sfsum 240 | locval += sfval * df 241 | regdfl.append(locval) 242 | 243 | angle = (sigmoid(ang[h * map_size[index][1] + w]) - 0.25) * math.pi 244 | 245 | left, top, right, bottom = regdfl[0], regdfl[1], regdfl[2], regdfl[3] 246 | cos, sin = math.cos(angle), math.sin(angle) 247 | fx = (right - left) / 2 248 | fy = (bottom - top) / 2 249 | 250 | cx = ((fx * cos - fy * sin) + meshgrid[gridIndex + 0]) * strides[index] 251 | cy = ((fx * sin + fy * cos) + meshgrid[gridIndex + 1])* strides[index] 252 | cw = (left + right) * strides[index] 253 | ch = (top + bottom) * strides[index] 254 | 255 | box = CSXYWHR(cls_index, cls_max, cx, cy, cw, ch, angle) 256 | 257 | detect_result.append(box) 258 | # NMS 259 | print('before nms num is:', len(detect_result)) 260 | pred_boxes = nms_rotated(detect_result, nms_thresh) 261 | 262 | print('after nms num is:', len(pred_boxes)) 263 | 264 | resutl = [] 265 | for i in range(len(pred_boxes)): 266 | classid = pred_boxes[i].classId 267 | score = pred_boxes[i].score 268 | cx = pred_boxes[i].x 269 | cy = pred_boxes[i].y 270 | cw = pred_boxes[i].w 271 | ch = pred_boxes[i].h 272 | angle = pred_boxes[i].angle 273 | 274 | bw_ = cw if cw > ch else ch 275 | bh_ = ch if cw > ch else cw 276 | bt = angle % math.pi if cw > ch else (angle + math.pi / 2) % math.pi 277 | 278 | pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y = xywhr2xyxyxyxy(cx, cy, bw_, bh_, bt) 279 | 280 | bbox = DetectBox(classid, score, pt1x, pt1y, pt2x, pt2y, pt3x, pt3y, pt4x, pt4y, angle) 281 | resutl.append(bbox) 282 | 283 | return resutl 284 | 285 | 286 | def precess_image(img_src, resize_w, resize_h): 287 | image = cv2.resize(img_src, (resize_w, resize_h), interpolation=cv2.INTER_LINEAR) 288 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 289 | image = image.astype(np.float32) 290 | image /= 255.0 291 | image = image.transpose(2, 0, 1) 292 | image = image.copy() 293 | return image 294 | 295 | 296 | 297 | def main(): 298 | engine_file_path = './yolov8n-obb.trt' 299 | input_image_path = './test.jpg' 300 | 301 | origin_image = cv2.imread(input_image_path) 302 | image_h, image_w = origin_image.shape[:2] 303 | image = precess_image(origin_image, input_width, input_height) 304 | 305 | with get_engine_from_bin(engine_file_path) as engine, engine.create_execution_context() as context: 306 | inputs, outputs, bindings, stream = allocate_buffers(engine) 307 | 308 | inputs[0].host = image 309 | t1 = time.time() 310 | trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=1) 311 | t2 = time.time() 312 | print('run tiems time:', (t2 - t1)) 313 | 314 | print('outputs heads num: ', len(trt_outputs)) 315 | 316 | out = [] 317 | for i in range(len(trt_outputs)): 318 | out.append(trt_outputs[i]) 319 | pred_boxes = postprocess(out) 320 | 321 | print('obj num is :', len(pred_boxes)) 322 | 323 | for i in range(len(pred_boxes)): 324 | classId = pred_boxes[i].classId 325 | score = pred_boxes[i].score 326 | pt1x = int(pred_boxes[i].pt1x / input_width * image_w) 327 | pt1y = int(pred_boxes[i].pt1y / input_width * image_h) 328 | pt2x = int(pred_boxes[i].pt2x / input_width * image_w) 329 | pt2y = int(pred_boxes[i].pt2y / input_width * image_h) 330 | pt3x = int(pred_boxes[i].pt3x / input_width * image_w) 331 | pt3y = int(pred_boxes[i].pt3y / input_width * image_h) 332 | pt4x = int(pred_boxes[i].pt4x / input_width * image_w) 333 | pt4y = int(pred_boxes[i].pt4y / input_width * image_h) 334 | angle = pred_boxes[i].angle 335 | 336 | cv2.line(origin_image, (pt1x, pt1y), (pt2x, pt2y), (0, 255, 0), 2) 337 | cv2.line(origin_image, (pt2x, pt2y), (pt3x, pt3y), (0, 255, 0), 2) 338 | cv2.line(origin_image, (pt3x, pt3y), (pt4x, pt4y), (0, 255, 0), 2) 339 | cv2.line(origin_image, (pt4x, pt4y), (pt1x, pt1y), (0, 255, 0), 2) 340 | 341 | title = CLASSES[classId] + "%.2f" % score 342 | cv2.putText(origin_image, title, (pt1x, pt1y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2, cv2.LINE_AA) 343 | 344 | cv2.imwrite('./test_onnx_result.jpg', origin_image) 345 | 346 | if __name__ == '__main__': 347 | print('This is main ...') 348 | GenerateMeshgrid() 349 | main() 350 | -------------------------------------------------------------------------------- /yolov8obb_tensorRT/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_tensorRT/test.jpg -------------------------------------------------------------------------------- /yolov8obb_tensorRT/test_onnx_result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_tensorRT/test_onnx_result.jpg -------------------------------------------------------------------------------- /yolov8obb_tensorRT/yolov8n-obb.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_tensorRT/yolov8n-obb.onnx -------------------------------------------------------------------------------- /yolov8obb_tensorRT/yolov8n-obb.trt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cqu20160901/yolov8-obb_onnx_rknn_horizon_tensonRT/01c4095697f8b2379a8f83a3699f602886f5f652/yolov8obb_tensorRT/yolov8n-obb.trt --------------------------------------------------------------------------------