本文旨在将自己微调的模型部署到jetson-Orin中并利用官方的isaac-ROS库isaac-ros-rtdetr库运行该模型做实时检测一、创建环境1.训练环境a.创建conda基础环境conda create -n rtdetr python3.10 -yconda activate rtdetrb.安装pytorchpip install torch torchvision --index-url https://download.pytorch.org/whl/cu118c.验证是否安装成功python -c import torch; print(torch.__version__); print(torch.cuda.is_available())2.源代码a.克隆仓库git clone https://github.com/lyuwenyu/RT-DETR.gitb.进入对应目录cd RT-DETRcd rtdetrv2_pytorchc.安装仓库依赖pip install -r requirements.txt二、准备数据集准备COCO格式数据集目录结构官方如下或者很多标注工具如roboflow输出的coco格式如下注要确保json文件中类别必须连续从0开始roboflow输出的数据集会多一个奇怪的父类如下所示正确如下三、训练模型1.写数据配置文件configs/dataset/your_detection.ymltask: detection evaluator: type: CocoEvaluator iou_types: [bbox, ] num_classes: 3 remap_mscoco_category: False train_dataloader: type: DataLoader dataset: type: CocoDetection img_folder: ./dataset/bolt_washer_spacer/train/ ann_file: ./dataset/bolt_washer_spacer/train/_annotations.coco.json return_masks: False transforms: type: Compose ops: ~ shuffle: True num_workers: 2 drop_last: True collate_fn: type: BatchImageCollateFunction val_dataloader: type: DataLoader dataset: type: CocoDetection img_folder: ./dataset/bolt_washer_spacer/valid/ ann_file: ./dataset/bolt_washer_spacer/valid/_annotations.coco.json return_masks: False transforms: type: Compose ops: ~ shuffle: False num_workers: 2 drop_last: False collate_fn: type: BatchImageCollateFunction2.修改训练配置文件configs/rtdetrv2/your_train.yml__include__: [ ../dataset/your_detection.yml, ../runtime.yml, ./include/dataloader.yml, ./include/optimizer.yml, ./include/rtdetrv2_r50vd.yml, ] output_dir: ./output/your_file eval_spatial_size: [640, 640] PResNet: depth: 18 freeze_at: -1 freeze_norm: False pretrained: True HybridEncoder: in_channels: [128, 256, 512] hidden_dim: 256 expansion: 0.5 RTDETRTransformerv2: num_layers: 3 use_amp: True use_ema: True epoches: 72 clip_max_norm: 0.1 optimizer: type: AdamW params: - params: ^(?.*(?:norm|bn)).*$ weight_decay: 0. lr: 0.0001 betas: [0.9, 0.999] weight_decay: 0.0001 lr_scheduler: type: MultiStepLR milestones: [60] gamma: 0.1 lr_warmup_scheduler: type: LinearWarmup warmup_duration: 500 train_dataloader: dataset: transforms: ops: - {type: RandomPhotometricDistort, p: 0.5} - {type: RandomZoomOut, fill: 0} - {type: RandomIoUCrop, p: 0.8} - {type: SanitizeBoundingBoxes, min_size: 1} - {type: RandomHorizontalFlip} - {type: Resize, size: [640, 640]} - {type: SanitizeBoundingBoxes, min_size: 1} - {type: ConvertPILImage, dtype: float32, scale: True} - {type: ConvertBoxes, fmt: cxcywh, normalize: True} policy: name: stop_epoch epoch: 60 ops: [RandomPhotometricDistort, RandomZoomOut, RandomIoUCrop] collate_fn: scales: ~ total_batch_size: 1 num_workers: 2 val_dataloader: dataset: transforms: ops: - {type: Resize, size: [640, 640]} - {type: ConvertPILImage, dtype: float32, scale: True} total_batch_size: 1 num_workers: 23.训练命令python tools/train.py -c configs/rtdetrv2/your_train.yml -t path/to/pretrained.pth --use-amp4.测试集评估a.新建test数据配置__include__: [ ./your_train.yml, ] output_dir: ./output/your_test val_dataloader: dataset: img_folder: ./dataset/your_dataset/test/ ann_file: ./dataset/your_dataset/test/_annotations.coco.jsonb.测试命令python tools/train.py -c configs/rtdetrv2/your_test.yml -r output/best.pth --test-only -d cuda:0c.可视化新建批量可视化脚本Visualize RT-DETRv2 predictions on test images. import argparse import json import os import sys sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ..)) import torch import torch.nn as nn import torchvision.transforms as T from PIL import Image, ImageDraw, ImageFont from src.core import YAMLConfig CLASS_NAMES { 0: screw, 1: washer, 2: spacer, } COLORS { 0: #FF3838, 1: #FF9D97, 2: #FF701F, } def load_gt_boxes(ann_file, image_stem): with open(ann_file, encodingutf-8) as f: data json.load(f) image_id None for image in data[images]: if os.path.splitext(image[file_name])[0] image_stem: image_id image[id] break if image_id is None: return [] boxes [] for ann in data[annotations]: if ann[image_id] ! image_id: continue x, y, w, h ann[bbox] boxes.append({ label: ann[category_id], box: [x, y, x w, y h], }) return boxes def draw_boxes(image, boxes, modepred, score_mapNone): draw ImageDraw.Draw(image) for item in boxes: if mode pred: label int(item[label].item()) score float(item[score].item()) box [float(v) for v in item[box]] text f{CLASS_NAMES.get(label, label)} {score:.2f} color COLORS.get(label, #00FF00) else: label int(item[label]) box [float(v) for v in item[box]] text fGT {CLASS_NAMES.get(label, label)} color #00FF00 draw.rectangle(box, outlinecolor, width3) draw.text((box[0], max(0, box[1] - 12)), texttext, fillcolor) return image def build_model(cfg, checkpoint_path, device): checkpoint torch.load(checkpoint_path, map_locationcpu) if ema in checkpoint: state checkpoint[ema][module] else: state checkpoint[model] cfg.model.load_state_dict(state) class DeployModel(nn.Module): def __init__(self): super().__init__() self.model cfg.model.deploy() self.postprocessor cfg.postprocessor.deploy() def forward(self, images, orig_target_sizes): outputs self.model(images) return self.postprocessor(outputs, orig_target_sizes) return DeployModel().to(device).eval() def main(args): cfg YAMLConfig(args.config, resumeargs.resume) device torch.device(args.device if torch.cuda.is_available() else cpu) model build_model(cfg, args.resume, device) img_size cfg.yaml_cfg.get(eval_spatial_size, [640, 640]) if isinstance(img_size, list): img_size img_size[0] transforms T.Compose([ T.Resize((img_size, img_size)), T.ToTensor(), ]) os.makedirs(args.output_dir, exist_okTrue) image_files [ f for f in os.listdir(args.image_dir) if f.lower().endswith((.jpg, .jpeg, .png)) ] image_files.sort() if args.max_images 0: image_files image_files[:args.max_images] ann_file args.ann_file threshold args.threshold for image_name in image_files: image_path os.path.join(args.image_dir, image_name) im_pil Image.open(image_path).convert(RGB) w, h im_pil.size orig_size torch.tensor([[w, h]], devicedevice) im_data transforms(im_pil)[None].to(device) with torch.no_grad(): labels, boxes, scores model(im_data, orig_size) pred_items [] scr scores[0] for label, box, score in zip(labels[0], boxes[0], scr): if score.item() threshold: continue pred_items.append({label: label, box: box, score: score}) pred_image im_pil.copy() draw_boxes(pred_image, pred_items, modepred) gt_items load_gt_boxes(ann_file, os.path.splitext(image_name)[0]) gt_image im_pil.copy() draw_boxes(gt_image, gt_items, modegt) merged Image.new(RGB, (w * 2, h)) merged.paste(gt_image, (0, 0)) merged.paste(pred_image, (w, 0)) save_path os.path.join(args.output_dir, fvis_{image_name}) merged.save(save_path) print(fsaved: {save_path}) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(-c, --config, typestr, requiredTrue) parser.add_argument(-r, --resume, typestr, requiredTrue) parser.add_argument(--image-dir, typestr, requiredTrue) parser.add_argument(--ann-file, typestr, requiredTrue) parser.add_argument(--output-dir, typestr, requiredTrue) parser.add_argument(--device, typestr, defaultcuda:0) parser.add_argument(--threshold, typefloat, default0.5) parser.add_argument(--max-images, typeint, default0, help0 means all images) args parser.parse_args() main(args)执行脚本命令python tools/visualize_test.py ^-c configs/rtdetrv2/your_test.yml ^-r output/best.pth ^--image-dir dataset/your_dataset/test ^--ann-file dataset/your_dataset/test/_annotations.coco.json ^--output-dir output/your_test/vis ^--device cuda:0 --threshold 0.55.导出onnx模型命令python tools/export_onnx.py ^-c configs/rtdetrv2/your_train.yml ^-r output/your_file/best.pth ^-o output/your_file/best.onnx ^-s 640 --check四、部署模型1.输出onnx-sim模型pip install onnxsimonnxsim /path/to/your_model.onnx /path/to/your_model_simplified.onnx注rt-detr模型需要做simplify操作以兼容tensorRT架构不进行onnx-sim输出engine模型会报错2.输出engine模型在isaac_ros容器中终端运行trtexec \--onnxbest_sim.onnx \--saveEnginert_detr.engine \--fp16 \--memPoolSizeworkspace:4096 \--verbose五、封装进节点1.下载rt-detr官方节点库源码isaac_ros_rtdetr — isaac_ros_docs documentation2.更改源码由于源码是针对工业相机我用的是logic c270并没有时间戳以及输出的图像分辨率也不是正方形所以需要修改几处地方新增launch文件需要提前准备好v4l2节点可以看我之前apriltag的文章# SPDX-License-Identifier: Apache-2.0 import os from ament_index_python.packages import get_package_share_directory import launch from launch.actions import DeclareLaunchArgument from launch.conditions import IfCondition from launch.substitutions import LaunchConfiguration from launch_ros.actions import ComposableNodeContainer, Node from launch_ros.descriptions import ComposableNode MODEL_INPUT_SIZE 640 MODEL_NUM_CHANNELS 3 IMAGE_WIDTH 640 IMAGE_HEIGHT 480 WORKSPACE_ROOT os.environ.get(ISAAC_ROS_WS, /workspaces/isaac_ros-dev) DEFAULT_ENGINE_PATH os.path.join( get_package_share_directory(isaac_ros_rtdetr), models, rt_detr.engine) DEFAULT_CAMERA_INFO_URL ( file:// os.path.join(WORKSPACE_ROOT, src/v4l2_camera/config/camera_info.yaml) ) def generate_launch_description(): launch_args [ DeclareLaunchArgument( video_device, default_value/dev/video0, descriptionV4L2 device path for the Logitech C270), DeclareLaunchArgument( engine_file_path, default_valueDEFAULT_ENGINE_PATH, descriptionAbsolute path to the RT-DETR TensorRT engine file), DeclareLaunchArgument( camera_info_url, default_valueDEFAULT_CAMERA_INFO_URL, descriptionCamera calibration URL for camera_info_manager), DeclareLaunchArgument( confidence_threshold, default_value0.7, descriptionMinimum detection confidence score), DeclareLaunchArgument( force_engine_update, default_valueFalse, descriptionWhether TensorRT should rebuild the engine file), DeclareLaunchArgument( launch_visualizer, default_valueTrue, descriptionLaunch the RT-DETR detection visualizer node), ] engine_file_path LaunchConfiguration(engine_file_path) confidence_threshold LaunchConfiguration(confidence_threshold) force_engine_update LaunchConfiguration(force_engine_update) launch_visualizer LaunchConfiguration(launch_visualizer) v4l2_camera_node Node( packagev4l2_camera, executablev4l2_camera_node, namev4l2_camera, parameters[{ video_device: LaunchConfiguration(video_device), image_size: [IMAGE_WIDTH, IMAGE_HEIGHT], output_encoding: rgb8, camera_info_url: LaunchConfiguration(camera_info_url), }], outputscreen, ) resize_node ComposableNode( nameresize_node, packageisaac_ros_image_proc, pluginnvidia::isaac_ros::image_proc::ResizeNode, parameters[{ input_width: IMAGE_WIDTH, input_height: IMAGE_HEIGHT, output_width: MODEL_INPUT_SIZE, output_height: MODEL_INPUT_SIZE, # RT-DETRv2 training uses direct resize to 640x640 (stretch). keep_aspect_ratio: False, encoding_desired: rgb8, disable_padding: True, }], remappings[ (image, /image_raw), (camera_info, /camera_info), ], ) image_to_tensor_node ComposableNode( nameimage_to_tensor_node, packageisaac_ros_tensor_proc, pluginnvidia::isaac_ros::dnn_inference::ImageToTensorNode, parameters[{ scale: True, tensor_name: image, }], remappings[ (image, resize/image), (tensor, image_tensor), ], ) normalize_node ComposableNode( namenormalize_node, packageisaac_ros_tensor_proc, pluginnvidia::isaac_ros::dnn_inference::ImageTensorNormalizeNode, parameters[{ mean: [0.485, 0.456, 0.406], stddev: [0.229, 0.224, 0.225], input_tensor_name: image, output_tensor_name: image, }], remappings[ (tensor, image_tensor), (normalized_tensor, normalized_tensor), ], ) interleave_to_planar_node ComposableNode( nameinterleaved_to_planar_node, packageisaac_ros_tensor_proc, pluginnvidia::isaac_ros::dnn_inference::InterleavedToPlanarNode, parameters[{ input_tensor_shape: [MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, MODEL_NUM_CHANNELS], }], remappings[(interleaved_tensor, normalized_tensor)], ) reshape_node ComposableNode( namereshape_node, packageisaac_ros_tensor_proc, pluginnvidia::isaac_ros::dnn_inference::ReshapeNode, parameters[{ output_tensor_name: input_tensor, input_tensor_shape: [MODEL_NUM_CHANNELS, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE], output_tensor_shape: [1, MODEL_NUM_CHANNELS, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE], }], remappings[(tensor, planar_tensor)], ) rtdetr_preprocessor_node ComposableNode( namertdetr_preprocessor, packageisaac_ros_rtdetr, pluginnvidia::isaac_ros::rtdetr::RtDetrPreprocessorNode, parameters[{ image_width: IMAGE_WIDTH, image_height: IMAGE_HEIGHT, }], remappings[(encoded_tensor, reshaped_tensor)], ) tensor_rt_node ComposableNode( nametensor_rt, packageisaac_ros_tensor_rt, pluginnvidia::isaac_ros::dnn_inference::TensorRTNode, parameters[{ model_file_path: , engine_file_path: engine_file_path, output_binding_names: [labels, boxes, scores], output_tensor_names: [labels, boxes, scores], input_tensor_names: [images, orig_target_sizes], input_binding_names: [images, orig_target_sizes], verbose: False, force_engine_update: force_engine_update, }], ) rtdetr_decoder_node ComposableNode( namertdetr_decoder, packageisaac_ros_rtdetr, pluginnvidia::isaac_ros::rtdetr::RtDetrDecoderNode, parameters[{ confidence_threshold: confidence_threshold, }], ) rtdetr_container ComposableNodeContainer( namertdetr_container, namespace, packagerclcpp_components, executablecomponent_container_mt, composable_node_descriptions[ resize_node, image_to_tensor_node, normalize_node, interleave_to_planar_node, reshape_node, rtdetr_preprocessor_node, tensor_rt_node, rtdetr_decoder_node, ], outputscreen, ) visualizer_node Node( packageisaac_ros_rtdetr, executablertdetr_c270_visualizer.py, namertdetr_c270_visualizer, parameters[{ image_topic: /image_raw, detections_topic: /detections_output, min_confidence: confidence_threshold, max_detections: 20, }], conditionIfCondition(launch_visualizer), ) return launch.LaunchDescription( launch_args [v4l2_camera_node, rtdetr_container, visualizer_node])新增可视化脚本源码可视化脚本需要相机输出时间戳#!/usr/bin/env python3 import cv2 import cv_bridge import rclpy from rclpy.node import Node from sensor_msgs.msg import Image from vision_msgs.msg import Detection2DArray class RtDetrC270Visualizer(Node): def __init__(self): super().__init__(rtdetr_c270_visualizer) self.declare_parameter(image_topic, /image_raw) self.declare_parameter(detections_topic, /detections_output) self.declare_parameter(min_confidence, 0.7) self.declare_parameter(max_detections, 20) self._bridge cv_bridge.CvBridge() self._latest_detections None self._color (0, 255, 0) self._bbox_thickness 2 self._font cv2.FONT_HERSHEY_SIMPLEX self._font_scale 0.5 image_topic self.get_parameter(image_topic).get_parameter_value().string_value detections_topic ( self.get_parameter(detections_topic).get_parameter_value().string_value ) self._min_confidence ( self.get_parameter(min_confidence).get_parameter_value().double_value ) self._max_detections ( self.get_parameter(max_detections).get_parameter_value().integer_value ) self._processed_image_pub self.create_publisher( Image, rtdetr_processed_image, 10) self.create_subscription( Detection2DArray, detections_topic, self._detections_callback, 10) self.create_subscription( Image, image_topic, self._image_callback, 10) self.get_logger().info( fVisualizing {image_topic}, min_confidence{self._min_confidence}) def _detections_callback(self, msg): self._latest_detections msg def _filter_detections(self, detections): scored [] for detection in detections: if not detection.results: continue score detection.results[0].hypothesis.score if score self._min_confidence: continue scored.append((score, detection)) scored.sort(keylambda item: item[0], reverseTrue) return [detection for _, detection in scored[:self._max_detections]] def _image_callback(self, img_msg): try: cv2_img self._bridge.imgmsg_to_cv2(img_msg, desired_encodingrgb8) except cv_bridge.CvBridgeError as exc: self.get_logger().warn(fFailed to convert image: {exc}) return if self._latest_detections is not None: filtered self._filter_detections(self._latest_detections.detections) if filtered: self.get_logger().info( fDrawing {len(filtered)} detections, throttle_duration_sec2.0) for detection in filtered: center_x detection.bbox.center.position.x center_y detection.bbox.center.position.y width detection.bbox.size_x height detection.bbox.size_y try: min_pt ( round(center_x - (width / 2.0)), round(center_y - (height / 2.0)), ) max_pt ( round(center_x (width / 2.0)), round(center_y (height / 2.0)), ) cv2.rectangle( cv2_img, min_pt, max_pt, self._color, self._bbox_thickness) label detection.results[0].hypothesis.class_id score detection.results[0].hypothesis.score cv2.putText( cv2_img, f{label} {score:.2f}, min_pt, self._font, self._font_scale, self._color, 2) except ValueError: pass processed_img self._bridge.cv2_to_imgmsg(cv2_img, encodingrgb8) processed_img.header img_msg.header self._processed_image_pub.publish(processed_img) def main(): rclpy.init() node RtDetrC270Visualizer() rclpy.spin(node) rclpy.shutdown() if __name__ __main__: main()增加依赖:cmakelist.txt# SPDX-FileCopyrightText: NVIDIA CORPORATION AFFILIATES # Copyright (c) 2024 NVIDIA CORPORATION AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the License); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an AS IS BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # SPDX-License-Identifier: Apache-2.0 cmake_minimum_required(VERSION 3.22.1) project(isaac_ros_rtdetr LANGUAGES C CXX) if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES Clang) add_compile_options(-Wall -Wextra -Wpedantic) endif() find_package(ament_cmake_auto REQUIRED) ament_auto_find_build_dependencies() # RtDetrPreprocessorNode ament_auto_add_library(rtdetr_preprocessor_node SHARED src/rtdetr_preprocessor_node.cpp) rclcpp_components_register_nodes(rtdetr_preprocessor_node nvidia::isaac_ros::rtdetr::RtDetrPreprocessorNode) set(node_plugins ${node_plugins}nvidia::isaac_ros::rtdetr::RtDetrPreprocessorNode;$TARGET_FILE:rtdetr_preprocessor_node\n) # RtDetrDecoderNode ament_auto_add_library(rtdetr_decoder_node SHARED src/rtdetr_decoder_node.cpp) rclcpp_components_register_nodes(rtdetr_decoder_node nvidia::isaac_ros::rtdetr::RtDetrDecoderNode) set(node_plugins ${node_plugins}nvidia::isaac_ros::rtdetr::RtDetrDecoderNode;$TARGET_FILE:rtdetr_decoder_node\n) if(BUILD_TESTING) find_package(ament_lint_auto REQUIRED) ament_lint_auto_find_test_dependencies() # The FindPythonInterp and FindPythonLibs modules are removed if(POLICY CMP0148) cmake_policy(SET CMP0148 OLD) endif() find_package(launch_testing_ament_cmake REQUIRED) add_launch_test(test/isaac_ros_rtdetr_pol_test.py TIMEOUT 600) endif() # Visualizer python scripts ament_python_install_package(${PROJECT_NAME}) install(PROGRAMS scripts/isaac_ros_rtdetr_visualizer.py scripts/rtdetr_c270_visualizer.py scripts/rtdetr_sam_visualizer.py DESTINATION lib/${PROJECT_NAME} ) install(DIRECTORY models DESTINATION share/${PROJECT_NAME} ) # Embed versioning information into installed files ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) include(${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake) generate_version_info(${PROJECT_NAME}) ament_auto_package(INSTALL_TO_SHARE launch)package.xml?xml version1.0? !-- SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION AFFILIATES. All rights reserved. Licensed under the Apache License, Version 2.0 (the License); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an AS IS BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. SPDX-License-Identifier: Apache-2.0 -- ?xml-model hrefhttp://download.ros.org/schema/package_format3.xsd schematypenshttp://www.w3.org/2001/XMLSchema? package format3 nameisaac_ros_rtdetr/name version3.2.5/version descriptionRT-DETR model processing/description maintainer emailisaac-ros-maintainersnvidia.comIsaac ROS Maintainers/maintainer licenseApache-2.0/license url typewebsitehttps://developer.nvidia.com/isaac-ros//url authorJaiveer Singh/author buildtool_dependament_cmake_auto/buildtool_depend dependrclcpp/depend dependrclcpp_components/depend dependvision_msgs/depend dependisaac_ros_nitros/depend dependisaac_ros_nitros_tensor_list_type/depend dependisaac_ros_managed_nitros/depend dependisaac_ros_tensor_list_interfaces/depend build_dependisaac_ros_common/build_depend exec_dependisaac_ros_dnn_image_encoder/exec_depend exec_dependisaac_ros_image_proc/exec_depend exec_dependisaac_ros_tensor_proc/exec_depend exec_dependisaac_ros_tensor_rt/exec_depend exec_dependisaac_ros_segment_anything/exec_depend exec_dependisaac_ros_triton/exec_depend exec_dependv4l2_camera/exec_depend exec_dependrclpy/exec_depend exec_dependcv_bridge/exec_depend exec_dependsensor_msgs/exec_depend exec_dependgxf_isaac_optimizer/exec_depend test_dependament_lint_auto/test_depend test_dependament_lint_common/test_depend test_dependisaac_ros_test/test_depend export build_typeament_cmake/build_type /export /package编译源码colcon build --packages-select isaac_ros_rtdetr --symlink-installsource install/setup.bash运行启动文件ros2 launch isaac_ros_rtdetr rtdetr_c270.launch.py可以通过rviz2或者rqt订阅节点话题六、检测效果
