138 lines
4.2 KiB
Python
138 lines
4.2 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
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 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)
|
|
|
|
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]]
|
|
camera.DC = [0, 0, 0, 0, 0]
|
|
camera.R = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
|
|
camera.T = [[1], [2], [3]]
|
|
camera.width = 640
|
|
camera.height = 480
|
|
camera.type = "pinhole"
|
|
|
|
rendered = repr(camera)
|
|
assert "Camera 1" in rendered
|
|
assert "pinhole" in rendered
|
|
|
|
|
|
@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)
|
|
poses_3d = triangulator.triangulate_poses(
|
|
poses_2d,
|
|
person_counts,
|
|
cameras,
|
|
np.asarray(roomparams, dtype=np.float32),
|
|
JOINT_NAMES,
|
|
)
|
|
|
|
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()
|
|
|
|
|
|
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()
|
|
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,
|
|
)
|
|
|
|
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]]],
|
|
]
|
|
)
|
|
|
|
|
|
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)])
|