Skip to content

Commit d4675a1

Browse files
yuecidengyuecideng
committed
wip
1 parent 9db1183 commit d4675a1

File tree

3 files changed

+190
-1
lines changed

3 files changed

+190
-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: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
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+
20+
from scipy.spatial.transform import Rotation as R
21+
22+
from embodichain.lab.sim.sensors import Camera
23+
from embodichain.utils.logger import log_info, log_error, log_warning
24+
25+
26+
def run_keyboard_control_for_camera(
27+
sensor: Camera, trans_step: float = 0.01, rot_step: float = 1.0
28+
) -> None:
29+
if sensor.num_instances > 1:
30+
log_warning(
31+
"Multiple sensor instances detected. Keyboard control will only work for one instance."
32+
)
33+
return
34+
35+
log_info("\n=== Camera Pose Control ===")
36+
log_info("Translation controls:")
37+
log_info(" W/S: Move forward/backward (Z-axis)")
38+
log_info(" A/D: Move left/left (Y-axis)")
39+
log_info(" Q/E: Move up/down (X-axis)")
40+
log_info("\nRotation controls:")
41+
log_info(" I/K: Pitch up/down (X-rotation)")
42+
log_info(" J/L: Yaw left/left (Z-rotation)")
43+
log_info(" U/O: Roll left/left (Y-rotation)")
44+
log_info("\nOther controls:")
45+
log_info(" R: Reset to initial pose")
46+
log_info(" P: Print current pose")
47+
log_info(" ESC: Exit control mode")
48+
49+
init_pose = sensor.get_local_pose(to_matrix=True).squeeze().numpy()
50+
51+
try:
52+
while True:
53+
current_pose = sensor.get_local_pose(to_matrix=True).squeeze().numpy()
54+
55+
sensor.update()
56+
image = sensor.get_data()["color"].squeeze(0).cpu().numpy()
57+
image_vis = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
58+
cv2.putText(
59+
image_vis,
60+
"Press keys to control camera pose",
61+
(10, 30),
62+
cv2.FONT_HERSHEY_SIMPLEX,
63+
0.7,
64+
(0, 255, 0),
65+
2,
66+
)
67+
cv2.putText(
68+
image_vis,
69+
"ESC to exit control mode",
70+
(10, 60),
71+
cv2.FONT_HERSHEY_SIMPLEX,
72+
0.7,
73+
(0, 255, 0),
74+
2,
75+
)
76+
cv2.imshow("cam view", image_vis)
77+
key = cv2.waitKey(1) & 0xFF
78+
79+
if key == 255:
80+
continue
81+
elif key == 27:
82+
log_info("Exiting keyboard control mode...")
83+
break
84+
85+
pose_changed = False
86+
new_pose = current_pose.copy()
87+
88+
# controlling translation
89+
if key == ord("w"):
90+
new_pose[2, 3] += trans_step
91+
pose_changed = True
92+
log_info(f"Moving forward: Z += {trans_step}")
93+
elif key == ord("s"):
94+
new_pose[2, 3] -= trans_step
95+
pose_changed = True
96+
log_info(f"Moving backward: Z -= {trans_step}")
97+
elif key == ord("a"):
98+
new_pose[1, 3] -= trans_step
99+
pose_changed = True
100+
log_info(f"Moving left: Y -= {trans_step}")
101+
elif key == ord("d"):
102+
new_pose[1, 3] += trans_step
103+
pose_changed = True
104+
log_info(f"Moving left: Y += {trans_step}")
105+
elif key == ord("q"):
106+
new_pose[0, 3] += trans_step
107+
pose_changed = True
108+
log_info(f"Moving up: X += {trans_step}")
109+
elif key == ord("e"):
110+
new_pose[0, 3] -= trans_step
111+
pose_changed = True
112+
log_info(f"Moving down: X -= {trans_step}")
113+
114+
# controlling rotation
115+
elif key == ord("i"):
116+
current_rotation = R.from_matrix(new_pose[:3, :3])
117+
delta_rotation = R.from_euler("x", rot_step, degrees=True)
118+
new_rotation = delta_rotation * current_rotation
119+
new_pose[:3, :3] = new_rotation.as_matrix()
120+
pose_changed = True
121+
log_info(f"Pitch up: X rotation += {rot_step}°")
122+
elif key == ord("k"):
123+
current_rotation = R.from_matrix(new_pose[:3, :3])
124+
delta_rotation = R.from_euler("x", -rot_step, degrees=True)
125+
new_rotation = delta_rotation * current_rotation
126+
new_pose[:3, :3] = new_rotation.as_matrix()
127+
pose_changed = True
128+
log_info(f"Pitch down: X rotation -= {rot_step}°")
129+
elif key == ord("j"):
130+
current_rotation = R.from_matrix(new_pose[:3, :3])
131+
delta_rotation = R.from_euler("z", rot_step, degrees=True)
132+
new_rotation = delta_rotation * current_rotation
133+
new_pose[:3, :3] = new_rotation.as_matrix()
134+
pose_changed = True
135+
log_info(f"Yaw left: Z rotation += {rot_step}°")
136+
elif key == ord("l"):
137+
current_rotation = R.from_matrix(new_pose[:3, :3])
138+
delta_rotation = R.from_euler("z", -rot_step, degrees=True)
139+
new_rotation = delta_rotation * current_rotation
140+
new_pose[:3, :3] = new_rotation.as_matrix()
141+
pose_changed = True
142+
log_info(f"Yaw left: Z rotation -= {rot_step}°")
143+
elif key == ord("u"):
144+
current_rotation = R.from_matrix(new_pose[:3, :3])
145+
delta_rotation = R.from_euler("y", rot_step, degrees=True)
146+
new_rotation = delta_rotation * current_rotation
147+
new_pose[:3, :3] = new_rotation.as_matrix()
148+
pose_changed = True
149+
log_info(f"Roll left: Y rotation += {rot_step}°")
150+
elif key == ord("o"):
151+
current_rotation = R.from_matrix(new_pose[:3, :3])
152+
delta_rotation = R.from_euler("y", -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"Roll left: Y rotation -= {rot_step}°")
157+
158+
# other controls
159+
elif key == ord("r"):
160+
new_pose = init_pose.copy()
161+
pose_changed = True
162+
log_info("Reset to initial pose")
163+
elif key == ord("p"):
164+
translation = new_pose[:3, 3]
165+
rot = R.from_matrix(new_pose[:3, :3])
166+
quaternion = rot.as_quat()
167+
log_info("Current Camera pose:")
168+
log_info(f"Translation: {translation}")
169+
quat_wxyz = [quaternion[3], quaternion[0], quaternion[1], quaternion[2]]
170+
log_info(f"Quaternion (w, x, y, z): {quat_wxyz}")
171+
172+
rotation_euler = rot.as_euler("xyz", degrees=True)
173+
log_info(f"Rotation (XYZ Euler, degrees): {rotation_euler}")
174+
175+
if pose_changed:
176+
cam_pose = new_pose.copy()
177+
cam_pose = torch.as_tensor(cam_pose, dtype=torch.float32).unsqueeze_(0)
178+
sensor.set_local_pose(cam_pose)
179+
180+
except KeyboardInterrupt:
181+
log_error("Keyboard control interrupted by user. Exiting control mode...")
182+
finally:
183+
try:
184+
cv2.destroyAllWindows()
185+
except Exception as e:
186+
log_warning(f"cv2.destroyAllWindows() failed: {e}")

0 commit comments

Comments
 (0)