diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 858950af..c564e40a 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -107,7 +107,9 @@ def generate_function( def main(args, env, gym_config): if getattr(args, "preview", False): - log_warning("Preview mode enabled. Launching environment preview...") + log_info( + "Preview mode enabled. Launching environment preview...", color="green" + ) preview(env) log_info("Start offline data generation.", color="green") diff --git a/embodichain/lab/sim/utility/__init__.py b/embodichain/lab/sim/utility/__init__.py index 94cbbc59..dd671cba 100644 --- a/embodichain/lab/sim/utility/__init__.py +++ b/embodichain/lab/sim/utility/__init__.py @@ -17,3 +17,4 @@ from .sim_utils import * from .mesh_utils import * from .gizmo_utils import * +from .keyboard_utils import * diff --git a/embodichain/lab/sim/utility/keyboard_utils.py b/embodichain/lab/sim/utility/keyboard_utils.py new file mode 100644 index 00000000..f5997c03 --- /dev/null +++ b/embodichain/lab/sim/utility/keyboard_utils.py @@ -0,0 +1,226 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import cv2 +import torch +import numpy as np + +from scipy.spatial.transform import Rotation as R + +from embodichain.lab.sim.sensors import Camera +from embodichain.utils.logger import log_info, log_error, log_warning + + +def run_keyboard_control_for_camera( + sensor: Camera, + trans_step: float = 0.01, + rot_step: float = 1.0, + vis_pose: bool = False, +) -> None: + """Run keyboard control loop for camera pose adjustment. + + Args: + sensor (Camera): Camera sensor to control. + trans_step (float, optional): Translation step size. Defaults to 0.01. + rot_step (float, optional): Rotation step size in degrees. Defaults to 1.0. + vis_pose (bool, optional): Whether to visualize the camera pose in axis form. Defaults to False. + """ + if sensor.num_instances > 1: + log_warning( + "Multiple sensor instances detected. Keyboard control will only work for one instance." + ) + return + + log_info("\n=== Camera Pose Control ===") + log_info("Translation controls:") + log_info(" W/S: Move forward/backward (Z-axis)") + log_info(" A/D: Move left/left (Y-axis)") + log_info(" Q/E: Move up/down (X-axis)") + log_info("\nRotation controls:") + log_info(" I/K: Pitch up/down (X-rotation)") + log_info(" J/L: Yaw left/left (Z-rotation)") + log_info(" U/O: Roll left/left (Y-rotation)") + log_info("\nOther controls:") + log_info(" R: Reset to initial pose") + log_info(" P: Print current pose") + log_info(" ESC: Exit control mode") + + init_pose = sensor.get_local_pose(to_matrix=True).squeeze().numpy() + + marker = None + if vis_pose: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.cfg import MarkerCfg + + init_axis_pose = sensor.get_arena_pose(to_matrix=True).squeeze().numpy() + + sim = SimulationManager.get_instance() + marker = sim.draw_marker( + cfg=MarkerCfg( + name="camera_axis", + marker_type="axis", + axis_xpos=[init_axis_pose], + axis_size=0.002, + axis_len=0.05, + ) + ) + + # TODO: We may add node to BatchEntity object. + marker[0].node.attach_node(sensor._entities[0].get_node()) + + try: + while True: + current_pose = sensor.get_local_pose(to_matrix=True).squeeze().numpy() + + sensor.update() + image = sensor.get_data()["color"].squeeze(0).cpu().numpy() + image_vis = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + cv2.putText( + image_vis, + "Press keys to control camera pose", + (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (0, 255, 0), + 2, + ) + cv2.putText( + image_vis, + "ESC to exit control mode", + (10, 60), + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (0, 255, 0), + 2, + ) + cv2.imshow("cam view", image_vis) + key = cv2.waitKey(1) & 0xFF + + if key == 255: + continue + elif key == 27: + if vis_pose: + sim.remove_marker("camera_axis") + log_info("Exiting keyboard control mode...") + break + + pose_changed = False + new_pose = current_pose.copy() + + # controlling translation + if key == ord("w"): + new_pose[2, 3] += trans_step + pose_changed = True + log_info(f"Moving forward: Z += {trans_step}") + elif key == ord("s"): + new_pose[2, 3] -= trans_step + pose_changed = True + log_info(f"Moving backward: Z -= {trans_step}") + elif key == ord("a"): + new_pose[1, 3] -= trans_step + pose_changed = True + log_info(f"Moving left: Y -= {trans_step}") + elif key == ord("d"): + new_pose[1, 3] += trans_step + pose_changed = True + log_info(f"Moving left: Y += {trans_step}") + elif key == ord("q"): + new_pose[0, 3] += trans_step + pose_changed = True + log_info(f"Moving up: X += {trans_step}") + elif key == ord("e"): + new_pose[0, 3] -= trans_step + pose_changed = True + log_info(f"Moving down: X -= {trans_step}") + + # controlling rotation + elif key == ord("i"): + current_rotation = R.from_matrix(new_pose[:3, :3]) + delta_rotation = R.from_euler("x", rot_step, degrees=True) + new_rotation = delta_rotation * current_rotation + new_pose[:3, :3] = new_rotation.as_matrix() + pose_changed = True + log_info(f"Pitch up: X rotation += {rot_step}°") + elif key == ord("k"): + current_rotation = R.from_matrix(new_pose[:3, :3]) + delta_rotation = R.from_euler("x", -rot_step, degrees=True) + new_rotation = delta_rotation * current_rotation + new_pose[:3, :3] = new_rotation.as_matrix() + pose_changed = True + log_info(f"Pitch down: X rotation -= {rot_step}°") + elif key == ord("j"): + current_rotation = R.from_matrix(new_pose[:3, :3]) + delta_rotation = R.from_euler("z", rot_step, degrees=True) + new_rotation = delta_rotation * current_rotation + new_pose[:3, :3] = new_rotation.as_matrix() + pose_changed = True + log_info(f"Yaw left: Z rotation += {rot_step}°") + elif key == ord("l"): + current_rotation = R.from_matrix(new_pose[:3, :3]) + delta_rotation = R.from_euler("z", -rot_step, degrees=True) + new_rotation = delta_rotation * current_rotation + new_pose[:3, :3] = new_rotation.as_matrix() + pose_changed = True + log_info(f"Yaw left: Z rotation -= {rot_step}°") + elif key == ord("u"): + current_rotation = R.from_matrix(new_pose[:3, :3]) + delta_rotation = R.from_euler("y", rot_step, degrees=True) + new_rotation = delta_rotation * current_rotation + new_pose[:3, :3] = new_rotation.as_matrix() + pose_changed = True + log_info(f"Roll left: Y rotation += {rot_step}°") + elif key == ord("o"): + current_rotation = R.from_matrix(new_pose[:3, :3]) + delta_rotation = R.from_euler("y", -rot_step, degrees=True) + new_rotation = delta_rotation * current_rotation + new_pose[:3, :3] = new_rotation.as_matrix() + pose_changed = True + log_info(f"Roll left: Y rotation -= {rot_step}°") + + # other controls + elif key == ord("r"): + new_pose = init_pose.copy() + pose_changed = True + log_info("Reset to initial pose") + elif key == ord("p"): + translation = new_pose[:3, 3] + rot = R.from_matrix(new_pose[:3, :3]) + quaternion = rot.as_quat() + log_info("Current Camera pose:") + log_info(f"Translation: {translation}") + quat_wxyz = [quaternion[3], quaternion[0], quaternion[1], quaternion[2]] + log_info(f"Quaternion (w, x, y, z): {quat_wxyz}") + + rotation_euler = rot.as_euler("xyz", degrees=True) + log_info(f"Rotation (XYZ Euler, degrees): {rotation_euler}") + + if pose_changed: + cam_pose = new_pose.copy() + cam_pose = torch.as_tensor(cam_pose, dtype=torch.float32).unsqueeze_(0) + sensor.set_local_pose(cam_pose) + + if vis_pose: + sim.update(step=1) + + except KeyboardInterrupt: + if vis_pose: + sim.remove_marker("camera_axis") + log_error("Keyboard control interrupted by user. Exiting control mode...") + finally: + try: + cv2.destroyAllWindows() + except Exception as e: + log_warning(f"cv2.destroyAllWindows() failed: {e}")