odin_ros_driver1/include/target_observation_processing.hpp

#pragma once

#include <array>
#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 "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 debug = false;
};

struct TargetObservation {
    bool valid = false;
    int 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 poses_count = 0;
    int tracks_count = 0;
    int selected_track_id = -1;
    int detection_index = -1;
    int projected_cloud_points = 0;
    int depth_sample_count = 0;
    bool found_existing_target = false;
    bool target_selected = false;
    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 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;

    TargetObservationConfig config_;
    std::unique_ptr<yolos::pose::YOLOPoseDetector> yolo_;
    std::unique_ptr<motcpp::trackers::ByteTrack> tracker_;
    int current_target_id_ = -1;
    bool initialized_ = false;
};

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

}  // namespace odin_ros_driver