yolov8-obb 旋转目标检测瑞芯微RKNN芯片部署、地平线Horizon芯片部署、TensorRT部署

2024/10/25 12:24:28 来源：https://blog.csdn.net/zhangqian_1/article/details/139437315 浏览: 次关键词：yolov8-obb 旋转目标检测瑞芯微RKNN芯片部署、地平线Horizon芯片部署、TensorRT部署

特别说明：参考官方开源的yolov8代码、瑞芯微官方文档、地平线的官方文档，如有侵权告知删，谢谢。

模型和完整仿真测试代码，放在github上参考链接模型和代码。

折腾旋转目标检测的小伙伴们看过来，yolov8旋转目标检测部署来了。之前一直想玩一下旋转目标检测，迟迟未动手，看到yolov8可以进行旋转目标检测，最近也一直在折腾yolov8，就拿来再整一下yolov8的旋转目标检测。

1 模型和训练

老规矩只搞部署，训练参考官方开源的yolov8代码。

2 导出 onnx

导出onnx增加以下几行代码：

第一处增加：
在这里插入图片描述

        # 导出 onnx 增加y = []for i in range(self.nl):t1 = self.cv2[i](x[i])t2 = self.cv3[i](x[i])y.append(t1)y.append(t2)return y

第二处增加：
在这里插入图片描述

        # 导出onnx 增加angle_res = []for i in range(self.nl):t1 = self.cv4[i](x[i])angle_res.append(t1)x = Detect.forward(self, x)return x, angle_res

第三处增加保存onnx代码
在这里插入图片描述

        print("===========  onnx =========== ")import torchdummy_input = torch.randn(1, 3, 640, 640)input_names = ["data"]output_names = ["reg1", "cls2", "reg3", "cls4", "reg5", "cls6", "ang1", "ang2", "ang3"]torch.onnx.export(self.model, dummy_input, "./weights/yolov8n-obb.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=11)print("======================== convert onnx Finished! .... ")

以上三处增加完，运行以下代码（运行会报错，但不影响onnx文件的生成）

from ultralytics import YOLO# Load a model
model = YOLO("./weights/yolov8n-obb.pt")
results = model(task='detect', mode='predict', source='./test.jpg', imgsz=640, line_width=3, show=True, save=True, device='cpu')