odin_ros_driver1/include/target_observation_processing.hpp

#pragma once

#include <array>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include <Eigen/Dense>
#include <opencv2/core.hpp>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>

#include "cloud_reprojector.hpp"
#include "reid_trt_extractor.hpp"
#include "motcpp/trackers/bytetrack.hpp"
#include "yolos/tasks/pose.hpp"

namespace odin_ros_driver {

struct TargetObservationConfig {
    std::string yolo_engine_path;
    std::string yolo_labels_path;
    float yolo_conf = 0.45f;
    float yolo_nms = 0.50f;
    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;
    float depth_confidence = 0.0f;
    std::array<float, 4> bbox_xyxy{0.0f, 0.0f, 0.0f, 0.0f};
    std::array<float, 17 * 3> keypoints_xyc{};
    Eigen::Vector3f target_pos_cam = Eigen::Vector3f::Zero();
    Eigen::Vector3f target_pos_world = Eigen::Vector3f::Zero();
};

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;
    float total_ms = 0.0f;
    std::vector<yolos::pose::PoseResult> poses;
    std::vector<cv::Point> valid_projected_pixels;
};

class TargetObservationProcessor {
public:
    TargetObservationProcessor() = default;

    void initialize(const TargetObservationConfig& config);

    TargetObservation process(
        const cv::Mat& camera_bgr,
        const pcl::PointCloud<pcl::PointXYZRGB>& cloud_in_cam,
        const CloudReprojector& reprojector,
        const CloudReprojector::OdomPose& odom_pose,
        TargetObservationDebugInfo* debug_info = nullptr);

    bool initialized() const { return initialized_; }

    void draw_detected_poses(
        cv::Mat& image_bgr,
        const std::vector<yolos::pose::PoseResult>& poses) const;

private:
    Eigen::MatrixXf format_detections(
        const std::vector<yolos::pose::PoseResult>& poses) const;

    TargetObservation select_target(
        const cv::Mat& camera_bgr,
        const std::vector<yolos::pose::PoseResult>& poses,
        const Eigen::MatrixXf& tracks,
        int image_width,
        int image_height,
        TargetObservationDebugInfo* debug_info);

    TargetObservation enrich_target_with_cloud(
        const TargetObservation& target,
        const std::vector<yolos::pose::PoseResult>& poses,
        const pcl::PointCloud<pcl::PointXYZRGB>& cloud_in_cam,
        const CloudReprojector& reprojector,
        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<yolos::pose::YOLOPoseDetector> yolo_;
    std::unique_ptr<motcpp::trackers::ByteTrack> tracker_;
    std::unique_ptr<ReIDTensorRTExtractor> 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;
};

void draw_target_observation_overlay(
    cv::Mat& image_bgr,
    const TargetObservation& observation);

}  // namespace odin_ros_driver