crosstyan/RapidPoseTriangulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make triangulation a zero-copy pure function

Browse Source
This commit is contained in:
crosstyan
2026-03-11 22:29:21 +08:00
parent 5bed0f0aaf
commit 24f74c87f1
10 changed files with 596 additions and 947 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bindings/rpt_module.cpp
+33 -40
View File
@@ -1,7 +1,5 @@
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <nanobind/nanobind.h>
@@ -23,31 +21,28 @@ using CountArray = nb::ndarray<nb::numpy, const uint32_t, nb::shape<-1>, nb::c_c
using RoomArray = nb::ndarray<nb::numpy, const float, nb::shape<2, 3>, nb::c_contig>;
using PoseArray3D = nb::ndarray<nb::numpy, float, nb::shape<-1, -1, 4>, nb::c_contig>;
PoseBatch2D pose_batch_from_numpy(const PoseArray2D &poses_2d, const CountArray &person_counts)
PoseBatch2DView pose_batch_view_from_numpy(const PoseArray2D &poses_2d, const CountArray &person_counts)
{
if (poses_2d.shape(0) != person_counts.shape(0))
{
throw std::invalid_argument("poses_2d and person_counts must have the same number of views.");
}
PoseBatch2D batch;
batch.num_views = static_cast<size_t>(poses_2d.shape(0));
batch.max_persons = static_cast<size_t>(poses_2d.shape(1));
batch.num_joints = static_cast<size_t>(poses_2d.shape(2));
batch.person_counts.assign(person_counts.data(), person_counts.data() + batch.num_views);
for (size_t i = 0; i < batch.person_counts.size(); ++i)
for (size_t i = 0; i < static_cast<size_t>(person_counts.shape(0)); ++i)
{
if (batch.person_counts[i] > batch.max_persons)
if (person_counts(i) > poses_2d.shape(1))
{
throw std::invalid_argument("person_counts entries must not exceed the padded person dimension.");
}
}
const size_t total_size = batch.num_views * batch.max_persons * batch.num_joints * 3;
batch.data.resize(total_size);
std::memcpy(batch.data.data(), poses_2d.data(), total_size * sizeof(float));
return batch;
return PoseBatch2DView {
poses_2d.data(),
person_counts.data(),
static_cast<size_t>(poses_2d.shape(0)),
static_cast<size_t>(poses_2d.shape(1)),
static_cast<size_t>(poses_2d.shape(2)),
};
}
std::array<std::array<float, 3>, 2> roomparams_from_numpy(const RoomArray &roomparams)
@@ -93,29 +88,27 @@ NB_MODULE(_core, m)
return camera.to_string();
});
nb::class_<Triangulator>(m, "Triangulator")
.def(nb::init<float, size_t>(),
"min_match_score"_a = 0.95f,
"min_group_size"_a = 1)
.def(
"triangulate_poses",
[](Triangulator &self,
const PoseArray2D &poses_2d,
const CountArray &person_counts,
const std::vector<Camera> &cameras,
const RoomArray &roomparams,
const std::vector<std::string> &joint_names)
{
PoseBatch2D pose_batch = pose_batch_from_numpy(poses_2d, person_counts);
auto room = roomparams_from_numpy(roomparams);
PoseBatch3D poses_3d = self.triangulate_poses(pose_batch, cameras, room, joint_names);
return pose_batch_to_numpy(std::move(poses_3d));
},
"poses_2d"_a,
"person_counts"_a,
"cameras"_a,
"roomparams"_a,
"joint_names"_a)
.def("reset", &Triangulator::reset)
.def("print_stats", &Triangulator::print_stats);
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,
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));
},
"poses_2d"_a,
"person_counts"_a,
"cameras"_a,
"roomparams"_a,
"joint_names"_a,
"min_match_score"_a = 0.95f,
"min_group_size"_a = 1);
}