RoboMaster- RDK X5能量机关识别案例如何实现？

作者：SkyXZ CSDN：https://blog.csdn.net/xiongqi123123 博客园：https://www.cnblogs.com/SkyXZ 在RoboMaster的25赛季，我主要负责了能量机关的视觉方案开发，目前整体算法已经搭建完成，实际方案上我使用的上位机是Jetson Orin NX 16GB，其具备100TOPS的算力，在经过TensorRT优化部署后，实现了1080P原始图像从识别到PNP解算再到得到预测结果的每帧耗时仅为2.5ms，由于算法部分已经完成，且正值寒假，我准备利用手头的RDK X5进行能量机关的识别的算法验证，试试RDK X5作为RoboMaster的上位机可不可行 一、训练模型 我采用的模型是Yolov8n-Pose，数据集我们使用的是西交利物浦GMaster战队开源的数据集：zRzRzRzRzRzRzR/YOLO-of-RoboMaster-Keypoints-Detection-2023: 2023年西交利物浦大学动云科技GMaster战队yolo 装甲板四点模型，能量机关五点模型，区域赛视觉识别板目标检测其标注格式为4类别5点，具体介绍如下： 训练部分没什么好说的，配置好环境使用如下数据配置及命令进行训练即可： # buff.yaml path: buff_format train: train val: test kpt_shape: [5, 2] names: 0: RR 1: RW 2: BR 3: BW yolo pose train data=buff.yaml model=yolov8n-pose.pt epochs=200 batch=32 imgsz=640 iou=0.7 max_det=10 kobj=10 rect=True name=buff 二、量化模型 完成了模型的训练便到了我们最关键的一步模型量化啦，我们首先需要修改模型的输出头使得三个特征层的Bounding Box信息和Classify信息分开输出，具体而言，我们找到Yolov8的源码中的./ultralytics/ultralytics/nn/modules/head.py文件，接着在大约第64行的地方用如下代码替代Detect类的forward方法： def forward(self, x): # Detect result = [] for i in range(self.nl): result.append(self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()) result.append(self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()) return result 然后继续用如下代码来替换大约第242行的Pose类的forward方法： def forward(self, x): detect_results = Detect.forward(self, x) kpt_results = [] for i in range(self.nl): kpt_results.append(self.cv4[i](x[i]).permute(0, 2, 3, 1).contiguous()) return (detect_results, kpt_results) 修改完上述部分后我们便可以使用如下命令来导出ONNX模型啦： yolo export model=/path/to/your/model format=onnx simplify=True opset=11 imgsz=640 # 注意，如果生成的onnx模型显示ir版本过高，可以将simplify=False 然后我们进入地瓜的RDK算法工具链的Docker镜像（具体安装配置可见我的另外一篇Blogs：学弟一看就会的RDKX5模型转换及部署，你确定不学？），使用如下命令对我们的ONNX进行验证，之后终端便会打印出我们这个模型的基本信息、结构信息以及算子信息 hb_mapper checker --model-type onnx --march bayes-e --model /path/to/your/model.onnx 我们根据如下的打印信息可以知道我们这个模型的所有算子均可以放到BPU上 接着我们便可以开始配置我们的模型量化配置文件啦，我们开启了calibration_parameters校准数据类的preprocess_on功能来开启图片校准样本自动处