Improve the gui_app for data collection

lerobot_aloha/README.MD

@@ -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

lerobot_aloha/configs/agilex.yaml

@@ -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:

lerobot_aloha/gui_app.py

@@ -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.single_task_widget = general_group.add_config("Single Task", "")
-        self.resume_widget = general_group.add_config("Resume", self.cfg.resume, "checkbox")
+        self.single_task_widget.setMinimumWidth(300)  # Wider text box
-        self.repo_id_widget = general_group.add_config("Repo ID", self.cfg.repo_id)
+        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_widget = recording_group.add_config("Episode", 0, "spinbox")
-        self.episode_time_widget = recording_group.add_config("Episode Time (ms)", self.cfg.episode_time_s, "spinbox")
+        self.num_episodes_widget = recording_group.add_config("Number of Episodes", 100, "spinbox")
-        self.video_widget = recording_group.add_config("Save Video", self.cfg.video, "checkbox")
+        self.episode_time_widget = recording_group.add_config("Episode Time (ms)", 36000, "spinbox")
-        self.display_cameras_widget = recording_group.add_config("Display Cameras", self.cfg.display_cameras, "checkbox")
+        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_processes_widget = advanced_group.add_config("Writer Processes", 0, "spinbox")
-        self.num_writer_threads_widget = advanced_group.add_config("Threads per Camera", self.cfg.num_image_writer_threads_per_camera, "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", self.cfg.use_depth_image, "checkbox")
+        self.use_depth_widget = advanced_group.add_config("Use Depth Image", False, "checkbox")
-        self.use_base_widget = advanced_group.add_config("Use Base", self.cfg.use_base, "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", self.cfg.push_to_hub, "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:

lerobot_aloha/main.py

@@ -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