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

news2024/11/18 14:00:19

特别说明：参考官方开源的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 torch
        dummy_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')