import sys import os import yaml from collect_data import main from types import SimpleNamespace from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QSpinBox, QCheckBox, QGroupBox, QTabWidget, QMessageBox, QFileDialog) from PyQt5.QtCore import Qt class AlohaDataCollectionGUI(QMainWindow): def __init__(self): super().__init__() self.setWindowTitle("MindRobot-V1 Data Collection") self.setGeometry(100, 100, 800, 600) self.central_widget = QWidget() self.setCentralWidget(self.central_widget) self.main_layout = QVBoxLayout(self.central_widget) self.config_path = os.path.expanduser("/home/ubuntu/LYT/lerobot_aloha/collect_data/aloha.yaml") self.create_ui() self.setup_connections() self.load_default_config() def create_ui(self): # Create tabs self.tabs = QTabWidget() self.main_layout.addWidget(self.tabs) # General Settings Tab self.general_tab = QWidget() self.tabs.addTab(self.general_tab, "General Settings") self.create_general_tab() # Camera Settings Tab self.camera_tab = QWidget() self.tabs.addTab(self.camera_tab, "Camera Settings") self.create_camera_tab() # Arm Settings Tab self.arm_tab = QWidget() self.tabs.addTab(self.arm_tab, "Arm Settings") self.create_arm_tab() # Control Buttons self.control_group = QGroupBox("Control") control_layout = QHBoxLayout() self.load_config_button = QPushButton("Load Config") self.save_config_button = QPushButton("Save Config") self.start_button = QPushButton("Start Recording") self.stop_button = QPushButton("Stop Recording") self.exit_button = QPushButton("Exit") control_layout.addWidget(self.load_config_button) control_layout.addWidget(self.save_config_button) control_layout.addWidget(self.start_button) control_layout.addWidget(self.stop_button) control_layout.addWidget(self.exit_button) self.control_group.setLayout(control_layout) self.main_layout.addWidget(self.control_group) def create_general_tab(self): layout = QVBoxLayout(self.general_tab) # Config File Path config_group = QGroupBox("Configuration File") config_layout = QHBoxLayout() self.config_path_edit = QLineEdit(self.config_path) self.browse_config_button = QPushButton("Browse...") config_layout.addWidget(QLabel("Config File:")) config_layout.addWidget(self.config_path_edit) config_layout.addWidget(self.browse_config_button) config_group.setLayout(config_layout) layout.addWidget(config_group) # Dataset Directory dir_group = QGroupBox("Dataset Directory") dir_layout = QHBoxLayout() self.dataset_dir_edit = QLineEdit() self.browse_dir_button = QPushButton("Browse...") dir_layout.addWidget(QLabel("Dataset Directory:")) dir_layout.addWidget(self.dataset_dir_edit) dir_layout.addWidget(self.browse_dir_button) dir_group.setLayout(dir_layout) layout.addWidget(dir_group) # Task Settings task_group = QGroupBox("Task Settings") task_layout = QVBoxLayout() self.task_name_edit = QLineEdit() self.episode_idx_spin = QSpinBox() self.episode_idx_spin.setRange(0, 9999) self.max_timesteps_spin = QSpinBox() self.max_timesteps_spin.setRange(1, 10000) self.num_episodes_spin = QSpinBox() self.num_episodes_spin.setRange(1, 1000) self.frame_rate_spin = QSpinBox() self.frame_rate_spin.setRange(1, 60) task_layout.addWidget(QLabel("Task Name:")) task_layout.addWidget(self.task_name_edit) task_layout.addWidget(QLabel("Episode Index:")) task_layout.addWidget(self.episode_idx_spin) task_layout.addWidget(QLabel("Max Timesteps:")) task_layout.addWidget(self.max_timesteps_spin) task_layout.addWidget(QLabel("Number of Episodes:")) task_layout.addWidget(self.num_episodes_spin) task_layout.addWidget(QLabel("Frame Rate:")) task_layout.addWidget(self.frame_rate_spin) task_group.setLayout(task_layout) layout.addWidget(task_group) # Options options_group = QGroupBox("Options") options_layout = QVBoxLayout() self.use_robot_base_check = QCheckBox("Use Robot Base") self.use_depth_image_check = QCheckBox("Use Depth Image") options_layout.addWidget(self.use_robot_base_check) options_layout.addWidget(self.use_depth_image_check) options_group.setLayout(options_layout) layout.addWidget(options_group) layout.addStretch() def create_camera_tab(self): layout = QVBoxLayout(self.camera_tab) # Color Image Topics color_group = QGroupBox("Color Image Topics") color_layout = QVBoxLayout() self.img_front_topic_edit = QLineEdit() self.img_left_topic_edit = QLineEdit() self.img_right_topic_edit = QLineEdit() color_layout.addWidget(QLabel("Front Camera Topic:")) color_layout.addWidget(self.img_front_topic_edit) color_layout.addWidget(QLabel("Left Camera Topic:")) color_layout.addWidget(self.img_left_topic_edit) color_layout.addWidget(QLabel("Right Camera Topic:")) color_layout.addWidget(self.img_right_topic_edit) color_group.setLayout(color_layout) layout.addWidget(color_group) # Depth Image Topics depth_group = QGroupBox("Depth Image Topics") depth_layout = QVBoxLayout() self.img_front_depth_topic_edit = QLineEdit() self.img_left_depth_topic_edit = QLineEdit() self.img_right_depth_topic_edit = QLineEdit() depth_layout.addWidget(QLabel("Front Depth Topic:")) depth_layout.addWidget(self.img_front_depth_topic_edit) depth_layout.addWidget(QLabel("Left Depth Topic:")) depth_layout.addWidget(self.img_left_depth_topic_edit) depth_layout.addWidget(QLabel("Right Depth Topic:")) depth_layout.addWidget(self.img_right_depth_topic_edit) depth_group.setLayout(depth_layout) layout.addWidget(depth_group) layout.addStretch() def create_arm_tab(self): layout = QVBoxLayout(self.arm_tab) # Master Arm Topics master_group = QGroupBox("Master Arm Topics") master_layout = QVBoxLayout() self.master_arm_left_topic_edit = QLineEdit() self.master_arm_right_topic_edit = QLineEdit() master_layout.addWidget(QLabel("Master Left Arm Topic:")) master_layout.addWidget(self.master_arm_left_topic_edit) master_layout.addWidget(QLabel("Master Right Arm Topic:")) master_layout.addWidget(self.master_arm_right_topic_edit) master_group.setLayout(master_layout) layout.addWidget(master_group) # Puppet Arm Topics puppet_group = QGroupBox("Puppet Arm Topics") puppet_layout = QVBoxLayout() self.puppet_arm_left_topic_edit = QLineEdit() self.puppet_arm_right_topic_edit = QLineEdit() puppet_layout.addWidget(QLabel("Puppet Left Arm Topic:")) puppet_layout.addWidget(self.puppet_arm_left_topic_edit) puppet_layout.addWidget(QLabel("Puppet Right Arm Topic:")) puppet_layout.addWidget(self.puppet_arm_right_topic_edit) puppet_group.setLayout(puppet_layout) layout.addWidget(puppet_group) # Robot Base Topic base_group = QGroupBox("Robot Base Topic") base_layout = QVBoxLayout() self.robot_base_topic_edit = QLineEdit() base_layout.addWidget(QLabel("Robot Base Topic:")) base_layout.addWidget(self.robot_base_topic_edit) base_group.setLayout(base_layout) layout.addWidget(base_group) layout.addStretch() def setup_connections(self): self.load_config_button.clicked.connect(self.load_config) self.save_config_button.clicked.connect(self.save_config) self.browse_config_button.clicked.connect(self.browse_config_file) self.browse_dir_button.clicked.connect(self.browse_dataset_dir) self.start_button.clicked.connect(self.start_recording) self.stop_button.clicked.connect(self.stop_recording) self.exit_button.clicked.connect(self.close) def load_default_config(self): default_config = { 'dataset_dir': '/home/ubuntu/LYT/lerobot_aloha/datasets/3camera', 'task_name': 'aloha_mobile_dummy', 'episode_idx': 0, 'max_timesteps': 500, 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'], 'num_episodes': 50, 'img_front_topic': '/camera_f/color/image_raw', 'img_left_topic': '/camera_l/color/image_raw', 'img_right_topic': '/camera_r/color/image_raw', 'img_front_depth_topic': '/camera_f/depth/image_raw', 'img_left_depth_topic': '/camera_l/depth/image_raw', 'img_right_depth_topic': '/camera_r/depth/image_raw', 'master_arm_left_topic': '/master/joint_left', 'master_arm_right_topic': '/master/joint_right', 'puppet_arm_left_topic': '/puppet/joint_left', 'puppet_arm_right_topic': '/puppet/joint_right', 'robot_base_topic': '/odom', 'use_robot_base': False, 'use_depth_image': False, 'frame_rate': 30 } # Update UI with default values self.update_ui_from_config(default_config) def update_ui_from_config(self, config): """Update UI elements from a config dictionary""" self.dataset_dir_edit.setText(config.get('dataset_dir', '')) self.task_name_edit.setText(config.get('task_name', '')) self.episode_idx_spin.setValue(config.get('episode_idx', 0)) self.max_timesteps_spin.setValue(config.get('max_timesteps', 500)) self.num_episodes_spin.setValue(config.get('num_episodes', 1)) self.frame_rate_spin.setValue(config.get('frame_rate', 30)) self.img_front_topic_edit.setText(config.get('img_front_topic', '')) self.img_left_topic_edit.setText(config.get('img_left_topic', '')) self.img_right_topic_edit.setText(config.get('img_right_topic', '')) self.img_front_depth_topic_edit.setText(config.get('img_front_depth_topic', '')) self.img_left_depth_topic_edit.setText(config.get('img_left_depth_topic', '')) self.img_right_depth_topic_edit.setText(config.get('img_right_depth_topic', '')) self.master_arm_left_topic_edit.setText(config.get('master_arm_left_topic', '')) self.master_arm_right_topic_edit.setText(config.get('master_arm_right_topic', '')) self.puppet_arm_left_topic_edit.setText(config.get('puppet_arm_left_topic', '')) self.puppet_arm_right_topic_edit.setText(config.get('puppet_arm_right_topic', '')) self.robot_base_topic_edit.setText(config.get('robot_base_topic', '')) self.use_robot_base_check.setChecked(config.get('use_robot_base', False)) self.use_depth_image_check.setChecked(config.get('use_depth_image', False)) def get_config_from_ui(self): """Get current UI values as a config dictionary""" config = { 'dataset_dir': self.dataset_dir_edit.text(), 'task_name': self.task_name_edit.text(), 'episode_idx': self.episode_idx_spin.value(), 'max_timesteps': self.max_timesteps_spin.value(), 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'], 'num_episodes': self.num_episodes_spin.value(), 'img_front_topic': self.img_front_topic_edit.text(), 'img_left_topic': self.img_left_topic_edit.text(), 'img_right_topic': self.img_right_topic_edit.text(), 'img_front_depth_topic': self.img_front_depth_topic_edit.text(), 'img_left_depth_topic': self.img_left_depth_topic_edit.text(), 'img_right_depth_topic': self.img_right_depth_topic_edit.text(), 'master_arm_left_topic': self.master_arm_left_topic_edit.text(), 'master_arm_right_topic': self.master_arm_right_topic_edit.text(), 'puppet_arm_left_topic': self.puppet_arm_left_topic_edit.text(), 'puppet_arm_right_topic': self.puppet_arm_right_topic_edit.text(), 'robot_base_topic': self.robot_base_topic_edit.text(), 'use_robot_base': self.use_robot_base_check.isChecked(), 'use_depth_image': self.use_depth_image_check.isChecked(), 'frame_rate': self.frame_rate_spin.value() } return config def browse_config_file(self): file_path, _ = QFileDialog.getOpenFileName( self, "Select Config File", "", "YAML Files (*.yaml *.yml)" ) if file_path: self.config_path_edit.setText(file_path) self.load_config() def browse_dataset_dir(self): dir_path = QFileDialog.getExistingDirectory( self, "Select Dataset Directory" ) if dir_path: self.dataset_dir_edit.setText(dir_path) def load_config(self): config_path = self.config_path_edit.text() if not os.path.exists(config_path): QMessageBox.warning(self, "Warning", f"Config file not found: {config_path}") return try: with open(config_path, 'r') as f: config = yaml.safe_load(f) self.update_ui_from_config(config) self.statusBar().showMessage(f"Config loaded from {config_path}", 3000) except Exception as e: QMessageBox.critical(self, "Error", f"Failed to load config: {str(e)}") def save_config(self): config_path = self.config_path_edit.text() if not config_path: QMessageBox.warning(self, "Warning", "Please specify a config file path") return try: config = self.get_config_from_ui() with open(config_path, 'w') as f: yaml.dump(config, f, default_flow_style=False) self.statusBar().showMessage(f"Config saved to {config_path}", 3000) except Exception as e: QMessageBox.critical(self, "Error", f"Failed to save config: {str(e)}") def start_recording(self): try: # Save current config to a temporary file temp_config_path = "/tmp/aloha_temp_config.yaml" config = self.get_config_from_ui() # Validate inputs if not config['dataset_dir']: QMessageBox.warning(self, "Warning", "Dataset directory cannot be empty!") return if not config['task_name']: QMessageBox.warning(self, "Warning", "Task name cannot be empty!") return with open(temp_config_path, 'w') as f: yaml.dump(config, f, default_flow_style=False) self.statusBar().showMessage("Recording started...") # Start recording with the temporary config file exit_code = main(temp_config_path) if exit_code == 0: self.statusBar().showMessage("Recording completed successfully!", 5000) else: self.statusBar().showMessage("Recording completed with errors", 5000) except Exception as e: QMessageBox.critical(self, "Error", f"An error occurred: {str(e)}") self.statusBar().showMessage("Recording failed", 5000) def stop_recording(self): # In a real application, this would signal the recording thread to stop self.statusBar().showMessage("Recording stopped", 5000) QMessageBox.information(self, "Info", "Stop recording requested. This would stop the recording in a real implementation.") if __name__ == "__main__": app = QApplication(sys.argv) window = AlohaDataCollectionGUI() window.show() sys.exit(app.exec_())