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;
    // Optional detection ROI applied to every frame before the YOLO call.
    // Trades coverage (edge strips are ignored) for effective resolution on
    // the subjects inside the ROI. Width/height are in image pixels; the
    // ROI is clamped to the image and silently disabled if width*height
    // covers the full image. Horizontal placement is always centered. The
    // vertical offset is pixels from the image top; set to -1 to auto-center.
    bool detection_roi_enabled = false;
    int detection_roi_width_px = 1440;
    int detection_roi_height_px = 1080;
    int detection_roi_y_offset_px = 0;
    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 = 8;
    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 = 8;
    // Edge margin (in pixels) that is_good_reid_crop requires between the
    // bbox and every image border. Set to 0 to allow crops touching the edge
    // (common when the person stands close to the camera and the head gets
    // cut off). Legacy value was 3.
    int reid_crop_edge_margin_px = 0;
    // 3D-proximity fallback: when a row lacks a ReID feature (e.g. bbox was
    // too close to the edge for is_good_reid_crop to pass) but has a fresh
    // 3D estimate, rebind it to the follow target's stable_id if it lies
    // within a strict (tighter than reid_3d_gate) distance of the target's
    // last known world pos. Guards against ID inflation during close-range
    // partial views that block ReID extraction.
    bool reid_3d_fallback_enabled = true;
    float reid_3d_fallback_base_m = 0.75f;
    float reid_3d_fallback_per_frame_m = 0.05f;
    int reid_confirm_hits_before_allocate = 5;
    bool debug = false;
    bool debug_reid = 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
    int reid_3d_fallback_rebinds = 0;    // follow-target rebinds by 3D fallback
    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;          // cosine sim from Phase 3 ReID (≥0), or -1 if N/A
    float reid_3d_fallback_distance = -1.0f; // metres from Phase 3.5, or -1 if N/A
    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;
    // Per-track-row snapshot from bind_and_update_tracks (aligned, one entry
    // per row of the tracks matrix from ByteTrack). raw_track_ids are what
    // ByteTrack emitted for this frame; row_stable_ids are what our
    // TrackManager ultimately bound each row to. Both empty when the
    // pipeline bailed before binding.
    std::vector<int> raw_track_ids;
    std::vector<int> row_stable_ids;
    std::vector<int> row_detection_indices;
    // Center-cropped ROI actually used for the primary YOLO pass this frame.
    // Empty (area()==0) when config_.detection_roi_enabled is false or when
    // the ROI was clamped away to nothing.
    cv::Rect detection_roi;
    // Flat record of every Phase 3 ReID similarity comparison that actually
    // ran this frame. Tuple is (row_raw_track_id, candidate_stable_id, sim).
    // Useful for debugging "why didn't this rebind" cases — read the log
    // and you see each candidate's sim against each active gallery.
    std::vector<std::tuple<int, int, float>> reid_comparisons;
};

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_; }
    int current_frame_primary_stable_id() const { return last_selected_primary_stable_id_; }
    cv::Rect last_detection_roi() const { return last_detection_roi_; }

    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;
        bool just_rebound_by_3d_fallback = false;
        float last_rebind_similarity = -1.0f;   // cosine sim from ReID (≥0), -1 = none
        float last_rebind_3d_distance = -1.0f;  // metres from 3D fallback, -1 = none
        // 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;
    cv::Rect build_detection_roi(const cv::Size& image_size) 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_;

    struct PendingRaw {
        int consecutive_hits = 0;
        int64_t last_hit_frame = -1;
    };

    std::unordered_map<int, TrackEntry> tracks_;
    std::unordered_map<int, PendingRaw> pending_raw_ids_;
    int next_stable_id_ = 0;
    int current_target_stable_id_ = -1;
    int last_selected_primary_stable_id_ = -1;
    int64_t frame_index_ = 0;
    cv::Rect last_detection_roi_;
    bool initialized_ = false;
};

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

}  // namespace odin_ros_driver