diff --git a/lerobot_aloha/__pycache__/main.cpython-310.pyc b/lerobot_aloha/__pycache__/main.cpython-310.pyc new file mode 100644 index 0000000..02d7524 Binary files /dev/null and b/lerobot_aloha/__pycache__/main.cpython-310.pyc differ diff --git a/lerobot_aloha/common/__pycache__/agilex_robot.cpython-310.pyc b/lerobot_aloha/common/__pycache__/agilex_robot.cpython-310.pyc index 3cb50e6..5f8b554 100644 Binary files a/lerobot_aloha/common/__pycache__/agilex_robot.cpython-310.pyc and b/lerobot_aloha/common/__pycache__/agilex_robot.cpython-310.pyc differ diff --git a/lerobot_aloha/common/__pycache__/robot_components.cpython-310.pyc b/lerobot_aloha/common/__pycache__/robot_components.cpython-310.pyc index 00c1081..f3225e1 100644 Binary files a/lerobot_aloha/common/__pycache__/robot_components.cpython-310.pyc and b/lerobot_aloha/common/__pycache__/robot_components.cpython-310.pyc differ diff --git a/lerobot_aloha/common/__pycache__/rosrobot.cpython-310.pyc b/lerobot_aloha/common/__pycache__/rosrobot.cpython-310.pyc index 9f5fa93..af49d76 100644 Binary files a/lerobot_aloha/common/__pycache__/rosrobot.cpython-310.pyc and b/lerobot_aloha/common/__pycache__/rosrobot.cpython-310.pyc differ diff --git a/lerobot_aloha/common/__pycache__/rosrobot_factory.cpython-310.pyc b/lerobot_aloha/common/__pycache__/rosrobot_factory.cpython-310.pyc index d058a8e..786d65c 100644 Binary files a/lerobot_aloha/common/__pycache__/rosrobot_factory.cpython-310.pyc and b/lerobot_aloha/common/__pycache__/rosrobot_factory.cpython-310.pyc differ diff --git a/lerobot_aloha/common/agilex_robot.py b/lerobot_aloha/common/agilex_robot.py index ce63d76..e7d11c5 100644 --- a/lerobot_aloha/common/agilex_robot.py +++ b/lerobot_aloha/common/agilex_robot.py @@ -125,7 +125,12 @@ class AgilexRobot(Robot): return None, None if any(len(q) == 0 for q in self.sync_arm_queues.values()): + # 遍历字典，检查并报告每个空队列 + for arm_name, queue in self.sync_arm_queues.items(): + if len(queue) == 0: + print(f"{arm_name} arm not connected or queue is empty") return None, None + # 计算最小时间戳 timestamps = [ diff --git a/lerobot_aloha/common/robot_components.py b/lerobot_aloha/common/robot_components.py index 7669f71..1d3957e 100644 --- a/lerobot_aloha/common/robot_components.py +++ b/lerobot_aloha/common/robot_components.py @@ -365,7 +365,7 @@ class RobotDataManager: self.sensors = sensors self.actuators = actuators - def warmup(self, timeout: float = 10.0) -> bool: + def warmup(self, timeout: float = 30.0) -> bool: """ Wait until all data queues have sufficient messages @@ -391,6 +391,7 @@ class RobotDataManager: all_ready = True # Check camera image queues + rospy.loginfo(f"Nums of camera is {len(self.sensors.cameras)}") for cam_name in self.sensors.cameras: if len(self.sensors.sync_img_queues[cam_name]) < 50: rospy.loginfo(f"Waiting for camera {cam_name} (current: {len(self.sensors.sync_img_queues[cam_name])}/50)") diff --git a/lerobot_aloha/common/rosrobot.py b/lerobot_aloha/common/rosrobot.py index 30f80dd..fe3bc26 100644 --- a/lerobot_aloha/common/rosrobot.py +++ b/lerobot_aloha/common/rosrobot.py @@ -112,7 +112,7 @@ class Robot: - def warmup(self, timeout: float = 10.0) -> bool: + def warmup(self, timeout: float = 30.0) -> bool: """Wait until all data queues have at least 20 messages. Args: diff --git a/lerobot_aloha/common/utils/__pycache__/control_utils.cpython-310.pyc b/lerobot_aloha/common/utils/__pycache__/control_utils.cpython-310.pyc index 7fa908c..b2af2e1 100644 Binary files a/lerobot_aloha/common/utils/__pycache__/control_utils.cpython-310.pyc and b/lerobot_aloha/common/utils/__pycache__/control_utils.cpython-310.pyc differ diff --git a/lerobot_aloha/common/utils/control_utils.py b/lerobot_aloha/common/utils/control_utils.py index 11ae4d9..990d60b 100644 --- a/lerobot_aloha/common/utils/control_utils.py +++ b/lerobot_aloha/common/utils/control_utils.py @@ -3,6 +3,7 @@ import time import torch import rospy import cv2 +import numpy as np from contextlib import nullcontext from copy import copy from lerobot.common.datasets.lerobot_dataset import LeRobotDataset @@ -134,42 +135,50 @@ def control_loop( dataset.add_frame(frame) if display_cameras and not is_headless(): - image_keys = [key for key in observation if "image" in key] - - # 获取屏幕分辨率（假设屏幕分辨率为 1920x1080，可以根据实际情况调整） - screen_width = 1920 - screen_height = 1080 - - # 计算窗口的排列方式 + image_keys = [key for key in observation if "image" in key and "depth" not in key] num_images = len(image_keys) - max_columns = int(screen_width / 640) # 假设每个窗口宽度为 640 - rows = (num_images + max_columns - 1) // max_columns # 计算需要的行数 - columns = min(num_images, max_columns) # 实际使用的列数 + + if num_images > 0: + # 设置每个图像的显示尺寸 + display_width = 426 # 更小的宽度 + display_height = 320 # 更小的高度 + + # 确定网格布局的行列数 (尽量接近正方形布局) + grid_cols = int(np.ceil(np.sqrt(num_images))) + grid_rows = int(np.ceil(num_images / grid_cols)) + + # 创建一个大的画布来容纳所有图像 + canvas = np.zeros((grid_rows * display_height, grid_cols * display_width, 3), dtype=np.uint8) + + # 在画布上放置每个图像 + for idx, key in enumerate(image_keys): + row = idx // grid_cols + col = idx % grid_cols + + # 获取图像并转换为BGR + image = observation[key].numpy() + image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + + # 调整图像大小 + resized_image = cv2.resize(image, (display_width, display_height)) + + # 计算在画布上的位置 + y_start = row * display_height + y_end = y_start + display_height + x_start = col * display_width + x_end = x_start + display_width + + # 将图像放置到画布上 + canvas[y_start:y_end, x_start:x_end] = resized_image + + # 添加图像标题 + title_position = (x_start + 5, y_start + 15) + cv2.putText(canvas, key, title_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) + + # 显示合并后的画布 + cv2.imshow("Camera Views", canvas) + cv2.waitKey(1) - # 遍历所有图像键并显示 - for idx, key in enumerate(image_keys): - if "depth" in key: - continue # 跳过深度图像 - - # 将图像从 RGB 转换为 BGR 格式 - image = cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR) - - # 创建窗口 - cv2.imshow(key, image) - - # 计算窗口位置 - window_width = 640 - window_height = 480 - row = idx // max_columns - col = idx % max_columns - x_position = col * window_width - y_position = row * window_height - - # 移动窗口到指定位置 - cv2.moveWindow(key, x_position, y_position) - - # 等待 1 毫秒以处理事件 - cv2.waitKey(1) if fps is not None: dt_s = time.perf_counter() - start_loop_t diff --git a/lerobot_aloha/gui_app.py b/lerobot_aloha/gui_app.py new file mode 100644 index 0000000..a4aa87f --- /dev/null +++ b/lerobot_aloha/gui_app.py @@ -0,0 +1,208 @@ +import sys +import numpy as np +import cv2 +from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, + QHBoxLayout, QLabel, QLineEdit, QSpinBox, QCheckBox, + QPushButton, QGroupBox, QTabWidget, QScrollArea, QGridLayout) +from PyQt5.QtCore import Qt, QTimer +import cv2 +from PyQt5.QtGui import QImage, QPixmap +from main import get_arguments, control_robot + +class ConfigGroup(QGroupBox): + """Group of configuration widgets""" + def __init__(self, title, parent=None): + super().__init__(title, parent) + self.layout = QVBoxLayout() + self.setLayout(self.layout) + + def add_config(self, name, value, widget_type="lineedit"): + """Add a configuration widget""" + row = QHBoxLayout() + label = QLabel(name) + + if isinstance(value, bool): + widget = QCheckBox() + widget.setChecked(value) + elif isinstance(value, int): + widget = QSpinBox() + widget.setRange(0, 999999) + widget.setValue(value) + else: # string or other + widget = QLineEdit(str(value)) + + row.addWidget(label) + row.addWidget(widget) + self.layout.addLayout(row) + return widget + +class MainWindow(QMainWindow): + def __init__(self): + super().__init__() + self.setWindowTitle("MindRobot-V1 Control GUI") + self.setGeometry(100, 100, 400, 600) # Adjusted window size + self.robot = None + + # Get default arguments + self.cfg = get_arguments() + + # Main layout + main_widget = QWidget() + main_layout = QHBoxLayout() + main_widget.setLayout(main_layout) + + # Left panel - configuration + config_scroll = QScrollArea() + config_widget = QWidget() + config_layout = QVBoxLayout() + config_widget.setLayout(config_layout) + + # Add configuration groups + general_group = ConfigGroup("General Settings") + self.fps_widget = general_group.add_config("FPS", self.cfg.fps, "spinbox") + self.resume_widget = general_group.add_config("Resume", self.cfg.resume, "checkbox") + self.repo_id_widget = general_group.add_config("Repo ID", self.cfg.repo_id) + + # Config file with browse button + config_row = QHBoxLayout() + config_label = QLabel("Config File") + self.config_widget = QLineEdit("/home/ubuntu/LYT/lerobot_aloha/lerobot_aloha/configs/agilex.yaml") + config_browse_button = QPushButton("Browse...") + config_browse_button.clicked.connect(self.browse_config_file) + config_row.addWidget(config_label) + config_row.addWidget(self.config_widget) + config_row.addWidget(config_browse_button) + general_group.layout.addLayout(config_row) + # Root directory with browse button + root_row = QHBoxLayout() + root_label = QLabel("Root Directory") + self.root_widget = QLineEdit(str(self.cfg.root)) + browse_button = QPushButton("Browse...") + browse_button.clicked.connect(self.browse_root_directory) + root_row.addWidget(root_label) + root_row.addWidget(self.root_widget) + root_row.addWidget(browse_button) + general_group.layout.addLayout(root_row) + config_layout.addWidget(general_group) + + recording_group = ConfigGroup("Recording Settings") + self.num_episodes_widget = recording_group.add_config("Number of Episodes", self.cfg.num_episodes, "spinbox") + self.episode_time_widget = recording_group.add_config("Episode Time (ms)", self.cfg.episode_time_s, "spinbox") + self.video_widget = recording_group.add_config("Save Video", self.cfg.video, "checkbox") + self.display_cameras_widget = recording_group.add_config("Display Cameras", self.cfg.display_cameras, "checkbox") + config_layout.addWidget(recording_group) + + advanced_group = ConfigGroup("Advanced Settings") + self.num_writer_processes_widget = advanced_group.add_config("Writer Processes", self.cfg.num_image_writer_processes, "spinbox") + self.num_writer_threads_widget = advanced_group.add_config("Threads per Camera", self.cfg.num_image_writer_threads_per_camera, "spinbox") + self.use_depth_widget = advanced_group.add_config("Use Depth Image", self.cfg.use_depth_image, "checkbox") + self.use_base_widget = advanced_group.add_config("Use Base", self.cfg.use_base, "checkbox") + self.push_to_hub_widget = advanced_group.add_config("Push to Hub", self.cfg.push_to_hub, "checkbox") + config_layout.addWidget(advanced_group) + + # Control buttons + control_buttons = QHBoxLayout() + self.record_button = QPushButton("Record") + self.record_button.clicked.connect(self.start_recording) + control_buttons.addWidget(self.record_button) + + self.stop_button = QPushButton("Stop") + self.stop_button.clicked.connect(self.stop_recording) + control_buttons.addWidget(self.stop_button) + config_layout.addLayout(control_buttons) + + config_scroll.setWidget(config_widget) + config_scroll.setWidgetResizable(True) + main_layout.addWidget(config_scroll, stretch=1) # Left panel stretch factor + + # Remove camera view panel completely + + self.setCentralWidget(main_widget) + + # Robot control flag + self.is_recording = False + + + def browse_config_file(self): + """Open file dialog to select config file""" + from PyQt5.QtWidgets import QFileDialog + file_path, _ = QFileDialog.getOpenFileName( + self, + "Select Config File", + self.config_widget.text(), + "YAML Files (*.yaml *.yml)" + ) + if file_path: + self.config_widget.setText(file_path) + + def browse_root_directory(self): + """Open file dialog to select root directory""" + from PyQt5.QtWidgets import QFileDialog + dir_path = QFileDialog.getExistingDirectory( + self, + "Select Root Directory", + self.root_widget.text() + ) + if dir_path: + self.root_widget.setText(dir_path) + + def get_config_values(self): + """Get current configuration values from UI""" + self.cfg.fps = self.fps_widget.value() + self.cfg.resume = self.resume_widget.isChecked() + self.cfg.repo_id = self.repo_id_widget.text() + self.cfg.root = self.root_widget.text() + self.cfg.num_episodes = self.num_episodes_widget.value() + self.cfg.episode_time_s = self.episode_time_widget.value() + self.cfg.video = self.video_widget.isChecked() + self.cfg.display_cameras = self.display_cameras_widget.isChecked() + self.cfg.num_image_writer_processes = self.num_writer_processes_widget.value() + self.cfg.num_image_writer_threads_per_camera = self.num_writer_threads_widget.value() + self.cfg.use_depth_image = self.use_depth_widget.isChecked() + self.cfg.use_base = self.use_base_widget.isChecked() + self.cfg.push_to_hub = self.push_to_hub_widget.isChecked() + self.cfg.control_type = "record" + + def start_recording(self): + """Start recording with current configuration""" + self.get_config_values() + + try: + # Create robot instance with current config + from common.rosrobot_factory import RobotFactory + self.cfg.config_file = self.config_widget.text() + self.robot = RobotFactory.create( + config_file=self.cfg.config_file, + args=self.cfg + ) + from common.utils.data_utils import record + record(self.robot, self.cfg) + + self.is_recording = True + self.record_button.setEnabled(False) + self.stop_button.setEnabled(True) + print("Recording started with configuration:", vars(self.cfg)) + except Exception as e: + print(f"Failed to start recording: {e}") + self.is_recording = False + + def stop_recording(self): + """Stop recording""" + self.is_recording = False + self.record_button.setEnabled(True) + self.stop_button.setEnabled(False) + + # 模拟ESC键按下事件 + from pynput.keyboard import Key, Controller + keyboard = Controller() + keyboard.press(Key.esc) + keyboard.release(Key.esc) + + print("Recording stopped") + + +if __name__ == "__main__": + app = QApplication(sys.argv) + window = MainWindow() + window.show() + sys.exit(app.exec_())