add gui & modify camera view

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 = [

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)")

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:

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

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_())