Clean up nanobind typing and source layout

2026-03-27 11:19:53 +08:00
parent 9d63177de0
commit 69e83d8430
19 changed files with 792 additions and 47 deletions
@@ -0,0 +1,311 @@
+#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 TrackedPoseBatch3DView
+{
+    const int64_t *track_ids = nullptr;
+    const float *data = nullptr;
+    size_t num_persons = 0;
+    size_t num_joints = 0;
+
+    int64_t track_id(size_t person) const;
+    const float &at(size_t person, size_t joint, size_t coord) const;
+};
+
+struct PoseBatch3DByViewView
+{
+    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 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 PoseBatch3DByView
+{
+    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;
+    PoseBatch3DByViewView view() const;
+};
+
+// =================================================================================================
+
+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;
+    int64_t previous_track_id = -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;
+};
+
+enum class AssociationStatus
+{
+    Matched,
+    New,
+    Ambiguous,
+};
+
+struct AssociationReport
+{
+    std::vector<int> pose_previous_indices;
+    std::vector<int64_t> pose_previous_track_ids;
+    std::vector<AssociationStatus> pose_status;
+    std::vector<std::vector<int>> pose_candidate_previous_indices;
+    std::vector<std::vector<int64_t>> pose_candidate_previous_track_ids;
+    std::vector<int> unmatched_previous_indices;
+    std::vector<int64_t> unmatched_previous_track_ids;
+    std::vector<int> new_pose_indices;
+    std::vector<int> ambiguous_pose_indices;
+};
+
+struct FinalPoseAssociationDebug
+{
+    int final_pose_index = -1;
+    std::vector<int> source_core_proposal_indices;
+    std::vector<int> source_pair_indices;
+    std::vector<int> candidate_previous_indices;
+    std::vector<int64_t> candidate_previous_track_ids;
+    int resolved_previous_index = -1;
+    int64_t resolved_previous_track_id = -1;
+    AssociationStatus status = AssociationStatus::New;
+};
+
+struct TriangulationTrace
+{
+    std::vector<PairCandidate> pairs;
+    PreviousPoseFilterDebug previous_filter;
+    std::vector<CoreProposalDebug> core_proposals;
+    GroupingDebug grouping;
+    std::vector<FullProposalDebug> full_proposals;
+    MergeDebug merge;
+    AssociationReport association;
+    std::vector<FinalPoseAssociationDebug> final_pose_associations;
+    PoseBatch3D final_poses;
+};
+
+struct TriangulationResult
+{
+    PoseBatch3D poses;
+    AssociationReport association;
+};
+
+// =================================================================================================
+
+struct TriangulationOptions
+{
+    float min_match_score = 0.95f;
+    size_t min_group_size = 1;
+};
+
+struct TriangulationConfig
+{
+    std::vector<Camera> cameras;
+    std::array<std::array<float, 3>, 2> roomparams {{{0.0f, 0.0f, 0.0f}, {0.0f, 0.0f, 0.0f}}};
+    std::vector<std::string> joint_names;
+    TriangulationOptions options;
+};
+
+// =================================================================================================
+
+std::vector<PairCandidate> build_pair_candidates(const PoseBatch2DView &poses_2d);
+
+PreviousPoseFilterDebug filter_pairs_with_previous_poses(
+    const PoseBatch2DView &poses_2d,
+    const TriangulationConfig &config,
+    const TrackedPoseBatch3DView &previous_poses_3d,
+    const TriangulationOptions *options_override = nullptr);
+
+TriangulationTrace triangulate_debug(
+    const PoseBatch2DView &poses_2d,
+    const TriangulationConfig &config,
+    const TrackedPoseBatch3DView *previous_poses_3d = nullptr,
+    const TriangulationOptions *options_override = nullptr);
+
+// =================================================================================================
+
+/**
+ * Calculate a triangulation using a padded pose tensor.
+ *
+ * @param poses_2d Padded poses of shape [views, max_persons, joints, 3].
+ * @param config Triangulation configuration (cameras, room parameters, joint names, options).
+ * @param options_override Optional per-call options override. Defaults to config.options.
+ *
+ * @return Pose tensor of shape [persons, joints, 4].
+ */
+PoseBatch3D triangulate_poses(
+    const PoseBatch2DView &poses_2d,
+    const TriangulationConfig &config,
+    const TriangulationOptions *options_override = nullptr);
+
+PoseBatch3D merge_rgbd_views(
+    const PoseBatch3DByViewView &poses_3d,
+    const TriangulationConfig &config,
+    float max_distance = 0.5f);
+
+TriangulationResult triangulate_with_report(
+    const PoseBatch2DView &poses_2d,
+    const TriangulationConfig &config,
+    const TrackedPoseBatch3DView &previous_poses_3d,
+    const TriangulationOptions *options_override = nullptr);
+
+inline PoseBatch3D triangulate_poses(
+    const PoseBatch2D &poses_2d,
+    const TriangulationConfig &config,
+    const TriangulationOptions *options_override = nullptr)
+{
+    return triangulate_poses(poses_2d.view(), config, options_override);
+}
+
+inline PoseBatch3D merge_rgbd_views(
+    const PoseBatch3DByView &poses_3d,
+    const TriangulationConfig &config,
+    float max_distance = 0.5f)
+{
+    return merge_rgbd_views(poses_3d.view(), config, max_distance);
+}
+
+inline TriangulationTrace triangulate_debug(
+    const PoseBatch2D &poses_2d,
+    const TriangulationConfig &config,
+    const TriangulationOptions *options_override = nullptr)
+{
+    return triangulate_debug(poses_2d.view(), config, nullptr, options_override);
+}
+
+inline TriangulationResult triangulate_with_report(
+    const PoseBatch2D &poses_2d,
+    const TriangulationConfig &config,
+    const TrackedPoseBatch3DView &previous_poses_3d,
+    const TriangulationOptions *options_override = nullptr)
+{
+    return triangulate_with_report(poses_2d.view(), config, previous_poses_3d, options_override);
+}