diff --git a/CMakeLists.txt b/CMakeLists.txt index a3a5244..20afd17 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -193,6 +193,7 @@ set(COMMON_LIBS if(ODIN_TARGET_OBSERVATION_READY) add_library(target_observation_processing STATIC src/target_observation_processing.cpp + src/reid_trt_extractor.cpp ${TARGET_PREDICTION_THIRDPARTY_DIR}/YOLOs-CPP-TensorRT/include/yolos/core/cuda_preprocessing.cu ) target_include_directories(target_observation_processing PUBLIC diff --git a/README_track.md b/README_track.md new file mode 100644 index 0000000..feab7c0 --- /dev/null +++ b/README_track.md @@ -0,0 +1,3 @@ +# yolo+ track使用说明 +选取第一帧中心的人物作为target,持续建立gallery并使用bytetrack跟踪 +如果丢了使用lazy reid,基于gallery进行匹配 diff --git a/config/control_command.yaml b/config/control_command.yaml index 924d4a1..bfb7a28 100644 --- a/config/control_command.yaml +++ b/config/control_command.yaml @@ -89,6 +89,21 @@ register_keys: # draw target observation result on overlay image and print debug info # 0: off; 1: on debug: 1 + # lazy ReID only runs after the current ByteTrack target is lost + # first valid frame still bootstraps by choosing the target closest to image center + target_reid_enable: 1 + target_reid_engine: "/home/hjy/HJY/Codes/AMP_ws/TargetPrediction/models/target_reid_osnet_x0_25_dukemtmcreid.trt" + target_reid_match_threshold: 0.65 + target_reid_gap_threshold: 0.10 + target_reid_min_crop_area: 2000.0 + target_reid_max_crop_aspect_ratio: 0.90 + target_reid_feature_update_interval: 5 + target_reid_lost_timeout_frames: 150 + target_reid_gallery_size: 10 + target_reid_input_height: 256 + target_reid_input_width: 128 + target_reid_feature_dim: 512 + target_reid_max_batch_size: 8 # record rgb, odometry, and slam cloud data as proprietary olx format for further processing in MindCloud(TM) software. # save path: ws/src/odin_ros_driver/recorddata/{record_start_time}/ diff --git a/config/odin_ros2.rviz b/config/odin_ros2.rviz index 0e4da57..ea02fd1 100644 --- a/config/odin_ros2.rviz +++ b/config/odin_ros2.rviz @@ -14,8 +14,11 @@ Panels: - /Prediction1 - /Prediction1/Image1 - /Prediction1/Marker2 + - /Planning1 + - /Planning1/GridMap1 + - /Planning1/path1 Splitter Ratio: 0.5 - Tree Height: 865 + Tree Height: 892 - Class: rviz_common/Selection Name: Selection - Class: rviz_common/Tool Properties @@ -426,7 +429,7 @@ Visualization Manager: Value: /target/pred_trajectory Value: true - Class: rviz_default_plugins/Image - Enabled: false + Enabled: true Max Value: 1 Median window: 5 Min Value: 0 @@ -438,7 +441,7 @@ Visualization Manager: History Policy: Keep Last Reliability Policy: Reliable Value: /target/pred_image - Value: false + Value: true - Alpha: 1 Class: rviz_default_plugins/PointStamped Color: 224; 27; 36 @@ -473,7 +476,7 @@ Visualization Manager: Enabled: true Name: Marker Namespaces: - traj_prediction: true + {} Topic: Depth: 5 Durability Policy: Volatile @@ -484,6 +487,49 @@ Visualization Manager: Value: true Enabled: true Name: Prediction + - Class: rviz_common/Group + Displays: + - Alpha: 1 + Autocompute Intensity Bounds: true + Class: grid_map_rviz_plugin/GridMap + Color: 200; 200; 200 + Color Layer: traversability + Color Transformer: GridMapLayer + Enabled: true + Height Layer: elevation + Height Transformer: GridMapLayer + History Length: 1 + Invert Rainbow: true + Max Color: 255; 255; 255 + Max Intensity: 10 + Min Color: 0; 0; 0 + Min Intensity: 0 + Name: GridMap + Show Grid Lines: true + Topic: + Depth: 5 + Durability Policy: Volatile + Filter size: 10 + History Policy: Keep Last + Reliability Policy: Reliable + Value: /local_grid_map + Use Rainbow: true + Value: true + - Class: rviz_default_plugins/MarkerArray + Enabled: true + Name: path + Namespaces: + future_path: true + future_yaw: true + Topic: + Depth: 5 + Durability Policy: Volatile + History Policy: Keep Last + Reliability Policy: Reliable + Value: /odin1/odometry_rts_future_markers + Value: true + Enabled: true + Name: Planning Enabled: true Global Options: Background Color: 48; 48; 48 @@ -530,37 +576,37 @@ Visualization Manager: Views: Current: Class: rviz_default_plugins/ThirdPersonFollower - Distance: 31.34433937072754 + Distance: 8.066532135009766 Enable Stereo Rendering: Stereo Eye Separation: 0.05999999865889549 Stereo Focal Distance: 1 Swap Stereo Eyes: false Value: false Focal Point: - X: -29.93891143798828 - Y: -29.894529342651367 - Z: -3.1263247365131974e-05 + X: -0.3127555847167969 + Y: 0.25075721740722656 + Z: 0 Focal Shape Fixed Size: false Focal Shape Size: 0.05000000074505806 Invert Z Axis: false Name: Current View Near Clip Distance: 0.009999999776482582 - Pitch: 0.545399010181427 + Pitch: 0.9853979349136353 Target Frame: odin1_base_link Value: ThirdPersonFollower (rviz_default_plugins) - Yaw: 0.24540430307388306 + Yaw: 2.8103888034820557 Saved: ~ Window Geometry: Displays: collapsed: false - Height: 1016 + Height: 1043 Hide Left Dock: false Hide Right Dock: false Image: collapsed: false Image_undistort: collapsed: false - QMainWindow State: 000000ff00000000fd0000000400000000000001d40000039efc020000000dfb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb0000000a0049006d006100670065000000003d000002040000002800fffffffb000000100044006900730070006c006100790073010000003d0000039e000000c900fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000000a0049006d00610067006500000001f7000001700000002800fffffffb0000002200640065006e00730065005f00640065007000740068005f0069006d006100670065000000023e0000007a0000002800fffffffb0000001e0049006d006100670065005f0075006e0064006900730074006f00720074000000028f0000014c0000002800fffffffb0000002a0063006c006f007500640073006c0061006d005f0072006500700072006f006a006500630074006500640000000310000000cb0000002800ffffff000000010000010f0000039efc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073010000003d0000039e000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000004420000003efc0100000002fb0000000800540069006d00650100000000000004420000000000000000fb0000000800540069006d006501000000000000045000000000000000000000044b0000039e00000004000000040000000800000008fc0000000100000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 + QMainWindow State: 000000ff00000000fd0000000400000000000001d4000003b9fc020000000dfb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb0000000a0049006d006100670065000000003d000002040000002800fffffffb000000100044006900730070006c006100790073010000003d000003b9000000c900fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000000a0049006d006100670065030000082a00000112000002bd00000249fb0000002200640065006e00730065005f00640065007000740068005f0069006d006100670065000000023e0000007a0000002800fffffffb0000001e0049006d006100670065005f0075006e0064006900730074006f00720074000000028f0000014c0000002800fffffffb0000002a0063006c006f007500640073006c0061006d005f0072006500700072006f006a006500630074006500640000000310000000cb0000002800ffffff000000010000010f000003b9fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073010000003d000003b9000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000004420000003efc0100000002fb0000000800540069006d00650100000000000004420000000000000000fb0000000800540069006d006501000000000000045000000000000000000000044b000003b900000004000000040000000800000008fc0000000100000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 Selection: collapsed: false Tool Properties: @@ -569,7 +615,7 @@ Window Geometry: collapsed: false Width: 1850 X: 70 - Y: 27 + Y: 0 cloudslam_reprojected: collapsed: false dense_depth_image: diff --git a/include/cloud_reprojection_ros_node.hpp b/include/cloud_reprojection_ros_node.hpp index 9dd792d..bbd3c47 100644 --- a/include/cloud_reprojection_ros_node.hpp +++ b/include/cloud_reprojection_ros_node.hpp @@ -88,6 +88,7 @@ private: std::string sync_wiwc_topic_; std::string sync_image_topic_; std::string sync_overlay_image_topic_; + std::string sync_detection_debug_image_topic_; std::string sync_target_observation_topic_; std::string sync_target_pos_cam_topic_; std::string sync_target_pos_world_topic_; @@ -99,6 +100,7 @@ private: rclcpp::Publisher::SharedPtr sync_image_pub_; rclcpp::Publisher::SharedPtr overlay_compressed_pub_; // optional if send_overlay_ + rclcpp::Publisher::SharedPtr detection_debug_compressed_pub_; // debug detections/tracks rclcpp::Publisher::SharedPtr combined_pub_; // optional if publish_combined_compressed_ rclcpp::Publisher::SharedPtr target_observation_pub_; rclcpp::Publisher::SharedPtr target_pos_cam_pub_; diff --git a/include/reid_trt_extractor.hpp b/include/reid_trt_extractor.hpp new file mode 100644 index 0000000..06d9a96 --- /dev/null +++ b/include/reid_trt_extractor.hpp @@ -0,0 +1,65 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace odin_ros_driver { + +using ReIDFeature = std::vector; + +struct ReIDTensorRTConfig { + std::string engine_path; + int input_height = 256; + int input_width = 128; + int feature_dim = 512; + int max_batch_size = 8; +}; + +class TrackGallery { +public: + explicit TrackGallery(size_t max_features = 10); + + void clear(); + void add(const ReIDFeature& feature, int64_t frame_index); + + bool empty() const { return features_.empty(); } + size_t size() const { return features_.size(); } + int64_t last_update_frame() const { return last_update_frame_; } + + float best_similarity(const ReIDFeature& query) const; + +private: + std::deque features_; + size_t max_features_ = 10; + int64_t last_update_frame_ = -1; +}; + +class ReIDTensorRTExtractor { +public: + ReIDTensorRTExtractor(); + ~ReIDTensorRTExtractor(); + + ReIDTensorRTExtractor(const ReIDTensorRTExtractor&) = delete; + ReIDTensorRTExtractor& operator=(const ReIDTensorRTExtractor&) = delete; + + void initialize(const ReIDTensorRTConfig& config); + bool initialized() const { return initialized_; } + int feature_dim() const; + + std::vector extract( + const cv::Mat& frame_bgr, + const std::vector& boxes) const; + +private: + struct Impl; + std::unique_ptr impl_; + ReIDTensorRTConfig config_; + bool initialized_ = false; +}; + +} // namespace odin_ros_driver diff --git a/include/target_observation_processing.hpp b/include/target_observation_processing.hpp index b79672f..96f6611 100644 --- a/include/target_observation_processing.hpp +++ b/include/target_observation_processing.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -11,6 +12,7 @@ #include #include "cloud_reprojector.hpp" +#include "reid_trt_extractor.hpp" #include "motcpp/trackers/bytetrack.hpp" #include "yolos/tasks/pose.hpp" @@ -24,12 +26,26 @@ struct TargetObservationConfig { float min_depth = 0.5f; float max_depth = 12.0f; float search_radius_px = 25.0f; + bool reid_enabled = false; + std::string reid_engine_path; + float reid_match_threshold = 0.65f; + float reid_gap_threshold = 0.10f; + float reid_min_crop_area = 2000.0f; + float reid_max_crop_aspect_ratio = 0.9f; + int reid_feature_update_interval = 5; + int reid_lost_timeout_frames = 150; + int reid_gallery_size = 10; + int reid_input_height = 256; + int reid_input_width = 128; + int reid_feature_dim = 512; + int reid_max_batch_size = 8; bool debug = false; }; struct TargetObservation { bool valid = false; int track_id = -1; + int raw_track_id = -1; int detection_index = -1; float confidence = 0.0f; float depth = -1.0f; @@ -42,14 +58,22 @@ struct TargetObservation { struct TargetObservationDebugInfo { int current_target_id_before = -1; + int current_raw_track_id_before = -1; int poses_count = 0; int tracks_count = 0; int selected_track_id = -1; + int selected_raw_track_id = -1; int detection_index = -1; int projected_cloud_points = 0; int depth_sample_count = 0; + int gallery_size = 0; + int lost_frames = 0; bool found_existing_target = false; bool target_selected = false; + bool selected_from_center_bootstrap = false; + bool reid_attempted = false; + bool recovered_by_reid = false; + float reid_similarity = -1.0f; float yolo_ms = 0.0f; float mot_ms = 0.0f; float depth_ms = 0.0f; @@ -82,6 +106,7 @@ private: const std::vector& poses) const; TargetObservation select_target( + const cv::Mat& camera_bgr, const std::vector& poses, const Eigen::MatrixXf& tracks, int image_width, @@ -96,10 +121,27 @@ private: const CloudReprojector::OdomPose& odom_pose, TargetObservationDebugInfo* debug_info) const; + int recover_target_with_reid( + const cv::Mat& camera_bgr, + const Eigen::MatrixXf& tracks, + float* best_similarity) const; + void update_target_gallery( + const cv::Mat& camera_bgr, + const TargetObservation& target, + bool force_update); + + bool should_update_gallery(bool force_update) const; + bool is_good_reid_crop(const cv::Rect& bbox, const cv::Size& image_size) const; + TargetObservationConfig config_; std::unique_ptr yolo_; std::unique_ptr tracker_; + std::unique_ptr reid_extractor_; + TrackGallery target_gallery_; int current_target_id_ = -1; + int current_raw_track_id_ = -1; + int64_t frame_index_ = 0; + int lost_track_frames_ = 0; bool initialized_ = false; }; diff --git a/launch_ROS2/odin1_driver_ros2.launch.py b/launch_ROS2/odin1_driver_ros2.launch.py new file mode 100644 index 0000000..430518e --- /dev/null +++ b/launch_ROS2/odin1_driver_ros2.launch.py @@ -0,0 +1,33 @@ +import os + +from ament_index_python.packages import get_package_share_directory +from launch import LaunchDescription +from launch.actions import DeclareLaunchArgument, SetEnvironmentVariable +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import Node + + +def generate_launch_description(): + package_dir = get_package_share_directory('odin_ros_driver') + + config_file_arg = DeclareLaunchArgument( + 'config_file', + default_value=os.path.join(package_dir, 'config', 'control_command.yaml'), + description='Path to the control config YAML file' + ) + + host_sdk_node = Node( + package='odin_ros_driver', + executable='host_sdk_sample', + name='host_sdk_sample', + output='screen', + parameters=[{ + 'config_file': LaunchConfiguration('config_file') + }] + ) + + ld = LaunchDescription() + ld.add_action(SetEnvironmentVariable('RCUTILS_COLORIZED_OUTPUT', '1')) + ld.add_action(config_file_arg) + ld.add_action(host_sdk_node) + return ld diff --git a/launch_ROS2/odin1_reprojection_ros2.launch.py b/launch_ROS2/odin1_reprojection_ros2.launch.py new file mode 100644 index 0000000..7b6d787 --- /dev/null +++ b/launch_ROS2/odin1_reprojection_ros2.launch.py @@ -0,0 +1,52 @@ +import os +import yaml + +from ament_index_python.packages import get_package_share_directory +from launch import LaunchDescription +from launch.actions import DeclareLaunchArgument, OpaqueFunction, SetEnvironmentVariable +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import Node + + +def create_nodes(context): + package_dir = get_package_share_directory('odin_ros_driver') + config_file = LaunchConfiguration('config_file').perform(context) + + with open(config_file, 'r', encoding='utf-8') as stream: + reprojection_params = yaml.safe_load(stream) or {} + + reprojection_params['calib_file_path'] = os.path.join(package_dir, 'config', 'calib.yaml') + + pcd2depth_node = Node( + package='odin_ros_driver', + executable='pcd2depth_ros2_node', + name='pcd2depth_ros2_node', + output='screen', + parameters=[reprojection_params] + ) + + cloud_reprojection_node = Node( + package='odin_ros_driver', + executable='cloud_reprojection_ros2_node', + name='cloud_reprojection_ros2_node', + output='screen', + parameters=[reprojection_params] + ) + + return [pcd2depth_node, cloud_reprojection_node] + + +def generate_launch_description(): + package_dir = get_package_share_directory('odin_ros_driver') + + config_file_arg = DeclareLaunchArgument( + 'config_file', + default_value=os.path.join(package_dir, 'config', 'control_command.yaml'), + description='Path to the control config YAML file' + ) + + ld = LaunchDescription() + ld.add_action(SetEnvironmentVariable('RCUTILS_COLORIZED_OUTPUT', '1')) + ld.add_action(config_file_arg) + ld.add_action(OpaqueFunction(function=lambda context: create_nodes(context))) + return ld diff --git a/launch_ROS2/odin1_ros2.launch.py b/launch_ROS2/odin1_ros2.launch.py index b582556..70ffd67 100644 --- a/launch_ROS2/odin1_ros2.launch.py +++ b/launch_ROS2/odin1_ros2.launch.py @@ -1,89 +1,47 @@ - -# USAGE: ros2 launch odin_ros_driver odin1_ros2.launch.py +# USAGE: +# ros2 launch odin_ros_driver odin1_driver_ros2.launch.py +# ros2 launch odin_ros_driver odin1_reprojection_ros2.launch.py +# ros2 launch odin_ros_driver odin1_ros2.launch.py import os -import yaml + from ament_index_python.packages import get_package_share_directory from launch import LaunchDescription -from launch.actions import DeclareLaunchArgument, SetEnvironmentVariable +from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription, SetEnvironmentVariable +from launch.launch_description_sources import PythonLaunchDescriptionSource from launch.substitutions import LaunchConfiguration -from launch_ros.actions import Node + def generate_launch_description(): - # Get package directory package_dir = get_package_share_directory('odin_ros_driver') - - # Declare configuration parameter + launch_dir = os.path.join(package_dir, 'launch') + config_file_arg = DeclareLaunchArgument( 'config_file', default_value=os.path.join(package_dir, 'config', 'control_command.yaml'), description='Path to the control config YAML file' ) - - # Add RViz2 configuration file parameter - rviz_config_arg = DeclareLaunchArgument( - 'rviz_config', - default_value=os.path.join(package_dir, 'config', 'odin_ros2.rviz'), - description='Path to RViz2 config file' - ) - - # Create main node - host_sdk_node = Node( - package='odin_ros_driver', - executable='host_sdk_sample', - name='host_sdk_sample', - output='screen', - # arguments=['--ros-args', '--log-level', 'debug'], - parameters=[{ + + driver_launch = IncludeLaunchDescription( + PythonLaunchDescriptionSource( + os.path.join(launch_dir, 'odin1_driver_ros2.launch.py') + ), + launch_arguments={ 'config_file': LaunchConfiguration('config_file') - }] + }.items() ) - pcd2depth_config_path = os.path.join(package_dir, 'config', 'control_command.yaml') - with open(pcd2depth_config_path, 'r') as f: - pcd2depth_params = yaml.safe_load(f) - pcd2depth_calib_path = os.path.join(package_dir, 'config', 'calib.yaml') - pcd2depth_params['calib_file_path'] = pcd2depth_calib_path - pcd2depth_node = Node( - package='odin_ros_driver', - executable='pcd2depth_ros2_node', - name='pcd2depth_ros2_node', - output='screen', - parameters=[pcd2depth_params] + reprojection_launch = IncludeLaunchDescription( + PythonLaunchDescriptionSource( + os.path.join(launch_dir, 'odin1_reprojection_ros2.launch.py') + ), + launch_arguments={ + 'config_file': LaunchConfiguration('config_file') + }.items() ) - # Cloud reprojection node - reprojection_config_path = os.path.join(package_dir, 'config', 'control_command.yaml') - with open(reprojection_config_path, 'r') as f: - reprojection_params = yaml.safe_load(f) - reprojection_calib_path = os.path.join(package_dir, 'config', 'calib.yaml') - reprojection_params['calib_file_path'] = reprojection_calib_path - cloud_reprojection_node = Node( - package='odin_ros_driver', - executable='cloud_reprojection_ros2_node', - name='cloud_reprojection_ros2_node', - output='screen', - parameters=[reprojection_params] - ) - - # Combined jpeg is published by cloud_reprojection_ros2_node (left=depth z, right=camera BGR) - - # Create RViz2 node - loads specified configuration file - rviz_node = Node( - package='rviz2', - executable='rviz2', - name='rviz2', - output='screen', - arguments=['-d', LaunchConfiguration('rviz_config')] - ) - - # Create launch description ld = LaunchDescription() ld.add_action(SetEnvironmentVariable('RCUTILS_COLORIZED_OUTPUT', '1')) ld.add_action(config_file_arg) - ld.add_action(rviz_config_arg) # Add RViz configuration argument - ld.add_action(host_sdk_node) - ld.add_action(pcd2depth_node) - ld.add_action(cloud_reprojection_node) - # ld.add_action(rviz_node) # Add RViz node - + ld.add_action(driver_launch) + ld.add_action(reprojection_launch) return ld diff --git a/src/cloud_reprojection_ros.cpp b/src/cloud_reprojection_ros.cpp index bc87468..511a3b9 100644 --- a/src/cloud_reprojection_ros.cpp +++ b/src/cloud_reprojection_ros.cpp @@ -131,6 +131,7 @@ CloudReprojectionRosNode::CloudReprojectionRosNode(const rclcpp::NodeOptions& op << "\n camera_image_topic: " << camera_image_topic_ << "\n sync_* topics under: " << sync_topic_prefix_ << "\n overlay_compressed_topic: " << sync_overlay_image_topic_ + << "\n detection_debug_topic: " << sync_detection_debug_image_topic_ << "\n combined_compressed_topic: " << combined_compressed_topic_ << "\n publish_combined_compressed: " << (publish_combined_compressed_ ? "on" : "off") << "\n send_overlay: " << (send_overlay_ ? "on" : "off") @@ -163,6 +164,10 @@ CloudReprojectionRosNode::CloudReprojectionRosNode(const rclcpp::NodeOptions& op overlay_compressed_pub_ = this->create_publisher(sync_overlay_image_topic_, 10); } + if (enable_target_observation_ && debug_target_observation_) { + detection_debug_compressed_pub_ = + this->create_publisher(sync_detection_debug_image_topic_, 10); + } if (publish_combined_compressed_) { combined_pub_ = this->create_publisher(combined_compressed_topic_, 10); } @@ -216,6 +221,19 @@ void CloudReprojectionRosNode::loadParameters() this->declare_parameter("register_keys.target_min_depth", 0.5); this->declare_parameter("register_keys.target_max_depth", 12.0); this->declare_parameter("register_keys.target_search_radius_px", 25.0); + this->declare_parameter("register_keys.target_reid_enable", 0); + this->declare_parameter("register_keys.target_reid_engine", ""); + this->declare_parameter("register_keys.target_reid_match_threshold", 0.65); + this->declare_parameter("register_keys.target_reid_gap_threshold", 0.10); + this->declare_parameter("register_keys.target_reid_min_crop_area", 2000.0); + this->declare_parameter("register_keys.target_reid_max_crop_aspect_ratio", 0.90); + this->declare_parameter("register_keys.target_reid_feature_update_interval", 5); + this->declare_parameter("register_keys.target_reid_lost_timeout_frames", 150); + this->declare_parameter("register_keys.target_reid_gallery_size", 10); + this->declare_parameter("register_keys.target_reid_input_height", 256); + this->declare_parameter("register_keys.target_reid_input_width", 128); + this->declare_parameter("register_keys.target_reid_feature_dim", 512); + this->declare_parameter("register_keys.target_reid_max_batch_size", 8); #endif cloud_slam_topic_ = this->get_parameter("cloud_slam_topic").as_string(); @@ -238,6 +256,7 @@ void CloudReprojectionRosNode::loadParameters() sync_wiwc_topic_ = sync_topic_prefix_ + "/wiwc"; sync_image_topic_ = sync_topic_prefix_ + "/image"; sync_overlay_image_topic_ = sync_topic_prefix_ + "/overlay_img/compressed"; + sync_detection_debug_image_topic_ = sync_topic_prefix_ + "/detection_img_debug/compressed"; sync_target_observation_topic_ = sync_topic_prefix_ + "/target_observation"; sync_target_pos_cam_topic_ = sync_topic_prefix_ + "/target_pos_cam"; sync_target_pos_world_topic_ = sync_topic_prefix_ + "/target_pos_world"; @@ -339,14 +358,42 @@ void CloudReprojectionRosNode::loadParameters() this->get_parameter("register_keys.target_max_depth").as_double()); target_config.search_radius_px = static_cast( this->get_parameter("register_keys.target_search_radius_px").as_double()); + target_config.reid_enabled = + (this->get_parameter("register_keys.target_reid_enable").as_int() != 0); + target_config.reid_engine_path = + this->get_parameter("register_keys.target_reid_engine").as_string(); + target_config.reid_match_threshold = static_cast( + this->get_parameter("register_keys.target_reid_match_threshold").as_double()); + target_config.reid_gap_threshold = static_cast( + this->get_parameter("register_keys.target_reid_gap_threshold").as_double()); + target_config.reid_min_crop_area = static_cast( + this->get_parameter("register_keys.target_reid_min_crop_area").as_double()); + target_config.reid_max_crop_aspect_ratio = static_cast( + this->get_parameter("register_keys.target_reid_max_crop_aspect_ratio").as_double()); + target_config.reid_feature_update_interval = + this->get_parameter("register_keys.target_reid_feature_update_interval").as_int(); + target_config.reid_lost_timeout_frames = + this->get_parameter("register_keys.target_reid_lost_timeout_frames").as_int(); + target_config.reid_gallery_size = + this->get_parameter("register_keys.target_reid_gallery_size").as_int(); + target_config.reid_input_height = + this->get_parameter("register_keys.target_reid_input_height").as_int(); + target_config.reid_input_width = + this->get_parameter("register_keys.target_reid_input_width").as_int(); + target_config.reid_feature_dim = + this->get_parameter("register_keys.target_reid_feature_dim").as_int(); + target_config.reid_max_batch_size = + this->get_parameter("register_keys.target_reid_max_batch_size").as_int(); target_config.debug = debug_target_observation_; target_observation_processor_ = std::make_unique(); target_observation_processor_->initialize(target_config); RCLCPP_INFO( this->get_logger(), - "Target observation enabled | engine=%s", - target_config.yolo_engine_path.c_str()); + "Target observation enabled | yolo_engine=%s | reid=%s | reid_engine=%s", + target_config.yolo_engine_path.c_str(), + target_config.reid_enabled ? "on" : "off", + target_config.reid_engine_path.empty() ? "" : target_config.reid_engine_path.c_str()); RCLCPP_INFO( this->get_logger(), "Target observation processor initialized successfully | labels=%s", @@ -354,14 +401,21 @@ void CloudReprojectionRosNode::loadParameters() if (debug_target_observation_) { RCLCPP_INFO( this->get_logger(), - "Target observation debug enabled | yolo_engine=%s | yolo_labels=%s | yolo_conf=%.3f | yolo_nms=%.3f | min_depth=%.2f | max_depth=%.2f | search_radius_px=%.1f", + "Target observation debug enabled | yolo_engine=%s | yolo_labels=%s | yolo_conf=%.3f | yolo_nms=%.3f | min_depth=%.2f | max_depth=%.2f | search_radius_px=%.1f | reid=%s | reid_engine=%s | reid_match=%.3f | reid_gap=%.3f | reid_gallery=%d | reid_update_interval=%d | reid_timeout=%d", target_config.yolo_engine_path.c_str(), target_config.yolo_labels_path.empty() ? "" : target_config.yolo_labels_path.c_str(), target_config.yolo_conf, target_config.yolo_nms, target_config.min_depth, target_config.max_depth, - target_config.search_radius_px); + target_config.search_radius_px, + target_config.reid_enabled ? "on" : "off", + target_config.reid_engine_path.empty() ? "" : target_config.reid_engine_path.c_str(), + target_config.reid_match_threshold, + target_config.reid_gap_threshold, + target_config.reid_gallery_size, + target_config.reid_feature_update_interval, + target_config.reid_lost_timeout_frames); } } catch (const std::exception& e) { enable_target_observation_ = false; @@ -488,9 +542,33 @@ void CloudReprojectionRosNode::syncCallback( cloud_cam, *reprojector_, odom_pose, - debug_target_observation_ ? &target_debug : nullptr); + &target_debug); const auto target_end = std::chrono::steady_clock::now(); + if (target_debug.reid_attempted) { + if (target_debug.recovered_by_reid) { + RCLCPP_WARN( + this->get_logger(), + "\033[1;33m[LazyReID]\033[0m \033[1;32mRECOVERED\033[0m stable_id=%d raw_id=%d det_ind=%d sim=%.3f gallery=%d lost=%d", + target_debug.current_target_id_before, + target_debug.selected_raw_track_id, + target_debug.detection_index, + target_debug.reid_similarity, + target_debug.gallery_size, + target_debug.lost_frames); + } else { + RCLCPP_WARN( + this->get_logger(), + "\033[1;33m[LazyReID]\033[0m \033[1;31mNO_MATCH\033[0m prev_id=%d sim=%.3f gallery=%d lost=%d tracked=%d detections=%d", + target_debug.current_target_id_before, + target_debug.reid_similarity, + target_debug.gallery_size, + target_debug.lost_frames, + target_debug.tracks_count, + target_debug.poses_count); + } + } + if (debug_target_observation_) { if (!target_observation.valid) { if (target_debug.poses_count == 0) { @@ -517,11 +595,18 @@ void CloudReprojectionRosNode::syncCallback( this->get_logger(), *this->get_clock(), 1000, - "Target observation | detections=%d tracked=%d selected_id=%d det_ind=%d cloud_pts=%d depth_samples=%d | yolo=%.2f ms | mot=%.2f ms | depth=%.2f ms | total=%.2f ms", + "Target observation | detections=%d tracked=%d selected_id=%d raw_id=%d det_ind=%d center_fallback=%s reid_attempted=%s reid_recovered=%s reid_sim=%.3f gallery=%d lost=%d cloud_pts=%d depth_samples=%d | yolo=%.2f ms | mot=%.2f ms | depth=%.2f ms | total=%.2f ms", target_debug.poses_count, target_debug.tracks_count, target_debug.selected_track_id, + target_debug.selected_raw_track_id, target_debug.detection_index, + target_debug.selected_from_center_bootstrap ? "yes" : "no", + target_debug.reid_attempted ? "yes" : "no", + target_debug.recovered_by_reid ? "yes" : "no", + target_debug.reid_similarity, + target_debug.gallery_size, + target_debug.lost_frames, target_debug.projected_cloud_points, target_debug.depth_sample_count, target_debug.yolo_ms, @@ -534,7 +619,7 @@ void CloudReprojectionRosNode::syncCallback( this->get_logger(), *this->get_clock(), 500, - "Target observation | stamp=%u.%u size=%dx%d cloud=%zu detections=%d tracked=%d current_target_id=%d selected_id=%d det_ind=%d reused=%s bbox=%s depth=%.3f conf=%.3f depth_conf=%.3f pos_cam=%s pos_world=%s cloud_pts=%d depth_samples=%d | yolo=%.2f ms | mot=%.2f ms | depth=%.2f ms | total=%.2f ms | node_total=%.2f ms", + "Target observation | stamp=%u.%u size=%dx%d cloud=%zu detections=%d tracked=%d current_target_id=%d current_raw_id=%d selected_id=%d selected_raw_id=%d det_ind=%d reused=%s center_fallback=%s reid_attempted=%s reid_recovered=%s reid_sim=%.3f gallery=%d lost=%d bbox=%s depth=%.3f conf=%.3f depth_conf=%.3f pos_cam=%s pos_world=%s cloud_pts=%d depth_samples=%d | yolo=%.2f ms | mot=%.2f ms | depth=%.2f ms | total=%.2f ms | node_total=%.2f ms", image_msg->header.stamp.sec, image_msg->header.stamp.nanosec, cam_bgr.cols, @@ -543,9 +628,17 @@ void CloudReprojectionRosNode::syncCallback( target_debug.poses_count, target_debug.tracks_count, target_debug.current_target_id_before, + target_debug.current_raw_track_id_before, target_observation.track_id, + target_debug.selected_raw_track_id, target_observation.detection_index, target_debug.found_existing_target ? "yes" : "no", + target_debug.selected_from_center_bootstrap ? "yes" : "no", + target_debug.reid_attempted ? "yes" : "no", + target_debug.recovered_by_reid ? "yes" : "no", + target_debug.reid_similarity, + target_debug.gallery_size, + target_debug.lost_frames, format_bbox_xyxy(target_observation.bbox_xyxy).c_str(), target_observation.depth, target_observation.confidence, @@ -617,25 +710,6 @@ void CloudReprojectionRosNode::syncCallback( if (send_overlay_ && overlay_compressed_pub_) { cv::Mat overlay_vis = odin_ros_driver::overlay_projected_cloud_on_image( cam_bgr, cloud_cam, reprojector_->getCameraParams()); -#ifdef ODIN_ROS_DRIVER_HAS_TARGET_OBSERVATION - if (debug_target_observation_ && target_observation_processor_) { - target_observation_processor_->draw_detected_poses( - overlay_vis, - target_debug.poses); - for (const auto& pixel : target_debug.valid_projected_pixels) { - if (pixel.x < 0 || pixel.x >= overlay_vis.cols || - pixel.y < 0 || pixel.y >= overlay_vis.rows) { - continue; - } - cv::circle(overlay_vis, pixel, 3, cv::Scalar(0, 255, 0), -1); - } - } - if (target_observation.valid && debug_target_observation_) { - odin_ros_driver::draw_target_observation_overlay( - overlay_vis, - target_observation); - } -#endif if (!overlay_vis.empty()) { std::vector obuf; const std::vector oenc = {cv::IMWRITE_JPEG_QUALITY, overlay_jpeg_quality_}; @@ -651,6 +725,32 @@ void CloudReprojectionRosNode::syncCallback( } } +#ifdef ODIN_ROS_DRIVER_HAS_TARGET_OBSERVATION + if ( + enable_target_observation_ && + debug_target_observation_ && + target_observation_processor_ && + detection_debug_compressed_pub_) { + cv::Mat detection_debug_vis = cam_bgr.clone(); + if (!detection_debug_vis.empty()) { + target_observation_processor_->draw_detected_poses(detection_debug_vis, target_debug.poses); + odin_ros_driver::draw_target_observation_overlay(detection_debug_vis, target_observation); + + std::vector dbuf; + const std::vector denc = {cv::IMWRITE_JPEG_QUALITY, overlay_jpeg_quality_}; + if (!cv::imencode(".jpg", detection_debug_vis, dbuf, denc)) { + RCLCPP_ERROR(this->get_logger(), "cv::imencode failed (detection debug)"); + return; + } + CompressedImage dmsg; + dmsg.header = image_msg->header; + dmsg.format = "jpeg"; + dmsg.data.assign(dbuf.begin(), dbuf.end()); + detection_debug_compressed_pub_->publish(dmsg); + } + } +#endif + if (combined_pub_) { const int H = std::max(depth_vis.rows, cam_bgr.rows); cv::Mat left = odin_ros_driver::resize_to_height(depth_vis, H); diff --git a/src/reid_trt_extractor.cpp b/src/reid_trt_extractor.cpp new file mode 100644 index 0000000..ffb38e9 --- /dev/null +++ b/src/reid_trt_extractor.cpp @@ -0,0 +1,375 @@ +#include "reid_trt_extractor.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace odin_ros_driver { + +namespace { + +constexpr float kImageNetMean[3] = {0.485f, 0.456f, 0.406f}; +constexpr float kImageNetStd[3] = {0.229f, 0.224f, 0.225f}; + +float l2_norm(const ReIDFeature& feature) +{ + float value = 0.0f; + for (const float v : feature) { + value += v * v; + } + return std::sqrt(value); +} + +void l2_normalize(ReIDFeature* feature) +{ + if (feature == nullptr || feature->empty()) { + return; + } + + const float norm = l2_norm(*feature); + if (norm <= 1e-6f) { + std::fill(feature->begin(), feature->end(), 0.0f); + return; + } + + for (float& value : *feature) { + value /= norm; + } +} + +nvinfer1::Dims make_dims4(int n, int c, int h, int w) +{ + nvinfer1::Dims dims{}; + dims.nbDims = 4; + dims.d[0] = n; + dims.d[1] = c; + dims.d[2] = h; + dims.d[3] = w; + return dims; +} + +} // namespace + +TrackGallery::TrackGallery(size_t max_features) + : max_features_(std::max(1, max_features)) +{ +} + +void TrackGallery::clear() +{ + features_.clear(); + last_update_frame_ = -1; +} + +void TrackGallery::add(const ReIDFeature& feature, int64_t frame_index) +{ + if (feature.empty()) { + return; + } + + if (features_.size() >= max_features_) { + features_.pop_front(); + } + features_.push_back(feature); + last_update_frame_ = frame_index; +} + +float TrackGallery::best_similarity(const ReIDFeature& query) const +{ + if (query.empty() || features_.empty()) { + return -1.0f; + } + + float best = -1.0f; + for (const auto& feature : features_) { + const size_t count = std::min(feature.size(), query.size()); + float similarity = 0.0f; + for (size_t i = 0; i < count; ++i) { + similarity += feature[i] * query[i]; + } + best = std::max(best, similarity); + } + return best; +} + +struct ReIDTensorRTExtractor::Impl { + struct Binding { + std::string name; + bool is_input = false; + nvinfer1::DataType dtype = nvinfer1::DataType::kFLOAT; + size_t max_elements = 0; + yolos::trt::DeviceBuffer device; + }; + + yolos::trt::TrtUniquePtr runtime{nullptr}; + yolos::trt::TrtUniquePtr engine{nullptr}; + yolos::trt::TrtUniquePtr context{nullptr}; + yolos::trt::CudaStream stream; + yolos::trt::PinnedBuffer host_input; + yolos::trt::PinnedBuffer host_output; + Binding input_binding; + Binding output_binding; + bool dynamic_batch = false; + int max_runtime_batch_size = 1; +}; + +ReIDTensorRTExtractor::ReIDTensorRTExtractor() = default; + +ReIDTensorRTExtractor::~ReIDTensorRTExtractor() = default; + +void ReIDTensorRTExtractor::initialize(const ReIDTensorRTConfig& config) +{ + if (config.engine_path.empty()) { + throw std::runtime_error("ReID TensorRT engine path is empty"); + } + + auto impl = std::make_unique(); + const std::vector engine_bytes = yolos::trt::readFile(config.engine_path); + + impl->runtime.reset(nvinfer1::createInferRuntime(yolos::trt::getLogger())); + if (!impl->runtime) { + throw std::runtime_error("Failed to create TensorRT runtime for ReID"); + } + + impl->engine.reset( + impl->runtime->deserializeCudaEngine(engine_bytes.data(), engine_bytes.size())); + if (!impl->engine) { + throw std::runtime_error("Failed to deserialize ReID TensorRT engine: " + config.engine_path); + } + + impl->context.reset(impl->engine->createExecutionContext()); + if (!impl->context) { + throw std::runtime_error("Failed to create ReID TensorRT execution context"); + } + + if (impl->engine->getNbOptimizationProfiles() > 0 && + !impl->context->setOptimizationProfileAsync(0, impl->stream.get())) { + throw std::runtime_error("Failed to select optimization profile 0 for ReID TensorRT engine"); + } + + bool found_input = false; + bool found_output = false; + for (int i = 0; i < impl->engine->getNbIOTensors(); ++i) { + const char* name = impl->engine->getIOTensorName(i); + if (name == nullptr) { + continue; + } + const auto mode = impl->engine->getTensorIOMode(name); + if (mode == nvinfer1::TensorIOMode::kINPUT && !found_input) { + impl->input_binding.name = name; + impl->input_binding.is_input = true; + impl->input_binding.dtype = impl->engine->getTensorDataType(name); + found_input = true; + } else if (mode == nvinfer1::TensorIOMode::kOUTPUT && !found_output) { + impl->output_binding.name = name; + impl->output_binding.is_input = false; + impl->output_binding.dtype = impl->engine->getTensorDataType(name); + found_output = true; + } + } + + if (!found_input || !found_output) { + throw std::runtime_error("Expected one input and one output tensor in ReID TensorRT engine"); + } + + if (impl->input_binding.dtype != nvinfer1::DataType::kFLOAT) { + throw std::runtime_error("Only float32 ReID TensorRT inputs are supported"); + } + if (impl->output_binding.dtype != nvinfer1::DataType::kFLOAT && + impl->output_binding.dtype != nvinfer1::DataType::kHALF) { + throw std::runtime_error("Unsupported ReID TensorRT output dtype"); + } + + nvinfer1::Dims input_dims = impl->engine->getTensorShape(impl->input_binding.name.c_str()); + if (input_dims.nbDims != 4) { + throw std::runtime_error("Expected ReID TensorRT input dims [N,C,H,W]"); + } + + config_ = config; + if (input_dims.d[2] > 0) { + config_.input_height = input_dims.d[2]; + } + if (input_dims.d[3] > 0) { + config_.input_width = input_dims.d[3]; + } + + impl->dynamic_batch = (input_dims.d[0] == -1); + impl->max_runtime_batch_size = std::max(1, config_.max_batch_size); + + nvinfer1::Dims resolved_input_dims = input_dims; + if (impl->dynamic_batch) { + const nvinfer1::Dims max_input_dims = make_dims4( + impl->max_runtime_batch_size, 3, config_.input_height, config_.input_width); + if (!impl->context->setInputShape(impl->input_binding.name.c_str(), max_input_dims)) { + throw std::runtime_error("Failed to set ReID TensorRT input shape"); + } + resolved_input_dims = max_input_dims; + } else { + impl->max_runtime_batch_size = std::max(1, static_cast(input_dims.d[0])); + } + + const nvinfer1::Dims resolved_output_dims = + impl->context->getTensorShape(impl->output_binding.name.c_str()); + if (resolved_output_dims.nbDims < 2) { + throw std::runtime_error("Unexpected ReID TensorRT output dimensions"); + } + if (resolved_output_dims.d[resolved_output_dims.nbDims - 1] > 0) { + config_.feature_dim = resolved_output_dims.d[resolved_output_dims.nbDims - 1]; + } + + impl->input_binding.max_elements = + static_cast(impl->max_runtime_batch_size) * 3U * + static_cast(config_.input_height) * static_cast(config_.input_width); + impl->output_binding.max_elements = + static_cast(impl->max_runtime_batch_size) * static_cast(config_.feature_dim); + + impl->host_input.resize(impl->input_binding.max_elements * sizeof(float)); + impl->host_output.resize( + impl->output_binding.max_elements * yolos::trt::dataTypeBytes(impl->output_binding.dtype)); + impl->input_binding.device.resize(impl->host_input.size()); + impl->output_binding.device.resize(impl->host_output.size()); + + if (!impl->context->setTensorAddress( + impl->input_binding.name.c_str(), impl->input_binding.device.data())) { + throw std::runtime_error("Failed to bind ReID TensorRT input tensor address"); + } + if (!impl->context->setTensorAddress( + impl->output_binding.name.c_str(), impl->output_binding.device.data())) { + throw std::runtime_error("Failed to bind ReID TensorRT output tensor address"); + } + + impl_ = std::move(impl); + initialized_ = true; +} + +int ReIDTensorRTExtractor::feature_dim() const +{ + return config_.feature_dim; +} + +std::vector ReIDTensorRTExtractor::extract( + const cv::Mat& frame_bgr, + const std::vector& boxes) const +{ + std::vector features( + boxes.size(), ReIDFeature(static_cast(config_.feature_dim), 0.0f)); + if (!initialized_ || impl_ == nullptr || frame_bgr.empty() || boxes.empty()) { + return features; + } + + std::vector valid_indices; + std::vector valid_boxes; + valid_indices.reserve(boxes.size()); + valid_boxes.reserve(boxes.size()); + for (size_t i = 0; i < boxes.size(); ++i) { + const cv::Rect clamped = boxes[i] & cv::Rect(0, 0, frame_bgr.cols, frame_bgr.rows); + if (clamped.width < 4 || clamped.height < 4) { + continue; + } + valid_indices.push_back(static_cast(i)); + valid_boxes.push_back(clamped); + } + + if (valid_boxes.empty()) { + return features; + } + + const size_t pixels_per_crop = + static_cast(config_.input_height) * static_cast(config_.input_width); + const size_t floats_per_crop = 3U * pixels_per_crop; + auto* input_ptr = impl_->host_input.as(); + const size_t max_batch = static_cast(std::max(1, impl_->max_runtime_batch_size)); + + for (size_t start = 0; start < valid_boxes.size(); start += max_batch) { + const size_t batch = + std::min(valid_boxes.size() - start, max_batch); + + for (size_t local = 0; local < batch; ++local) { + const cv::Rect crop_box = valid_boxes[start + local]; + cv::Mat resized; + cv::resize( + frame_bgr(crop_box), + resized, + cv::Size(config_.input_width, config_.input_height), + 0.0, + 0.0, + cv::INTER_LINEAR); + + cv::Mat rgb; + cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB); + rgb.convertTo(rgb, CV_32FC3, 1.0 / 255.0); + + std::vector channels(3); + cv::split(rgb, channels); + float* dst = input_ptr + local * floats_per_crop; + for (int channel = 0; channel < 3; ++channel) { + channels[channel] = + (channels[channel] - kImageNetMean[channel]) / kImageNetStd[channel]; + std::memcpy( + dst + static_cast(channel) * pixels_per_crop, + channels[channel].ptr(), + pixels_per_crop * sizeof(float)); + } + } + + if (impl_->dynamic_batch) { + const nvinfer1::Dims batch_dims = make_dims4( + static_cast(batch), 3, config_.input_height, config_.input_width); + if (!impl_->context->setInputShape(impl_->input_binding.name.c_str(), batch_dims)) { + throw std::runtime_error("Failed to set batch shape for ReID TensorRT inference"); + } + } + + const nvinfer1::Dims output_dims = + impl_->context->getTensorShape(impl_->output_binding.name.c_str()); + const size_t output_elements = static_cast(yolos::trt::volume(output_dims)); + const size_t output_bytes = + output_elements * yolos::trt::dataTypeBytes(impl_->output_binding.dtype); + + CUDA_CHECK(cudaMemcpyAsync( + impl_->input_binding.device.data(), + input_ptr, + batch * floats_per_crop * sizeof(float), + cudaMemcpyHostToDevice, + impl_->stream.get())); + + if (!impl_->context->enqueueV3(impl_->stream.get())) { + throw std::runtime_error("TensorRT enqueueV3 failed for ReID inference"); + } + + CUDA_CHECK(cudaMemcpyAsync( + impl_->host_output.data(), + impl_->output_binding.device.data(), + output_bytes, + cudaMemcpyDeviceToHost, + impl_->stream.get())); + impl_->stream.synchronize(); + + for (size_t local = 0; local < batch; ++local) { + ReIDFeature feature(static_cast(config_.feature_dim), 0.0f); + if (impl_->output_binding.dtype == nvinfer1::DataType::kFLOAT) { + const float* src = impl_->host_output.as() + local * config_.feature_dim; + std::copy(src, src + config_.feature_dim, feature.begin()); + } else { + const __half* src = + static_cast(impl_->host_output.data()) + local * config_.feature_dim; + for (int dim = 0; dim < config_.feature_dim; ++dim) { + feature[static_cast(dim)] = __half2float(src[dim]); + } + } + l2_normalize(&feature); + features[static_cast(valid_indices[start + local])] = std::move(feature); + } + } + + return features; +} + +} // namespace odin_ros_driver diff --git a/src/target_observation_processing.cpp b/src/target_observation_processing.cpp index ba2ec32..d75b008 100644 --- a/src/target_observation_processing.cpp +++ b/src/target_observation_processing.cpp @@ -18,6 +18,25 @@ namespace { constexpr std::array kHipKeypointIndices = {11, 12}; +cv::Rect track_row_to_rect(const Eigen::MatrixXf& tracks, int row) +{ + const int x1 = static_cast(std::floor(tracks(row, 0))); + const int y1 = static_cast(std::floor(tracks(row, 1))); + const int x2 = static_cast(std::ceil(tracks(row, 2))); + const int y2 = static_cast(std::ceil(tracks(row, 3))); + return cv::Rect(x1, y1, std::max(0, x2 - x1), std::max(0, y2 - y1)); +} + +bool feature_has_signal(const ReIDFeature& feature) +{ + for (const float value : feature) { + if (std::abs(value) > 1e-6f) { + return true; + } + } + return false; +} + float median_in_place(std::vector& values) { if (values.empty()) { @@ -57,7 +76,30 @@ void TargetObservationProcessor::initialize(const TargetObservationConfig& confi config_.yolo_labels_path); tracker_ = std::make_unique( 0.3f, 30, 50, 3, 0.3f, false, 1); + target_gallery_ = TrackGallery(static_cast(std::max(1, config_.reid_gallery_size))); current_target_id_ = -1; + current_raw_track_id_ = -1; + frame_index_ = 0; + lost_track_frames_ = 0; + reid_extractor_.reset(); + if (config_.reid_enabled && !config_.reid_engine_path.empty()) { + try { + ReIDTensorRTConfig reid_config; + reid_config.engine_path = config_.reid_engine_path; + reid_config.input_height = config_.reid_input_height; + reid_config.input_width = config_.reid_input_width; + reid_config.feature_dim = config_.reid_feature_dim; + reid_config.max_batch_size = config_.reid_max_batch_size; + + reid_extractor_ = std::make_unique(); + reid_extractor_->initialize(reid_config); + } catch (const std::exception&) { + reid_extractor_.reset(); + config_.reid_enabled = false; + } + } else { + config_.reid_enabled = false; + } initialized_ = true; } @@ -70,16 +112,33 @@ TargetObservation TargetObservationProcessor::process( { const auto total_start = std::chrono::steady_clock::now(); TargetObservation empty; - if (debug_info) { - *debug_info = TargetObservationDebugInfo{}; - debug_info->current_target_id_before = current_target_id_; - } - if (!initialized_ || camera_bgr.empty()) { - if (debug_info) { - debug_info->total_ms = + ++frame_index_; + auto finalize_empty = [&](TargetObservationDebugInfo* info) { + if (info) { + info->gallery_size = static_cast(target_gallery_.size()); + info->lost_frames = lost_track_frames_; + info->total_ms = std::chrono::duration(std::chrono::steady_clock::now() - total_start).count(); } return empty; + }; + auto mark_lost = [&]() { + ++lost_track_frames_; + if (config_.reid_lost_timeout_frames > 0 && + lost_track_frames_ >= config_.reid_lost_timeout_frames) { + current_target_id_ = -1; + current_raw_track_id_ = -1; + lost_track_frames_ = 0; + target_gallery_.clear(); + } + }; + if (debug_info) { + *debug_info = TargetObservationDebugInfo{}; + debug_info->current_target_id_before = current_target_id_; + debug_info->current_raw_track_id_before = current_raw_track_id_; + } + if (!initialized_ || camera_bgr.empty()) { + return finalize_empty(debug_info); } const auto yolo_start = std::chrono::steady_clock::now(); @@ -92,11 +151,8 @@ TargetObservation TargetObservationProcessor::process( std::chrono::duration(yolo_end - yolo_start).count(); } if (poses.empty()) { - if (debug_info) { - debug_info->total_ms = - std::chrono::duration(std::chrono::steady_clock::now() - total_start).count(); - } - return empty; + mark_lost(); + return finalize_empty(debug_info); } const Eigen::MatrixXf detections = format_detections(poses); @@ -109,35 +165,37 @@ TargetObservation TargetObservationProcessor::process( std::chrono::duration(mot_end - mot_start).count(); } if (tracks.rows() == 0) { - if (debug_info) { - debug_info->total_ms = - std::chrono::duration(std::chrono::steady_clock::now() - total_start).count(); - } - return empty; + mark_lost(); + return finalize_empty(debug_info); } TargetObservation target = select_target( - poses, tracks, camera_bgr.cols, camera_bgr.rows, debug_info); + camera_bgr, poses, tracks, camera_bgr.cols, camera_bgr.rows, debug_info); if (!target.valid) { - if (debug_info) { - debug_info->total_ms = - std::chrono::duration(std::chrono::steady_clock::now() - total_start).count(); - } - return empty; + mark_lost(); + return finalize_empty(debug_info); } + current_raw_track_id_ = target.raw_track_id; + lost_track_frames_ = 0; + if (current_target_id_ < 0) { + current_target_id_ = target.raw_track_id; + } + const bool force_gallery_update = + target_gallery_.empty() || + (debug_info != nullptr && + (debug_info->selected_from_center_bootstrap || debug_info->recovered_by_reid)); + update_target_gallery(camera_bgr, target, force_gallery_update); + target = enrich_target_with_cloud( target, poses, cloud_in_cam, reprojector, odom_pose, debug_info); if (!target.valid) { - if (debug_info) { - debug_info->total_ms = - std::chrono::duration(std::chrono::steady_clock::now() - total_start).count(); - } - return empty; + return finalize_empty(debug_info); } - current_target_id_ = target.track_id; if (debug_info) { + debug_info->gallery_size = static_cast(target_gallery_.size()); + debug_info->lost_frames = lost_track_frames_; debug_info->total_ms = std::chrono::duration(std::chrono::steady_clock::now() - total_start).count(); } @@ -151,7 +209,20 @@ void TargetObservationProcessor::draw_detected_poses( if (!initialized_ || image_bgr.empty() || poses.empty()) { return; } - yolo_->drawPoses(image_bgr, poses); + + yolo_->drawSkeletonsOnly(image_bgr, poses); + + const cv::Scalar detection_box_color(0, 255, 255); + for (size_t i = 0; i < poses.size(); ++i) { + const auto& pose = poses[i]; + const cv::Point top_left( + static_cast(std::lround(pose.box.x)), + static_cast(std::lround(pose.box.y))); + const cv::Point bottom_right( + static_cast(std::lround(pose.box.x + pose.box.width)), + static_cast(std::lround(pose.box.y + pose.box.height))); + cv::rectangle(image_bgr, top_left, bottom_right, detection_box_color, 3); + } } Eigen::MatrixXf TargetObservationProcessor::format_detections( @@ -170,6 +241,7 @@ Eigen::MatrixXf TargetObservationProcessor::format_detections( } TargetObservation TargetObservationProcessor::select_target( + const cv::Mat& camera_bgr, const std::vector& poses, const Eigen::MatrixXf& tracks, int image_width, @@ -179,28 +251,58 @@ TargetObservation TargetObservationProcessor::select_target( TargetObservation observation; int best_row = -1; bool found_existing = false; + bool attempted_reid = false; const float img_cx = static_cast(image_width) * 0.5f; const float img_cy = static_cast(image_height) * 0.5f; - float best_dist = std::numeric_limits::max(); + const bool has_tracking_history = + current_target_id_ >= 0 || current_raw_track_id_ >= 0 || !target_gallery_.empty(); + const bool can_attempt_reid = + config_.reid_enabled && + reid_extractor_ != nullptr && + reid_extractor_->initialized() && + !target_gallery_.empty(); for (int i = 0; i < tracks.rows(); ++i) { const int track_id = static_cast(tracks(i, 4)); - if (track_id == current_target_id_) { + if (track_id == current_raw_track_id_) { best_row = i; found_existing = true; break; } + } - const float cx = 0.5f * (tracks(i, 0) + tracks(i, 2)); - const float cy = 0.5f * (tracks(i, 1) + tracks(i, 3)); - const float dist = (cx - img_cx) * (cx - img_cx) + (cy - img_cy) * (cy - img_cy); - if (dist < best_dist) { - best_dist = dist; - best_row = i; + if (!found_existing && has_tracking_history && can_attempt_reid) { + float best_similarity = -1.0f; + attempted_reid = true; + best_row = recover_target_with_reid(camera_bgr, tracks, &best_similarity); + if (debug_info) { + debug_info->reid_attempted = true; + debug_info->reid_similarity = best_similarity; + debug_info->recovered_by_reid = (best_row >= 0); + } + } + + if (best_row < 0 && !has_tracking_history) { + float best_dist = std::numeric_limits::max(); + for (int i = 0; i < tracks.rows(); ++i) { + const float cx = 0.5f * (tracks(i, 0) + tracks(i, 2)); + const float cy = 0.5f * (tracks(i, 1) + tracks(i, 3)); + const float dist = (cx - img_cx) * (cx - img_cx) + (cy - img_cy) * (cy - img_cy); + if (dist < best_dist) { + best_dist = dist; + best_row = i; + } + } + if (debug_info) { + debug_info->selected_from_center_bootstrap = (best_row >= 0); } } if (best_row < 0) { + if (debug_info) { + debug_info->found_existing_target = found_existing; + debug_info->reid_attempted = attempted_reid; + } return observation; } @@ -214,12 +316,16 @@ TargetObservation TargetObservationProcessor::select_target( observation.valid = true; if (debug_info) { debug_info->found_existing_target = found_existing; + debug_info->reid_attempted = attempted_reid; debug_info->target_selected = true; } - observation.track_id = static_cast(tracks(best_row, 4)); + const int selected_raw_track_id = static_cast(tracks(best_row, 4)); + observation.raw_track_id = selected_raw_track_id; + observation.track_id = (current_target_id_ >= 0) ? current_target_id_ : selected_raw_track_id; observation.detection_index = detection_index; if (debug_info) { debug_info->selected_track_id = observation.track_id; + debug_info->selected_raw_track_id = selected_raw_track_id; debug_info->detection_index = detection_index; } observation.confidence = pose.conf; @@ -235,13 +341,144 @@ TargetObservation TargetObservationProcessor::select_target( observation.keypoints_xyc[i * 3 + 1] = pose.keypoints[i].y; observation.keypoints_xyc[i * 3 + 2] = pose.keypoints[i].confidence; } - - if (!found_existing) { - current_target_id_ = observation.track_id; - } return observation; } +int TargetObservationProcessor::recover_target_with_reid( + const cv::Mat& camera_bgr, + const Eigen::MatrixXf& tracks, + float* best_similarity) const +{ + if (best_similarity) { + *best_similarity = -1.0f; + } + if (!config_.reid_enabled || + reid_extractor_ == nullptr || + !reid_extractor_->initialized() || + target_gallery_.empty() || + camera_bgr.empty() || + tracks.rows() == 0) { + return -1; + } + + std::vector candidate_rows; + std::vector candidate_boxes; + candidate_rows.reserve(static_cast(tracks.rows())); + candidate_boxes.reserve(static_cast(tracks.rows())); + for (int row = 0; row < tracks.rows(); ++row) { + const cv::Rect bbox = track_row_to_rect(tracks, row); + if (!is_good_reid_crop(bbox, camera_bgr.size())) { + continue; + } + candidate_rows.push_back(row); + candidate_boxes.push_back(bbox); + } + if (candidate_boxes.empty()) { + return -1; + } + + const auto features = reid_extractor_->extract(camera_bgr, candidate_boxes); + int best_row = -1; + float top1 = -1.0f; + float top2 = -1.0f; + for (size_t i = 0; i < features.size(); ++i) { + if (!feature_has_signal(features[i])) { + continue; + } + const float similarity = target_gallery_.best_similarity(features[i]); + if (similarity > top1) { + top2 = top1; + top1 = similarity; + best_row = candidate_rows[i]; + } else if (similarity > top2) { + top2 = similarity; + } + } + + if (best_similarity) { + *best_similarity = top1; + } + if (best_row < 0 || top1 < config_.reid_match_threshold) { + return -1; + } + + const float second = (top2 < 0.0f) ? -1.0f : top2; + if (second >= 0.0f && + (top1 - second) < config_.reid_gap_threshold) { + return -1; + } + return best_row; +} + +void TargetObservationProcessor::update_target_gallery( + const cv::Mat& camera_bgr, + const TargetObservation& target, + bool force_update) +{ + if (!config_.reid_enabled || + reid_extractor_ == nullptr || + !reid_extractor_->initialized() || + !target.valid || + !should_update_gallery(force_update) || + camera_bgr.empty()) { + return; + } + + const int x1 = static_cast(std::floor(target.bbox_xyxy[0])); + const int y1 = static_cast(std::floor(target.bbox_xyxy[1])); + const int x2 = static_cast(std::ceil(target.bbox_xyxy[2])); + const int y2 = static_cast(std::ceil(target.bbox_xyxy[3])); + const cv::Rect bbox(x1, y1, std::max(0, x2 - x1), std::max(0, y2 - y1)); + if (!is_good_reid_crop(bbox, camera_bgr.size())) { + return; + } + + const auto features = reid_extractor_->extract(camera_bgr, {bbox}); + if (features.empty() || !feature_has_signal(features.front())) { + return; + } + target_gallery_.add(features.front(), frame_index_); +} + +bool TargetObservationProcessor::should_update_gallery(bool force_update) const +{ + if (force_update || target_gallery_.empty()) { + return true; + } + if (config_.reid_feature_update_interval <= 0) { + return true; + } + return (frame_index_ - target_gallery_.last_update_frame()) >= + static_cast(config_.reid_feature_update_interval); +} + +bool TargetObservationProcessor::is_good_reid_crop( + const cv::Rect& bbox, + const cv::Size& image_size) const +{ + if (bbox.width <= 0 || bbox.height <= 0) { + return false; + } + if (bbox.area() < config_.reid_min_crop_area) { + return false; + } + const float aspect_ratio = + static_cast(bbox.width) / static_cast(bbox.height); + if (aspect_ratio > config_.reid_max_crop_aspect_ratio) { + return false; + } + + const int margin = 3; + if (bbox.x < margin || bbox.y < margin) { + return false; + } + if ((bbox.x + bbox.width) > (image_size.width - margin) || + (bbox.y + bbox.height) > (image_size.height - margin)) { + return false; + } + return true; +} + TargetObservation TargetObservationProcessor::enrich_target_with_cloud( const TargetObservation& target, const std::vector& poses, @@ -427,12 +664,12 @@ void draw_target_observation_overlay( return; } - const cv::Scalar box_color(80, 220, 255); + const cv::Scalar box_color(0, 0, 255); const int x1 = static_cast(std::lround(observation.bbox_xyxy[0])); const int y1 = static_cast(std::lround(observation.bbox_xyxy[1])); const int x2 = static_cast(std::lround(observation.bbox_xyxy[2])); const int y2 = static_cast(std::lround(observation.bbox_xyxy[3])); - cv::rectangle(image_bgr, cv::Point(x1, y1), cv::Point(x2, y2), box_color, 2); + cv::rectangle(image_bgr, cv::Point(x1, y1), cv::Point(x2, y2), box_color, 4); for (size_t i = 0; i < 17; ++i) { const float conf = observation.keypoints_xyc[i * 3 + 2];