RapidPoseTriangulation/rpt/interface.hpp

#pragma once

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

#include "camera.hpp"

// =================================================================================================

using RaggedPoses2D = std::vector<std::vector<std::vector<std::array<float, 3>>>>;
using NestedPoses3D = std::vector<std::vector<std::array<float, 4>>>;

// =================================================================================================

struct PoseBatch2DView
{
    const float *data = nullptr;
    const uint32_t *person_counts = nullptr;
    size_t num_views = 0;
    size_t max_persons = 0;
    size_t num_joints = 0;

    const float &at(size_t view, size_t person, size_t joint, size_t coord) const;
};

struct PoseBatch3DView
{
    const float *data = nullptr;
    size_t num_persons = 0;
    size_t num_joints = 0;

    const float &at(size_t person, size_t joint, size_t coord) const;
};

struct PoseBatch2D
{
    std::vector<float> data;
    std::vector<uint32_t> person_counts;
    size_t num_views = 0;
    size_t max_persons = 0;
    size_t num_joints = 0;

    float &at(size_t view, size_t person, size_t joint, size_t coord);
    const float &at(size_t view, size_t person, size_t joint, size_t coord) const;
    PoseBatch2DView view() const;

    static PoseBatch2D from_nested(const RaggedPoses2D &poses_2d);
};

struct PoseBatch3D
{
    std::vector<float> data;
    size_t num_persons = 0;
    size_t num_joints = 0;

    float &at(size_t person, size_t joint, size_t coord);
    const float &at(size_t person, size_t joint, size_t coord) const;
    PoseBatch3DView view() const;

    NestedPoses3D to_nested() const;
    static PoseBatch3D from_nested(const NestedPoses3D &poses_3d);
};

// =================================================================================================

struct PairCandidate
{
    int view1 = -1;
    int view2 = -1;
    int person1 = -1;
    int person2 = -1;
    int global_person1 = -1;
    int global_person2 = -1;
};

struct PreviousPoseMatch
{
    int previous_pose_index = -1;
    float score_view1 = 0.0f;
    float score_view2 = 0.0f;
    bool matched_view1 = false;
    bool matched_view2 = false;
    bool kept = true;
    std::string decision = "unfiltered";
};

struct PreviousPoseFilterDebug
{
    bool used_previous_poses = false;
    std::vector<PreviousPoseMatch> matches;
    std::vector<int> kept_pair_indices;
    std::vector<PairCandidate> kept_pairs;
};

struct CoreProposalDebug
{
    int pair_index = -1;
    PairCandidate pair;
    std::vector<std::array<float, 4>> pose_3d;
    float score = 0.0f;
    bool kept = false;
    std::string drop_reason = "uninitialized";
};

struct ProposalGroupDebug
{
    std::array<float, 3> center = {0.0f, 0.0f, 0.0f};
    std::vector<std::array<float, 4>> pose_3d;
    std::vector<int> proposal_indices;
};

struct GroupingDebug
{
    std::vector<ProposalGroupDebug> initial_groups;
    std::vector<int> duplicate_pair_drops;
    std::vector<ProposalGroupDebug> groups;
};

struct FullProposalDebug
{
    int source_core_proposal_index = -1;
    PairCandidate pair;
    std::vector<std::array<float, 4>> pose_3d;
};

struct MergeDebug
{
    std::vector<std::vector<std::array<float, 4>>> merged_poses;
    std::vector<std::vector<int>> group_proposal_indices;
};

struct TriangulationTrace
{
    std::vector<PairCandidate> pairs;
    PreviousPoseFilterDebug previous_filter;
    std::vector<CoreProposalDebug> core_proposals;
    GroupingDebug grouping;
    std::vector<FullProposalDebug> full_proposals;
    MergeDebug merge;
    PoseBatch3D final_poses;
};

// =================================================================================================

struct TriangulationOptions
{
    float min_match_score = 0.95f;
    size_t min_group_size = 1;
};

// =================================================================================================

std::vector<PairCandidate> build_pair_candidates(const PoseBatch2DView &poses_2d);

PreviousPoseFilterDebug filter_pairs_with_previous_poses(
    const PoseBatch2DView &poses_2d,
    const std::vector<Camera> &cameras,
    const std::vector<std::string> &joint_names,
    const PoseBatch3DView &previous_poses_3d,
    const TriangulationOptions &options = {});

inline PreviousPoseFilterDebug filter_pairs_with_previous_poses(
    const PoseBatch2D &poses_2d,
    const std::vector<Camera> &cameras,
    const std::vector<std::string> &joint_names,
    const PoseBatch3D &previous_poses_3d,
    const TriangulationOptions &options = {})
{
    return filter_pairs_with_previous_poses(
        poses_2d.view(), cameras, joint_names, previous_poses_3d.view(), options);
}

TriangulationTrace triangulate_debug(
    const PoseBatch2DView &poses_2d,
    const std::vector<Camera> &cameras,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<std::string> &joint_names,
    const PoseBatch3DView *previous_poses_3d = nullptr,
    const TriangulationOptions &options = {});

inline TriangulationTrace triangulate_debug(
    const PoseBatch2D &poses_2d,
    const std::vector<Camera> &cameras,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<std::string> &joint_names,
    const PoseBatch3D *previous_poses_3d = nullptr,
    const TriangulationOptions &options = {})
{
    PoseBatch3DView previous_view_storage;
    const PoseBatch3DView *previous_view = nullptr;
    if (previous_poses_3d != nullptr)
    {
        previous_view_storage = previous_poses_3d->view();
        previous_view = &previous_view_storage;
    }
    return triangulate_debug(
        poses_2d.view(), cameras, roomparams, joint_names, previous_view, options);
}

// =================================================================================================

/**
 * Calculate a triangulation using a padded pose tensor.
 *
 * @param poses_2d Padded poses of shape [views, max_persons, joints, 3].
 * @param cameras List of cameras.
 * @param roomparams Room parameters (room size, room center).
 * @param joint_names List of 2D joint names.
 * @param options Triangulation options.
 *
 * @return Pose tensor of shape [persons, joints, 4].
 */
PoseBatch3D triangulate_poses(
    const PoseBatch2DView &poses_2d,
    const std::vector<Camera> &cameras,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<std::string> &joint_names,
    const PoseBatch3DView *previous_poses_3d = nullptr,
    const TriangulationOptions &options = {});

inline PoseBatch3D triangulate_poses(
    const PoseBatch2D &poses_2d,
    const std::vector<Camera> &cameras,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<std::string> &joint_names,
    const PoseBatch3D *previous_poses_3d = nullptr,
    const TriangulationOptions &options = {})
{
    PoseBatch3DView previous_view_storage;
    const PoseBatch3DView *previous_view = nullptr;
    if (previous_poses_3d != nullptr)
    {
        previous_view_storage = previous_poses_3d->view();
        previous_view = &previous_view_storage;
    }
    return triangulate_poses(
        poses_2d.view(), cameras, roomparams, joint_names, previous_view, options);
}