120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
|
|
from roboimi.utils import raw_action_trajectory_viewer as traj_view
|
|
|
|
|
|
class RawActionTrajectoryViewerTest(unittest.TestCase):
|
|
def test_set_transfer_box_pose_writes_joint_qpos(self):
|
|
joint_qpos = np.zeros(7, dtype=np.float64)
|
|
|
|
class _FakeJoint:
|
|
def __init__(self, qpos):
|
|
self.qpos = qpos
|
|
|
|
class _FakeData:
|
|
def joint(self, name):
|
|
assert name == "red_box_joint"
|
|
return _FakeJoint(joint_qpos)
|
|
|
|
traj_view.set_transfer_box_pose(_FakeData(), np.array([0.2, -0.1, 1.05], dtype=np.float64))
|
|
|
|
np.testing.assert_array_equal(
|
|
joint_qpos,
|
|
np.array([0.2, -0.1, 1.05, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
|
|
)
|
|
|
|
def test_disable_cv2_highgui_temporarily_replaces_gui_calls(self):
|
|
fake_cv2 = SimpleNamespace(
|
|
namedWindow=lambda *args, **kwargs: "named",
|
|
imshow=lambda *args, **kwargs: "imshow",
|
|
waitKey=lambda *args, **kwargs: "wait",
|
|
)
|
|
|
|
restore = traj_view.disable_cv2_highgui(fake_cv2)
|
|
self.assertIsNone(fake_cv2.namedWindow("x"))
|
|
self.assertIsNone(fake_cv2.imshow("x", None))
|
|
self.assertEqual(fake_cv2.waitKey(1), 1)
|
|
|
|
restore()
|
|
self.assertEqual(fake_cv2.namedWindow("x"), "named")
|
|
self.assertEqual(fake_cv2.imshow("x", None), "imshow")
|
|
self.assertEqual(fake_cv2.waitKey(1), "wait")
|
|
|
|
def test_load_raw_action_positions_from_npz(self):
|
|
raw_action = np.array(
|
|
[
|
|
[1.0, 2.0, 3.0, 0, 0, 0, 1, 11.0, 12.0, 13.0, 0, 0, 0, 1, -1, -1],
|
|
[4.0, 5.0, 6.0, 0, 0, 0, 1, 14.0, 15.0, 16.0, 0, 0, 0, 1, -1, -1],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path = Path(tmpdir) / "trajectory.npz"
|
|
np.savez(path, raw_action=raw_action)
|
|
|
|
positions = traj_view.load_raw_action_positions(path)
|
|
|
|
np.testing.assert_array_equal(
|
|
positions["left"],
|
|
np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32),
|
|
)
|
|
np.testing.assert_array_equal(
|
|
positions["right"],
|
|
np.array([[11.0, 12.0, 13.0], [14.0, 15.0, 16.0]], dtype=np.float32),
|
|
)
|
|
|
|
def test_build_red_capsule_segments_downsamples_to_fit_scene_limit(self):
|
|
left = np.stack([np.array([float(i), 0.0, 0.0], dtype=np.float32) for i in range(6)])
|
|
right = np.stack([np.array([float(i), 1.0, 0.0], dtype=np.float32) for i in range(6)])
|
|
|
|
markers = traj_view.build_trajectory_capsule_markers(
|
|
{"left": left, "right": right},
|
|
max_markers=4,
|
|
radius=0.01,
|
|
)
|
|
|
|
self.assertLessEqual(len(markers), 4)
|
|
self.assertTrue(all(marker["rgba"] == (1.0, 0.0, 0.0, 1.0) for marker in markers))
|
|
self.assertTrue(all(marker["radius"] == 0.01 for marker in markers))
|
|
|
|
def test_apply_capsule_markers_populates_user_scene(self):
|
|
fake_scene = SimpleNamespace(
|
|
maxgeom=3,
|
|
ngeom=99,
|
|
geoms=[object(), object(), object()],
|
|
)
|
|
markers = [
|
|
{
|
|
"from": np.array([0.0, 0.0, 0.0], dtype=np.float64),
|
|
"to": np.array([1.0, 0.0, 0.0], dtype=np.float64),
|
|
"rgba": (1.0, 0.0, 0.0, 1.0),
|
|
"radius": 0.01,
|
|
},
|
|
{
|
|
"from": np.array([0.0, 1.0, 0.0], dtype=np.float64),
|
|
"to": np.array([1.0, 1.0, 0.0], dtype=np.float64),
|
|
"rgba": (1.0, 0.0, 0.0, 1.0),
|
|
"radius": 0.01,
|
|
},
|
|
]
|
|
|
|
with mock.patch.object(traj_view.mujoco, "mjv_initGeom") as init_geom, mock.patch.object(
|
|
traj_view.mujoco,
|
|
"mjv_connector",
|
|
) as connector:
|
|
traj_view.apply_capsule_markers_to_scene(fake_scene, markers)
|
|
|
|
self.assertEqual(fake_scene.ngeom, 2)
|
|
self.assertEqual(init_geom.call_count, 2)
|
|
self.assertEqual(connector.call_count, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|