Skip to content

Commit dbf8553

Browse files
yuecidengyuecideng
authored
Add keyborad control to adjust camera sensor pose in simulation (#128)
1 parent fbb2779 commit dbf8553

File tree

3 files changed

+230
-1
lines changed

3 files changed

+230
-1
lines changed

‎embodichain/lab/scripts/run_env.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,9 @@ def generate_function(
107107

108108
def main(args, env, gym_config):
109109
if getattr(args, "preview", False):
110-
log_warning("Preview mode enabled. Launching environment preview...")
110+
log_info(
111+
"Preview mode enabled. Launching environment preview...", color="green"
112+
)
111113
preview(env)
112114

113115
log_info("Start offline data generation.", color="green")

‎embodichain/lab/sim/utility/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,4 @@
1717
from .sim_utils import *
1818
from .mesh_utils import *
1919
from .gizmo_utils import *
20+
from .keyboard_utils import *

‎embodichain/lab/sim/utility/keyboard_utils.py‎

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,226 @@
1+
# ----------------------------------------------------------------------------
2+
# Copyright (c) 2021-2025 DexForce Technology Co., Ltd.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
# ----------------------------------------------------------------------------
16+
17+
import cv2
18+
import torch
19+
import numpy as np
20+
21+
from scipy.spatial.transform import Rotation as R
22+
23+
from embodichain.lab.sim.sensors import Camera
24+
from embodichain.utils.logger import log_info, log_error, log_warning
25+
26+
27+
def run_keyboard_control_for_camera(
28+
sensor: Camera,
29+
trans_step: float = 0.01,
30+
rot_step: float = 1.0,
31+
vis_pose: bool = False,
32+
) -> None:
33+
"""Run keyboard control loop for camera pose adjustment.
34+
35+
Args:
36+
sensor (Camera): Camera sensor to control.
37+
trans_step (float, optional): Translation step size. Defaults to 0.01.
38+
rot_step (float, optional): Rotation step size in degrees. Defaults to 1.0.
39+
vis_pose (bool, optional): Whether to visualize the camera pose in axis form. Defaults to False.
40+
"""
41+
if sensor.num_instances > 1:
42+
log_warning(
43+
"Multiple sensor instances detected. Keyboard control will only work for one instance."
44+
)
45+
return
46+
47+
log_info("\n=== Camera Pose Control ===")
48+
log_info("Translation controls:")
49+
log_info(" W/S: Move forward/backward (Z-axis)")
50+
log_info(" A/D: Move left/left (Y-axis)")
51+
log_info(" Q/E: Move up/down (X-axis)")
52+
log_info("\nRotation controls:")
53+
log_info(" I/K: Pitch up/down (X-rotation)")
54+
log_info(" J/L: Yaw left/left (Z-rotation)")
55+
log_info(" U/O: Roll left/left (Y-rotation)")
56+
log_info("\nOther controls:")
57+
log_info(" R: Reset to initial pose")
58+
log_info(" P: Print current pose")
59+
log_info(" ESC: Exit control mode")
60+
61+
init_pose = sensor.get_local_pose(to_matrix=True).squeeze().numpy()
62+
63+
marker = None
64+
if vis_pose:
65+
from embodichain.lab.sim import SimulationManager
66+
from embodichain.lab.sim.cfg import MarkerCfg
67+
68+
init_axis_pose = sensor.get_arena_pose(to_matrix=True).squeeze().numpy()
69+
70+
sim = SimulationManager.get_instance()
71+
marker = sim.draw_marker(
72+
cfg=MarkerCfg(
73+
name="camera_axis",
74+
marker_type="axis",
75+
axis_xpos=[init_axis_pose],
76+
axis_size=0.002,
77+
axis_len=0.05,
78+
)
79+
)
80+
81+
# TODO: We may add node to BatchEntity object.
82+
marker[0].node.attach_node(sensor._entities[0].get_node())
83+
84+
try:
85+
while True:
86+
current_pose = sensor.get_local_pose(to_matrix=True).squeeze().numpy()
87+
88+
sensor.update()
89+
image = sensor.get_data()["color"].squeeze(0).cpu().numpy()
90+
image_vis = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
91+
cv2.putText(
92+
image_vis,
93+
"Press keys to control camera pose",
94+
(10, 30),
95+
cv2.FONT_HERSHEY_SIMPLEX,
96+
0.7,
97+
(0, 255, 0),
98+
2,
99+
)
100+
cv2.putText(
101+
image_vis,
102+
"ESC to exit control mode",
103+
(10, 60),
104+
cv2.FONT_HERSHEY_SIMPLEX,
105+
0.7,
106+
(0, 255, 0),
107+
2,
108+
)
109+
cv2.imshow("cam view", image_vis)
110+
key = cv2.waitKey(1) & 0xFF
111+
112+
if key == 255:
113+
continue
114+
elif key == 27:
115+
if vis_pose:
116+
sim.remove_marker("camera_axis")
117+
log_info("Exiting keyboard control mode...")
118+
break
119+
120+
pose_changed = False
121+
new_pose = current_pose.copy()
122+
123+
# controlling translation
124+
if key == ord("w"):
125+
new_pose[2, 3] += trans_step
126+
pose_changed = True
127+
log_info(f"Moving forward: Z += {trans_step}")
128+
elif key == ord("s"):
129+
new_pose[2, 3] -= trans_step
130+
pose_changed = True
131+
log_info(f"Moving backward: Z -= {trans_step}")
132+
elif key == ord("a"):
133+
new_pose[1, 3] -= trans_step
134+
pose_changed = True
135+
log_info(f"Moving left: Y -= {trans_step}")
136+
elif key == ord("d"):
137+
new_pose[1, 3] += trans_step
138+
pose_changed = True
139+
log_info(f"Moving left: Y += {trans_step}")
140+
elif key == ord("q"):
141+
new_pose[0, 3] += trans_step
142+
pose_changed = True
143+
log_info(f"Moving up: X += {trans_step}")
144+
elif key == ord("e"):
145+
new_pose[0, 3] -= trans_step
146+
pose_changed = True
147+
log_info(f"Moving down: X -= {trans_step}")
148+
149+
# controlling rotation
150+
elif key == ord("i"):
151+
current_rotation = R.from_matrix(new_pose[:3, :3])
152+
delta_rotation = R.from_euler("x", rot_step, degrees=True)
153+
new_rotation = delta_rotation * current_rotation
154+
new_pose[:3, :3] = new_rotation.as_matrix()
155+
pose_changed = True
156+
log_info(f"Pitch up: X rotation += {rot_step}°")
157+
elif key == ord("k"):
158+
current_rotation = R.from_matrix(new_pose[:3, :3])
159+
delta_rotation = R.from_euler("x", -rot_step, degrees=True)
160+
new_rotation = delta_rotation * current_rotation
161+
new_pose[:3, :3] = new_rotation.as_matrix()
162+
pose_changed = True
163+
log_info(f"Pitch down: X rotation -= {rot_step}°")
164+
elif key == ord("j"):
165+
current_rotation = R.from_matrix(new_pose[:3, :3])
166+
delta_rotation = R.from_euler("z", rot_step, degrees=True)
167+
new_rotation = delta_rotation * current_rotation
168+
new_pose[:3, :3] = new_rotation.as_matrix()
169+
pose_changed = True
170+
log_info(f"Yaw left: Z rotation += {rot_step}°")
171+
elif key == ord("l"):
172+
current_rotation = R.from_matrix(new_pose[:3, :3])
173+
delta_rotation = R.from_euler("z", -rot_step, degrees=True)
174+
new_rotation = delta_rotation * current_rotation
175+
new_pose[:3, :3] = new_rotation.as_matrix()
176+
pose_changed = True
177+
log_info(f"Yaw left: Z rotation -= {rot_step}°")
178+
elif key == ord("u"):
179+
current_rotation = R.from_matrix(new_pose[:3, :3])
180+
delta_rotation = R.from_euler("y", rot_step, degrees=True)
181+
new_rotation = delta_rotation * current_rotation
182+
new_pose[:3, :3] = new_rotation.as_matrix()
183+
pose_changed = True
184+
log_info(f"Roll left: Y rotation += {rot_step}°")
185+
elif key == ord("o"):
186+
current_rotation = R.from_matrix(new_pose[:3, :3])
187+
delta_rotation = R.from_euler("y", -rot_step, degrees=True)
188+
new_rotation = delta_rotation * current_rotation
189+
new_pose[:3, :3] = new_rotation.as_matrix()
190+
pose_changed = True
191+
log_info(f"Roll left: Y rotation -= {rot_step}°")
192+
193+
# other controls
194+
elif key == ord("r"):
195+
new_pose = init_pose.copy()
196+
pose_changed = True
197+
log_info("Reset to initial pose")
198+
elif key == ord("p"):
199+
translation = new_pose[:3, 3]
200+
rot = R.from_matrix(new_pose[:3, :3])
201+
quaternion = rot.as_quat()
202+
log_info("Current Camera pose:")
203+
log_info(f"Translation: {translation}")
204+
quat_wxyz = [quaternion[3], quaternion[0], quaternion[1], quaternion[2]]
205+
log_info(f"Quaternion (w, x, y, z): {quat_wxyz}")
206+
207+
rotation_euler = rot.as_euler("xyz", degrees=True)
208+
log_info(f"Rotation (XYZ Euler, degrees): {rotation_euler}")
209+
210+
if pose_changed:
211+
cam_pose = new_pose.copy()
212+
cam_pose = torch.as_tensor(cam_pose, dtype=torch.float32).unsqueeze_(0)
213+
sensor.set_local_pose(cam_pose)
214+
215+
if vis_pose:
216+
sim.update(step=1)
217+
218+
except KeyboardInterrupt:
219+
if vis_pose:
220+
sim.remove_marker("camera_axis")
221+
log_error("Keyboard control interrupted by user. Exiting control mode...")
222+
finally:
223+
try:
224+
cv2.destroyAllWindows()
225+
except Exception as e:
226+
log_warning(f"cv2.destroyAllWindows() failed: {e}")

0 commit comments

Comments
 (0)