404 lines
17 KiB
Python
404 lines
17 KiB
Python
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_())
|