Refactor triangulation stages and camera model

bindings/rpt_module.cpp

@@ -1,6 +1,8 @@
 #include <algorithm>
+#include <array>
 #include <cstdint>
 #include <stdexcept>
+#include <vector>
 
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
@@ -19,7 +21,10 @@ using PoseArray2D =
     nb::ndarray<nb::numpy, const float, nb::shape<-1, -1, -1, 3>, nb::c_contig>;
 using CountArray = nb::ndarray<nb::numpy, const uint32_t, nb::shape<-1>, nb::c_contig>;
 using RoomArray = nb::ndarray<nb::numpy, const float, nb::shape<2, 3>, nb::c_contig>;
+using PoseArray3DConst =
+    nb::ndarray<nb::numpy, const float, nb::shape<-1, -1, 4>, nb::c_contig>;
 using PoseArray3D = nb::ndarray<nb::numpy, float, nb::shape<-1, -1, 4>, nb::c_contig>;
+using PoseArray2DOut = nb::ndarray<nb::numpy, float, nb::shape<-1, 4>, nb::c_contig>;
 
 PoseBatch2DView pose_batch_view_from_numpy(const PoseArray2D &poses_2d, const CountArray &person_counts)
 {
@@ -45,6 +50,15 @@ PoseBatch2DView pose_batch_view_from_numpy(const PoseArray2D &poses_2d, const Co
     };
 }
 
+PoseBatch3DView pose_batch3d_view_from_numpy(const PoseArray3DConst &poses_3d)
+{
+    return PoseBatch3DView {
+        poses_3d.data(),
+        static_cast<size_t>(poses_3d.shape(0)),
+        static_cast<size_t>(poses_3d.shape(1)),
+    };
+}
+
 std::array<std::array<float, 3>, 2> roomparams_from_numpy(const RoomArray &roomparams)
 {
     std::array<std::array<float, 3>, 2> result {};
@@ -69,10 +83,75 @@ PoseArray3D pose_batch_to_numpy(PoseBatch3D batch)
     const size_t shape[3] = {batch.num_persons, batch.num_joints, 4};
     return PoseArray3D(storage->data(), 3, shape, owner);
 }
+
+PoseArray3D pose_batch_to_numpy_copy(const PoseBatch3D &batch)
+{
+    PoseBatch3D copy = batch;
+    return pose_batch_to_numpy(std::move(copy));
+}
+
+PoseArray2DOut pose_rows_to_numpy_copy(const std::vector<std::array<float, 4>> &rows)
+{
+    auto *storage = new std::vector<float>(rows.size() * 4, 0.0f);
+    for (size_t row = 0; row < rows.size(); ++row)
+    {
+        for (size_t coord = 0; coord < 4; ++coord)
+        {
+            (*storage)[row * 4 + coord] = rows[row][coord];
+        }
+    }
+
+    nb::capsule owner(storage, [](void *value) noexcept
+    {
+        delete static_cast<std::vector<float> *>(value);
+    });
+
+    const size_t shape[2] = {rows.size(), 4};
+    return PoseArray2DOut(storage->data(), 2, shape, owner);
+}
+
+PoseArray3D merged_poses_to_numpy_copy(const std::vector<std::vector<std::array<float, 4>>> &poses)
+{
+    size_t num_poses = poses.size();
+    size_t num_joints = 0;
+    if (!poses.empty())
+    {
+        num_joints = poses[0].size();
+    }
+
+    auto *storage = new std::vector<float>(num_poses * num_joints * 4, 0.0f);
+    for (size_t pose = 0; pose < num_poses; ++pose)
+    {
+        if (poses[pose].size() != num_joints)
+        {
+            delete storage;
+            throw std::invalid_argument("Merged poses must use a consistent joint count.");
+        }
+        for (size_t joint = 0; joint < num_joints; ++joint)
+        {
+            for (size_t coord = 0; coord < 4; ++coord)
+            {
+                (*storage)[((pose * num_joints) + joint) * 4 + coord] = poses[pose][joint][coord];
+            }
+        }
+    }
+
+    nb::capsule owner(storage, [](void *value) noexcept
+    {
+        delete static_cast<std::vector<float> *>(value);
+    });
+
+    const size_t shape[3] = {num_poses, num_joints, 4};
+    return PoseArray3D(storage->data(), 3, shape, owner);
+}
 } // namespace
 
 NB_MODULE(_core, m)
 {
+    nb::enum_<CameraModel>(m, "CameraModel")
+        .value("PINHOLE", CameraModel::Pinhole)
+        .value("FISHEYE", CameraModel::Fisheye);
+
     nb::class_<Camera>(m, "Camera")
         .def(nb::init<>())
         .def_rw("name", &Camera::name)
@@ -82,12 +161,184 @@ NB_MODULE(_core, m)
         .def_rw("T", &Camera::T)
         .def_rw("width", &Camera::width)
         .def_rw("height", &Camera::height)
-        .def_rw("type", &Camera::type)
+        .def_rw("model", &Camera::model)
         .def("__repr__", [](const Camera &camera)
         {
             return camera.to_string();
         });
 
+    nb::class_<PairCandidate>(m, "PairCandidate")
+        .def(nb::init<>())
+        .def_rw("view1", &PairCandidate::view1)
+        .def_rw("view2", &PairCandidate::view2)
+        .def_rw("person1", &PairCandidate::person1)
+        .def_rw("person2", &PairCandidate::person2)
+        .def_rw("global_person1", &PairCandidate::global_person1)
+        .def_rw("global_person2", &PairCandidate::global_person2);
+
+    nb::class_<PreviousPoseMatch>(m, "PreviousPoseMatch")
+        .def(nb::init<>())
+        .def_rw("previous_pose_index", &PreviousPoseMatch::previous_pose_index)
+        .def_rw("score_view1", &PreviousPoseMatch::score_view1)
+        .def_rw("score_view2", &PreviousPoseMatch::score_view2)
+        .def_rw("matched_view1", &PreviousPoseMatch::matched_view1)
+        .def_rw("matched_view2", &PreviousPoseMatch::matched_view2)
+        .def_rw("kept", &PreviousPoseMatch::kept)
+        .def_rw("decision", &PreviousPoseMatch::decision);
+
+    nb::class_<PreviousPoseFilterDebug>(m, "PreviousPoseFilterDebug")
+        .def(nb::init<>())
+        .def_rw("used_previous_poses", &PreviousPoseFilterDebug::used_previous_poses)
+        .def_rw("matches", &PreviousPoseFilterDebug::matches)
+        .def_rw("kept_pair_indices", &PreviousPoseFilterDebug::kept_pair_indices)
+        .def_rw("kept_pairs", &PreviousPoseFilterDebug::kept_pairs);
+
+    nb::class_<CoreProposalDebug>(m, "CoreProposalDebug")
+        .def(nb::init<>())
+        .def_rw("pair_index", &CoreProposalDebug::pair_index)
+        .def_rw("pair", &CoreProposalDebug::pair)
+        .def_rw("score", &CoreProposalDebug::score)
+        .def_rw("kept", &CoreProposalDebug::kept)
+        .def_rw("drop_reason", &CoreProposalDebug::drop_reason)
+        .def_prop_ro("pose_3d", [](const CoreProposalDebug &proposal)
+        {
+            return pose_rows_to_numpy_copy(proposal.pose_3d);
+        }, nb::rv_policy::move);
+
+    nb::class_<ProposalGroupDebug>(m, "ProposalGroupDebug")
+        .def(nb::init<>())
+        .def_rw("center", &ProposalGroupDebug::center)
+        .def_rw("proposal_indices", &ProposalGroupDebug::proposal_indices)
+        .def_prop_ro("pose_3d", [](const ProposalGroupDebug &group)
+        {
+            return pose_rows_to_numpy_copy(group.pose_3d);
+        }, nb::rv_policy::move);
+
+    nb::class_<GroupingDebug>(m, "GroupingDebug")
+        .def(nb::init<>())
+        .def_rw("initial_groups", &GroupingDebug::initial_groups)
+        .def_rw("duplicate_pair_drops", &GroupingDebug::duplicate_pair_drops)
+        .def_rw("groups", &GroupingDebug::groups);
+
+    nb::class_<FullProposalDebug>(m, "FullProposalDebug")
+        .def(nb::init<>())
+        .def_rw("source_core_proposal_index", &FullProposalDebug::source_core_proposal_index)
+        .def_rw("pair", &FullProposalDebug::pair)
+        .def_prop_ro("pose_3d", [](const FullProposalDebug &proposal)
+        {
+            return pose_rows_to_numpy_copy(proposal.pose_3d);
+        }, nb::rv_policy::move);
+
+    nb::class_<MergeDebug>(m, "MergeDebug")
+        .def(nb::init<>())
+        .def_rw("group_proposal_indices", &MergeDebug::group_proposal_indices)
+        .def_prop_ro("merged_poses", [](const MergeDebug &merge)
+        {
+            return merged_poses_to_numpy_copy(merge.merged_poses);
+        }, nb::rv_policy::move);
+
+    nb::class_<TriangulationTrace>(m, "TriangulationTrace")
+        .def(nb::init<>())
+        .def_rw("pairs", &TriangulationTrace::pairs)
+        .def_rw("previous_filter", &TriangulationTrace::previous_filter)
+        .def_rw("core_proposals", &TriangulationTrace::core_proposals)
+        .def_rw("grouping", &TriangulationTrace::grouping)
+        .def_rw("full_proposals", &TriangulationTrace::full_proposals)
+        .def_rw("merge", &TriangulationTrace::merge)
+        .def_prop_ro("final_poses", [](const TriangulationTrace &trace)
+        {
+            return pose_batch_to_numpy_copy(trace.final_poses);
+        }, nb::rv_policy::move);
+
+    m.def(
+        "build_pair_candidates",
+        [](const PoseArray2D &poses_2d, const CountArray &person_counts)
+        {
+            return build_pair_candidates(pose_batch_view_from_numpy(poses_2d, person_counts));
+        },
+        "poses_2d"_a,
+        "person_counts"_a);
+
+    m.def(
+        "filter_pairs_with_previous_poses",
+        [](const PoseArray2D &poses_2d,
+           const CountArray &person_counts,
+           const std::vector<Camera> &cameras,
+           const std::vector<std::string> &joint_names,
+           const PoseArray3DConst &previous_poses_3d,
+           float min_match_score)
+        {
+            return filter_pairs_with_previous_poses(
+                pose_batch_view_from_numpy(poses_2d, person_counts),
+                cameras,
+                joint_names,
+                pose_batch3d_view_from_numpy(previous_poses_3d),
+                min_match_score);
+        },
+        "poses_2d"_a,
+        "person_counts"_a,
+        "cameras"_a,
+        "joint_names"_a,
+        "previous_poses_3d"_a,
+        "min_match_score"_a = 0.95f);
+
+    m.def(
+        "triangulate_debug",
+        [](const PoseArray2D &poses_2d,
+           const CountArray &person_counts,
+           const std::vector<Camera> &cameras,
+           const RoomArray &roomparams,
+           const std::vector<std::string> &joint_names,
+           float min_match_score,
+           size_t min_group_size)
+        {
+            return triangulate_debug(
+                pose_batch_view_from_numpy(poses_2d, person_counts),
+                cameras,
+                roomparams_from_numpy(roomparams),
+                joint_names,
+                nullptr,
+                min_match_score,
+                min_group_size);
+        },
+        "poses_2d"_a,
+        "person_counts"_a,
+        "cameras"_a,
+        "roomparams"_a,
+        "joint_names"_a,
+        "min_match_score"_a = 0.95f,
+        "min_group_size"_a = 1);
+
+    m.def(
+        "triangulate_debug",
+        [](const PoseArray2D &poses_2d,
+           const CountArray &person_counts,
+           const std::vector<Camera> &cameras,
+           const RoomArray &roomparams,
+           const std::vector<std::string> &joint_names,
+           const PoseArray3DConst &previous_poses_3d,
+           float min_match_score,
+           size_t min_group_size)
+        {
+            const PoseBatch3DView previous_view = pose_batch3d_view_from_numpy(previous_poses_3d);
+            return triangulate_debug(
+                pose_batch_view_from_numpy(poses_2d, person_counts),
+                cameras,
+                roomparams_from_numpy(roomparams),
+                joint_names,
+                &previous_view,
+                min_match_score,
+                min_group_size);
+        },
+        "poses_2d"_a,
+        "person_counts"_a,
+        "cameras"_a,
+        "roomparams"_a,
+        "joint_names"_a,
+        "previous_poses_3d"_a,
+        "min_match_score"_a = 0.95f,
+        "min_group_size"_a = 1);
+
     m.def(
         "triangulate_poses",
         [](const PoseArray2D &poses_2d,
@@ -98,11 +349,15 @@ NB_MODULE(_core, m)
            float min_match_score,
            size_t min_group_size)
         {
-            const PoseBatch2DView pose_batch = pose_batch_view_from_numpy(poses_2d, person_counts);
-            const auto room = roomparams_from_numpy(roomparams);
-            PoseBatch3D poses_3d = triangulate_poses(
-                pose_batch, cameras, room, joint_names, min_match_score, min_group_size);
-            return pose_batch_to_numpy(std::move(poses_3d));
+            const PoseBatch3D poses_3d = triangulate_poses(
+                pose_batch_view_from_numpy(poses_2d, person_counts),
+                cameras,
+                roomparams_from_numpy(roomparams),
+                joint_names,
+                nullptr,
+                min_match_score,
+                min_group_size);
+            return pose_batch_to_numpy(poses_3d);
         },
         "poses_2d"_a,
         "person_counts"_a,
@@ -111,4 +366,35 @@ NB_MODULE(_core, m)
         "joint_names"_a,
         "min_match_score"_a = 0.95f,
         "min_group_size"_a = 1);
+
+    m.def(
+        "triangulate_poses",
+        [](const PoseArray2D &poses_2d,
+           const CountArray &person_counts,
+           const std::vector<Camera> &cameras,
+           const RoomArray &roomparams,
+           const std::vector<std::string> &joint_names,
+           const PoseArray3DConst &previous_poses_3d,
+           float min_match_score,
+           size_t min_group_size)
+        {
+            const PoseBatch3DView previous_view = pose_batch3d_view_from_numpy(previous_poses_3d);
+            const PoseBatch3D poses_3d = triangulate_poses(
+                pose_batch_view_from_numpy(poses_2d, person_counts),
+                cameras,
+                roomparams_from_numpy(roomparams),
+                joint_names,
+                &previous_view,
+                min_match_score,
+                min_group_size);
+            return pose_batch_to_numpy(poses_3d);
+        },
+        "poses_2d"_a,
+        "person_counts"_a,
+        "cameras"_a,
+        "roomparams"_a,
+        "joint_names"_a,
+        "previous_poses_3d"_a,
+        "min_match_score"_a = 0.95f,
+        "min_group_size"_a = 1);
 }