Migrate Python bindings from SWIG to nanobind

2026-03-11 21:56:30 +08:00
parent 0ec6a85921
commit d7769794fb
20 changed files with 997 additions and 243 deletions
@@ -1,35 +1,216 @@
+#include <stdexcept>
+#include <utility>
+
 #include "triangulator.hpp"
 #include "interface.hpp"

+// =================================================================================================
+
+namespace
+{
+size_t pose2d_offset(
+    size_t view,
+    size_t person,
+    size_t joint,
+    size_t coord,
+    size_t max_persons,
+    size_t num_joints)
+{
+    return ((((view * max_persons) + person) * num_joints) + joint) * 3 + coord;
+}
+
+size_t pose3d_offset(size_t person, size_t joint, size_t coord, size_t num_joints)
+{
+    return (((person * num_joints) + joint) * 4) + coord;
+}
+} // namespace
+
+// =================================================================================================
+// =================================================================================================
+
+float &PoseBatch2D::at(size_t view, size_t person, size_t joint, size_t coord)
+{
+    return data[pose2d_offset(view, person, joint, coord, max_persons, num_joints)];
+}
+
+const float &PoseBatch2D::at(size_t view, size_t person, size_t joint, size_t coord) const
+{
+    return data[pose2d_offset(view, person, joint, coord, max_persons, num_joints)];
+}
+
+RaggedPoses2D PoseBatch2D::to_nested() const
+{
+    if (person_counts.size() != num_views)
+    {
+        throw std::invalid_argument("PoseBatch2D person_counts size must match num_views.");
+    }
+
+    RaggedPoses2D poses_2d(num_views);
+    for (size_t view = 0; view < num_views; ++view)
+    {
+        const size_t num_persons = person_counts[view];
+        poses_2d[view].resize(num_persons);
+        for (size_t person = 0; person < num_persons; ++person)
+        {
+            poses_2d[view][person].resize(num_joints);
+            for (size_t joint = 0; joint < num_joints; ++joint)
+            {
+                for (size_t coord = 0; coord < 3; ++coord)
+                {
+                    poses_2d[view][person][joint][coord] = at(view, person, joint, coord);
+                }
+            }
+        }
+    }
+    return poses_2d;
+}
+
+PoseBatch2D PoseBatch2D::from_nested(const RaggedPoses2D &poses_2d)
+{
+    PoseBatch2D batch;
+    batch.num_views = poses_2d.size();
+    for (const auto &view : poses_2d)
+    {
+        batch.max_persons = std::max(batch.max_persons, view.size());
+        if (!view.empty())
+        {
+            if (batch.num_joints == 0)
+            {
+                batch.num_joints = view[0].size();
+            }
+            else if (batch.num_joints != view[0].size())
+            {
+                throw std::invalid_argument("All views must use the same joint count.");
+            }
+
+            for (const auto &person : view)
+            {
+                if (person.size() != batch.num_joints)
+                {
+                    throw std::invalid_argument("All persons must use the same joint count.");
+                }
+            }
+        }
+    }
+
+    batch.person_counts.resize(batch.num_views);
+    batch.data.assign(batch.num_views * batch.max_persons * batch.num_joints * 3, 0.0f);
+
+    for (size_t view = 0; view < batch.num_views; ++view)
+    {
+        batch.person_counts[view] = static_cast<uint32_t>(poses_2d[view].size());
+        for (size_t person = 0; person < poses_2d[view].size(); ++person)
+        {
+            for (size_t joint = 0; joint < batch.num_joints; ++joint)
+            {
+                for (size_t coord = 0; coord < 3; ++coord)
+                {
+                    batch.at(view, person, joint, coord) = poses_2d[view][person][joint][coord];
+                }
+            }
+        }
+    }
+
+    return batch;
+}
+
+// =================================================================================================
+
+float &PoseBatch3D::at(size_t person, size_t joint, size_t coord)
+{
+    return data[pose3d_offset(person, joint, coord, num_joints)];
+}
+
+const float &PoseBatch3D::at(size_t person, size_t joint, size_t coord) const
+{
+    return data[pose3d_offset(person, joint, coord, num_joints)];
+}
+
+NestedPoses3D PoseBatch3D::to_nested() const
+{
+    NestedPoses3D poses_3d(num_persons);
+    for (size_t person = 0; person < num_persons; ++person)
+    {
+        poses_3d[person].resize(num_joints);
+        for (size_t joint = 0; joint < num_joints; ++joint)
+        {
+            for (size_t coord = 0; coord < 4; ++coord)
+            {
+                poses_3d[person][joint][coord] = at(person, joint, coord);
+            }
+        }
+    }
+    return poses_3d;
+}
+
+PoseBatch3D PoseBatch3D::from_nested(const NestedPoses3D &poses_3d)
+{
+    PoseBatch3D batch;
+    batch.num_persons = poses_3d.size();
+    if (!poses_3d.empty())
+    {
+        batch.num_joints = poses_3d[0].size();
+    }
+
+    batch.data.resize(batch.num_persons * batch.num_joints * 4);
+    for (size_t person = 0; person < batch.num_persons; ++person)
+    {
+        if (poses_3d[person].size() != batch.num_joints)
+        {
+            throw std::invalid_argument("All 3D poses must use the same joint count.");
+        }
+        for (size_t joint = 0; joint < batch.num_joints; ++joint)
+        {
+            for (size_t coord = 0; coord < 4; ++coord)
+            {
+                batch.at(person, joint, coord) = poses_3d[person][joint][coord];
+            }
+        }
+    }
+    return batch;
+}
+
 // =================================================================================================
 // =================================================================================================

 Triangulator::Triangulator(float min_match_score, size_t min_group_size)
+    : triangulator(std::make_unique<TriangulatorInternal>(min_match_score, min_group_size))
 {
-    this->triangulator = new TriangulatorInternal(min_match_score, min_group_size);
 }

+Triangulator::~Triangulator() = default;
+
 // =================================================================================================

-std::vector<std::vector<std::array<float, 4>>> Triangulator::triangulate_poses(
-    const std::vector<std::vector<std::vector<std::array<float, 3>>>> &poses_2d,
+PoseBatch3D Triangulator::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)
 {
-    return this->triangulator->triangulate_poses(poses_2d, cameras, roomparams, joint_names);
+    return PoseBatch3D::from_nested(
+        triangulator->triangulate_poses(poses_2d.to_nested(), cameras, roomparams, joint_names));
+}
+
+NestedPoses3D Triangulator::triangulate_poses(
+    const RaggedPoses2D &poses_2d,
+    const std::vector<Camera> &cameras,
+    const std::array<std::array<float, 3>, 2> &roomparams,
+    const std::vector<std::string> &joint_names)
+{
+    return triangulator->triangulate_poses(poses_2d, cameras, roomparams, joint_names);
 }

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

 void Triangulator::reset()
 {
-    this->triangulator->reset();
+    triangulator->reset();
 }

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

 void Triangulator::print_stats()
 {
-    this->triangulator->print_stats();
+    triangulator->print_stats();
 }
@@ -1,5 +1,7 @@
 #pragma once

+#include <cstdint>
+#include <memory>
 #include <string>
 #include <vector>

@@ -7,7 +9,41 @@

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

-// Forward declaration of the class, that swig does not try to parse all its dependencies.
+using RaggedPoses2D = std::vector<std::vector<std::vector<std::array<float, 3>>>>;
+using NestedPoses3D = std::vector<std::vector<std::array<float, 4>>>;
+
+// =================================================================================================
+
+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;
+
+    RaggedPoses2D to_nested() 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;
+
+    NestedPoses3D to_nested() const;
+    static PoseBatch3D from_nested(const NestedPoses3D &poses_3d);
+};
+
+// =================================================================================================
+
 class TriangulatorInternal;

 // =================================================================================================
@@ -18,7 +54,6 @@ public:
    /**
     * Triangulator to predict poses from multiple views.
     *
-     *
     * @param min_match_score Minimum score to consider a triangulated joint as valid.
     * @param min_group_size Minimum number of camera pairs that need to see a person.
     */
@@ -26,19 +61,29 @@ public:
        float min_match_score = 0.95,
        size_t min_group_size = 1);

+    ~Triangulator();
+
    /**
-     * Calculate a triangulation.
+     * Calculate a triangulation using a padded pose tensor.
     *
-     *
-     * @param poses_2d List of shape [views, persons, joints, 3], containing the 2D poses.
+     * @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.
     *
-     * @return List of shape [persons, joints, 4], containing the 3D poses.
+     * @return Pose tensor of shape [persons, joints, 4].
     */
-    std::vector<std::vector<std::array<float, 4>>> triangulate_poses(
-        const std::vector<std::vector<std::vector<std::array<float, 3>>>> &poses_2d,
+    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);
+
+    /**
+     * Compatibility overload for callers that still produce ragged nested vectors.
+     */
+    NestedPoses3D triangulate_poses(
+        const RaggedPoses2D &poses_2d,
        const std::vector<Camera> &cameras,
        const std::array<std::array<float, 3>, 2> &roomparams,
        const std::vector<std::string> &joint_names);
@@ -50,5 +95,5 @@ public:
    void print_stats();

 private:
-    TriangulatorInternal *triangulator;
+    std::unique_ptr<TriangulatorInternal> triangulator;
 };