特别说明:参考官方开源的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')
3 测试效果
pytorch效果
onnx 效果