odin_ros_driver1/include/target_observation_processing.hpp

#pragma once

#include <array>
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#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 lost_detection_compensation_enabled = true;
    float lost_detection_roi_scale = 2.0f;
    int lost_detection_roi_min_size_px = 192;
    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;
    int max_concurrent_tracks = 16;
    // 3D position gate for ReID rebind. dmax = base_m + per_frame_m *
    // frames_since_last_seen. soft mode (hard=false) logs via debug counters
    // but still allows the similarity match; hard mode rejects the pair.
    float reid_3d_gate_base_m = 1.5f;
    float reid_3d_gate_per_frame_m = 0.1f;
    bool reid_3d_gate_hard = false;
    bool debug = false;
};

struct TargetObservation {
    bool valid = false;
    int track_id = -1;
    int raw_track_id = -1;
    int detection_index = -1;
    bool is_primary_target = false;
    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 detection_pos_cam = Eigen::Vector3f::Zero();
    Eigen::Vector3f detection_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 active_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;
    int reid_3d_gate_rejections = 0;     // hard-mode rejections this frame
    int reid_3d_gate_soft_warnings = 0;  // soft-mode over-threshold this frame
    bool found_existing_target = false;
    bool target_selected = false;
    bool selected_from_center_bootstrap = false;
    bool reid_attempted = false;
    bool recovered_by_reid = false;
    bool lost_detection_compensation_attempted = false;
    bool recovered_by_lost_detection_compensation = false;
    float reid_similarity = -1.0f;
    float yolo_ms = 0.0f;
    float mot_ms = 0.0f;
    float cloud_project_ms = 0.0f;
    float bind_ms = 0.0f;        // entire bind_and_update_tracks (covers depth_ms)
    float depth_ms = 0.0f;       // Phase 2A per-entry 3D compute (subset of bind_ms)
    float reid_extract_ms = 0.0f; // Phase 2B TRT feature extraction (subset of bind_ms)
    float age_out_ms = 0.0f;
    float select_ms = 0.0f;
    float enrich_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);

    // All active tracks with a fresh 3D estimate from the most recent
    // process() call. One TargetObservation per entry; is_primary_target
    // marks the follow target. Empty if process() hasn't run yet.
    std::vector<TargetObservation> snapshot_active_tracks() const;
    int current_target_stable_id() const { return current_target_stable_id_; }

    bool initialized() const { return initialized_; }

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

private:
    struct ProjectedCloud {
        std::vector<cv::Point2f> uv;
        std::vector<Eigen::Vector3f> cam;
    };

    struct EntryCloudEstimate {
        bool valid = false;
        float depth = -1.0f;
        float depth_confidence = 0.0f;
        Eigen::Vector3f pos_cam = Eigen::Vector3f::Zero();
        Eigen::Vector3f pos_world = Eigen::Vector3f::Zero();
        int depth_sample_count = 0;
        std::vector<cv::Point> valid_projected_pixels;
    };

    struct TrackEntry {
        int stable_id = -1;
        int last_raw_track_id = -1;
        TrackGallery gallery;
        cv::Rect last_bbox;
        bool has_last_bbox = false;
        int64_t last_seen_frame = -1;
        bool just_rebound_by_reid = false;
        float last_rebind_similarity = -1.0f;
        // Per-entry 3D state; refreshed whenever compute_entry_3d succeeds.
        float last_depth = -1.0f;
        float last_depth_confidence = 0.0f;
        Eigen::Vector3f last_pos_cam = Eigen::Vector3f::Zero();
        Eigen::Vector3f last_pos_world = Eigen::Vector3f::Zero();
        bool has_last_pos = false;
        int64_t last_pos_frame = -1;
        // 2D observation cache — refreshed every frame in Phase 5 while the
        // entry is seen. Consumed by snapshot_active_tracks() to populate
        // TargetObservation messages for all tracked people.
        std::array<float, 17 * 3> last_keypoints_xyc{};
        int last_detection_index = -1;
        float last_confidence = 0.0f;
        // Debug snapshot from the most recent successful estimate — only the
        // follow target's values are copied into TargetObservationDebugInfo.
        int last_depth_sample_count = 0;
        std::vector<cv::Point> last_valid_projected_pixels;
    };

    Eigen::MatrixXf format_detections(
        const std::vector<yolos::pose::PoseResult>& poses) const;
    std::vector<yolos::pose::PoseResult> detect_poses_with_lost_compensation(
        const cv::Mat& camera_bgr,
        TargetObservationDebugInfo* debug_info) const;
    cv::Rect build_lost_detection_roi(
        const cv::Size& image_size,
        const cv::Rect& last_bbox) const;
    static void offset_poses(
        std::vector<yolos::pose::PoseResult>& poses,
        const cv::Point& offset);

    ProjectedCloud project_cloud_to_image(
        const pcl::PointCloud<pcl::PointXYZRGB>& cloud_in_cam,
        const CloudReprojector& reprojector) const;

    EntryCloudEstimate compute_entry_3d(
        const yolos::pose::PoseResult& pose,
        const std::array<float, 4>& bbox_xyxy,
        const ProjectedCloud& projected,
        const CloudReprojector& reprojector,
        const CloudReprojector::OdomPose& odom_pose) const;

    std::vector<int> bind_and_update_tracks(
        const cv::Mat& camera_bgr,
        const Eigen::MatrixXf& tracks,
        const std::vector<yolos::pose::PoseResult>& poses,
        const ProjectedCloud& projected,
        const CloudReprojector& reprojector,
        const CloudReprojector::OdomPose& odom_pose,
        TargetObservationDebugInfo* debug_info);

    void age_out_tracks();

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

    TargetObservation enrich_target_with_cloud(
        const TargetObservation& target,
        TargetObservationDebugInfo* debug_info) 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_;

    std::unordered_map<int, TrackEntry> tracks_;
    int next_stable_id_ = 0;
    int current_target_stable_id_ = -1;
    int64_t frame_index_ = 0;
    bool initialized_ = false;
};

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

}  // namespace odin_ros_driver