Improve the gui_app for data collection

This commit is contained in:
2025-04-08 22:52:27 +08:00
parent 88885a6a25
commit 3df284ddd1
5 changed files with 56 additions and 22 deletions

View File

@@ -1,3 +1,5 @@
export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtiff.so.5
python collect_data.py --robot.type=aloha --control.type=record --control.fps=30 --control.single_task="Grasp a lego block and put it in the bin." --control.repo_id=tangger/test --control.num_episodes=1 --control.root=./data
python lerobot/scripts/train.py --dataset.repo_id=maic/move_tube_on_scale --policy.type=act --output_dir=outputs/train/act_move_tube_on_scale --job_name=act_move_tube_on_scale --policy.device=cuda --wandb.enable=true --dataset.root=/home/ubuntu/LYT/aloha_lerobot/data1

View File

@@ -1,6 +1,12 @@
robot_type: aloha_agilex
ros_node_name: record_episodes
cameras:
cam_high:
img_topic_name: /camera/color/image_raw
depth_topic_name: /camera/depth/image_rect_raw
rgb_shape: [480, 640, 3]
width: 480
height: 640
cam_front:
img_topic_name: /camera_f/color/image_raw
depth_topic_name: /camera_f/depth/image_raw
@@ -19,12 +25,7 @@ cameras:
rgb_shape: [480, 640, 3]
width: 480
height: 640
cam_high:
img_topic_name: /camera/color/image_raw
depth_topic_name: /camera/depth/image_rect_raw
rgb_shape: [480, 640, 3]
width: 480
height: 640
arm:
master_left:

View File

@@ -40,11 +40,11 @@ class MainWindow(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("MindRobot-V1 Control GUI")
self.setGeometry(100, 100, 400, 600) # Adjusted window size
self.setGeometry(100, 100, 600, 800) # Adjusted window size
self.robot = None
# Get default arguments
self.cfg = get_arguments()
self.cfg = None
# Main layout
main_widget = QWidget()
@@ -59,9 +59,16 @@ class MainWindow(QMainWindow):
# 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)
self.single_task_widget = general_group.add_config("Single Task", "")
self.single_task_widget.setMinimumWidth(300) # Wider text box
self.single_task_widget.textChanged.connect(self.update_repo_id_from_task)
self.fps_widget = general_group.add_config("FPS", 30, "spinbox")
self.resume_widget = general_group.add_config("Resume", False, "checkbox")
self.repo_id_widget = general_group.add_config("Repo ID", "")
self.repo_id_widget.textChanged.connect(self.update_root_from_repo_id)
self.original_repo_id = "" # Store original value
self.play_sounds_widget = general_group.add_config("Play Sounds", False, "checkbox")
# Config file with browse button
config_row = QHBoxLayout()
@@ -76,7 +83,7 @@ class MainWindow(QMainWindow):
# Root directory with browse button
root_row = QHBoxLayout()
root_label = QLabel("Root Directory")
self.root_widget = QLineEdit(str(self.cfg.root))
self.root_widget = QLineEdit(str("/home/ubuntu/LYT/lerobot_aloha/datasets/"+self.repo_id_widget.text()))
browse_button = QPushButton("Browse...")
browse_button.clicked.connect(self.browse_root_directory)
root_row.addWidget(root_label)
@@ -86,18 +93,20 @@ class MainWindow(QMainWindow):
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")
self.episode_widget = recording_group.add_config("Episode", 0, "spinbox")
self.num_episodes_widget = recording_group.add_config("Number of Episodes", 100, "spinbox")
self.episode_time_widget = recording_group.add_config("Episode Time (ms)", 36000, "spinbox")
self.video_widget = recording_group.add_config("Save Video", True, "checkbox")
self.display_cameras_widget = recording_group.add_config("Display Cameras", True, "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")
self.num_writer_processes_widget = advanced_group.add_config("Writer Processes", 0, "spinbox")
self.num_writer_threads_widget = advanced_group.add_config("Threads per Camera", 4, "spinbox")
self.use_depth_widget = advanced_group.add_config("Use Depth Image", False, "checkbox")
self.use_base_widget = advanced_group.add_config("Use Base", False, "checkbox")
self.push_to_hub_widget = advanced_group.add_config("Push to Hub", False, "checkbox")
self.policy_widget = advanced_group.add_config("Policy", None)
config_layout.addWidget(advanced_group)
# Control buttons
@@ -135,6 +144,16 @@ class MainWindow(QMainWindow):
if file_path:
self.config_widget.setText(file_path)
def update_repo_id_from_task(self, text):
"""Update repo_id from single_task text, replacing spaces with underscores"""
if not hasattr(self, 'repo_id_edited') or not self.repo_id_edited:
self.repo_id_widget.setText(text.replace(" ", "_"))
def update_root_from_repo_id(self, text):
"""Update root directory based on repo_id"""
if self.root_widget:
self.root_widget.setText(f"/home/ubuntu/LYT/lerobot_aloha/datasets/{text}")
def browse_root_directory(self):
"""Open file dialog to select root directory"""
from PyQt5.QtWidgets import QFileDialog
@@ -151,20 +170,32 @@ class MainWindow(QMainWindow):
self.cfg.fps = self.fps_widget.value()
self.cfg.resume = self.resume_widget.isChecked()
self.cfg.repo_id = self.repo_id_widget.text()
self.repo_id_edited = (self.cfg.repo_id != self.original_repo_id)
self.cfg.root = self.root_widget.text()
self.cfg.episode = self.episode_widget.value()
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.play_sounds = self.play_sounds_widget.isChecked()
self.cfg.single_task = self.single_task_widget.text()
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.policy = None
self.cfg.control_type = "record"
def start_recording(self):
"""Start recording with current configuration"""
if not hasattr(self, 'cfg') or self.cfg is None:
from main import get_arguments
self.cfg = get_arguments()
if self.fps_widget is None:
return
self.get_config_values()
try:

View File

@@ -26,7 +26,7 @@ def get_arguments():
args.play_sounds = False
args.display_cameras = True
args.single_task = "move the bottle from the right to the scale right"
args.use_depth_image = False
args.use_depth_image = True
args.use_base = False
args.push_to_hub = False
args.policy = None