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
@@ -5,7 +5,9 @@ Various module tests
 ### Triangulator

  ```bash
-  cd /RapidPoseTriangulation/tests/ && python3 test_interface.py && cd ..
+  cd /RapidPoseTriangulation/
+  uv sync --group dev
+  uv run pytest tests/test_interface.py
  ```

 ### Onnx C++ Interface
@@ -1,19 +1,49 @@
 import json
-import sys
-import time
+from pathlib import Path

 import numpy as np
+import pytest

-sys.path.append("../swig/")
 import rpt

-# ==================================================================================================
+ROOT = Path(__file__).resolve().parents[1]
+
+JOINT_NAMES = [
+    "nose",
+    "eye_left",
+    "eye_right",
+    "ear_left",
+    "ear_right",
+    "shoulder_left",
+    "shoulder_right",
+    "elbow_left",
+    "elbow_right",
+    "wrist_left",
+    "wrist_right",
+    "hip_left",
+    "hip_right",
+    "knee_left",
+    "knee_right",
+    "ankle_left",
+    "ankle_right",
+    "hip_middle",
+    "shoulder_middle",
+    "head",
+]


-def main():
-    print("")
+def load_case(camera_path: str, pose_path: str):
+    with (ROOT / camera_path).open("r", encoding="utf-8") as file:
+        camera_data = json.load(file)
+    with (ROOT / pose_path).open("r", encoding="utf-8") as file:
+        pose_data = json.load(file)

-    # Test camera structure
+    poses_2d, person_counts = rpt.pack_poses_2d(pose_data["2D"], joint_count=len(JOINT_NAMES))
+    cameras = rpt.convert_cameras(camera_data["cameras"])
+    return poses_2d, person_counts, cameras
+
+
+def test_camera_structure_repr():
    camera = rpt.Camera()
    camera.name = "Camera 1"
    camera.K = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
@@ -22,111 +52,86 @@ def main():
    camera.T = [[1], [2], [3]]
    camera.width = 640
    camera.height = 480
-    print(camera)
-    print("")
+    camera.type = "pinhole"

-    # Load input data
-    roomparams = [[4.8, 6.0, 2.0], [0, 0, 1.0]]
-    joint_names = [
-        "nose",
-        "eye_left",
-        "eye_right",
-        "ear_left",
-        "ear_right",
-        "shoulder_left",
-        "shoulder_right",
-        "elbow_left",
-        "elbow_right",
-        "wrist_left",
-        "wrist_right",
-        "hip_left",
-        "hip_right",
-        "knee_left",
-        "knee_right",
-        "ankle_left",
-        "ankle_right",
-        "hip_middle",
-        "shoulder_middle",
-        "head",
-    ]
-    cpath = "/RapidPoseTriangulation/data/h1/sample.json"
-    ppath = "/RapidPoseTriangulation/tests/poses_h1.json"
-    with open(cpath, "r") as file:
-        cdata = json.load(file)
-    with open(ppath, "r") as file:
-        pdata = json.load(file)
-    cams = cdata["cameras"]
-    poses_2d = pdata["2D"]
-    cameras = rpt.convert_cameras(cams)
+    rendered = repr(camera)
+    assert "Camera 1" in rendered
+    assert "pinhole" in rendered

-    # Run triangulation
+
+@pytest.mark.parametrize(
+    ("camera_path", "pose_path", "roomparams"),
+    [
+        ("data/h1/sample.json", "tests/poses_h1.json", [[4.8, 6.0, 2.0], [0.0, 0.0, 1.0]]),
+        ("data/p1/sample.json", "tests/poses_p1.json", [[5.6, 6.4, 2.4], [0.0, -0.5, 1.2]]),
+        ("data/e1/sample.json", "tests/poses_e1.json", [[6.0, 5.0, 2.0], [1.5, 1.0, -0.5]]),
+    ],
+)
+def test_triangulate_samples(camera_path: str, pose_path: str, roomparams):
+    poses_2d, person_counts, cameras = load_case(camera_path, pose_path)
    triangulator = rpt.Triangulator(min_match_score=0.95)
-    stime = time.time()
    poses_3d = triangulator.triangulate_poses(
-        poses_2d, cameras, roomparams, joint_names
+        poses_2d,
+        person_counts,
+        cameras,
+        np.asarray(roomparams, dtype=np.float32),
+        JOINT_NAMES,
    )
-    print("3D time:", time.time() - stime)
-    print(np.array(poses_3d))
-    print("")

-    # Load input data
-    roomparams = [[5.6, 6.4, 2.4], [0, -0.5, 1.2]]
-    cpath = "/RapidPoseTriangulation/data/p1/sample.json"
-    ppath = "/RapidPoseTriangulation/tests/poses_p1.json"
-    with open(cpath, "r") as file:
-        cdata = json.load(file)
-    with open(ppath, "r") as file:
-        pdata = json.load(file)
-    cams = cdata["cameras"]
-    poses_2d = pdata["2D"]
-    cameras = rpt.convert_cameras(cams)
+    assert isinstance(poses_3d, np.ndarray)
+    assert poses_3d.dtype == np.float32
+    assert poses_3d.ndim == 3
+    assert poses_3d.shape[1:] == (len(JOINT_NAMES), 4)
+    assert poses_3d.shape[0] > 0
+    assert np.isfinite(poses_3d).all()

-    # Run triangulation
+
+def test_triangulate_repeatability_after_reset():
+    poses_2d, person_counts, cameras = load_case("data/p1/sample.json", "tests/poses_p1.json")
+    roomparams = np.asarray([[5.6, 6.4, 2.4], [0.0, -0.5, 1.2]], dtype=np.float32)
+    triangulator = rpt.Triangulator(min_match_score=0.95)
+
+    first = triangulator.triangulate_poses(poses_2d, person_counts, cameras, roomparams, JOINT_NAMES)
    triangulator.reset()
-    stime = time.time()
-    poses_3d = triangulator.triangulate_poses(
-        poses_2d, cameras, roomparams, joint_names
+    second = triangulator.triangulate_poses(poses_2d, person_counts, cameras, roomparams, JOINT_NAMES)
+
+    np.testing.assert_allclose(first, second, rtol=1e-5, atol=1e-5)
+
+
+def test_pack_poses_2d_from_ragged_inputs():
+    packed, counts = rpt.pack_poses_2d(
+        [
+            [[[1, 2, 0.5], [3, 4, 0.6]]],
+            np.asarray(
+                [
+                    [[5, 6, 0.7], [7, 8, 0.8]],
+                    [[9, 10, 0.9], [11, 12, 1.0]],
+                ],
+                dtype=np.float32,
+            ),
+            [],
+        ],
+        joint_count=2,
    )
-    print("3D time:", time.time() - stime)
-    print(np.array(poses_3d))
-    print("")

-    # Run again to test last pose cache
-    stime = time.time()
-    poses_3d = triangulator.triangulate_poses(
-        poses_2d, cameras, roomparams, joint_names
-    )
-    print("3D time:", time.time() - stime)
-    print(np.array(poses_3d))
-    print("")
-
-    # Load input data
-    roomparams = [[6.0, 5.0, 2.0], [1.5, 1.0, -0.5]]
-    cpath = "/RapidPoseTriangulation/data/e1/sample.json"
-    ppath = "/RapidPoseTriangulation/tests/poses_e1.json"
-    with open(cpath, "r") as file:
-        cdata = json.load(file)
-    with open(ppath, "r") as file:
-        pdata = json.load(file)
-    cams = cdata["cameras"]
-    poses_2d = pdata["2D"]
-    cameras = rpt.convert_cameras(cams)
-
-    # Run triangulation
-    triangulator.reset()
-    stime = time.time()
-    poses_3d = triangulator.triangulate_poses(
-        poses_2d, cameras, roomparams, joint_names
-    )
-    print("3D time:", time.time() - stime)
-    print(np.array(poses_3d))
-    print("")
-
-    triangulator.print_stats()
-    print("")
+    assert packed.shape == (3, 2, 2, 3)
+    assert packed.dtype == np.float32
+    np.testing.assert_array_equal(counts, np.asarray([1, 2, 0], dtype=np.uint32))
+    np.testing.assert_allclose(packed[0, 0], np.asarray([[1, 2, 0.5], [3, 4, 0.6]], dtype=np.float32))
+    np.testing.assert_allclose(packed[1, 1], np.asarray([[9, 10, 0.9], [11, 12, 1.0]], dtype=np.float32))
+    np.testing.assert_array_equal(packed[2], np.zeros((2, 2, 3), dtype=np.float32))


-# ==================================================================================================
+def test_pack_poses_2d_rejects_inconsistent_joint_count():
+    with pytest.raises(ValueError, match="joint count"):
+        rpt.pack_poses_2d(
+            [
+                [[[1, 2, 0.5], [3, 4, 0.6]]],
+                [[[5, 6, 0.7]]],
+            ]
+        )

-if __name__ == "__main__":
-    main()
+
+def test_pack_poses_2d_rejects_invalid_last_dimension():
+    with pytest.raises(ValueError, match="shape"):
+        rpt.pack_poses_2d([np.zeros((1, 2, 2), dtype=np.float32)])