feat(collect): added continuous sensor data collection and imu modality

#1000

Implemented the ability to collect data in a continuous manner, i.e.,
collect all sensor messages and record them all in the dataset during
data collection. To do this, a new 'continuous' field in the configuration for a
given sensor was created and included in the template.

To maintain backwards compatibility with older datasets, if there is no
'continuous' field, it defaults to False and works as it did previously.

The 'continuous' field was also added to the configuration YAML
template.

Additionally, support for the IMU sensor modality in data collection was also added.

atom_calibration/src/atom_calibration/collect/data_collector.py

@@ -255,16 +255,34 @@ def __init__(self, args, server, menu_handler):
                                                        self.callbackDeleteCollection)
 
         # Create the label_msgs dictionary and fill it with the first messages
+        # Also create the msg_buffer dict for continuous data collection
+        # NOTE: IMUs don't make use of labelers, so the label_msg_type and the labels_topic are simply the regular IMU msg type and message topic
         self.label_msgs = {}
+        self.msg_buffer = {}
         for sensor_key, sensor in self.config['sensors'].items():
             if sensor['modality'] == 'rgb':
                 label_msg_type = ImageWithRGBLabels
             elif sensor['modality'] == 'lidar3d':
                 label_msg_type = PointCloudWithLidar3DLabels
             elif sensor['modality'] == 'depth':
                 label_msg_type = DepthImageWithDepthLabels
+            elif sensor['modality'] == "imu":
+                label_msg_type = sensor_msgs.msg.Imu
+
+            if sensor['modality'] == "imu":
+                labels_topic = sensor["topic_name"]
+            else:
+                labels_topic = sensor_key + '/labels'
+
+            # Check if the sensor has the 'continuous' key. If it does not, default it to False
+            if 'continuous' not in sensor.keys():
+                sensor['continuous'] = False
+
+            # Create an empty list of data for each sensor with continuous collection for the msg_buffer
+            if sensor['continuous'] == True:
+                self.msg_buffer[sensor_key] = []
+                print('Created message buffer for continuous data collection for sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL + "!")
 
-            labels_topic = sensor_key + '/labels'
 
             print("Waiting for first message on topic " +
                   labels_topic + ' ... ', end='')
@@ -302,8 +320,13 @@ def __init__(self, args, server, menu_handler):
                 label_msg_type = PointCloudWithLidar3DLabels
             elif sensor['modality'] == 'depth':
                 label_msg_type = DepthImageWithDepthLabels
+            elif sensor['modality'] == 'imu':
+                label_msg_type = sensor_msgs.msg.Imu
 
-            labels_topic = sensor_key + '/labels'
+            if sensor["modality"] == "imu":
+                labels_topic = sensor["topic_name"]
+            else:
+                labels_topic = sensor_key + '/labels'
 
             print('Setting up subscriber for ' + Fore.BLUE + sensor_key + Style.RESET_ALL +
                   ' label msgs on topic ' + Fore.GREEN + labels_topic + Style.RESET_ALL)
@@ -325,8 +348,12 @@ def callbackReceivedAdditionalDataMsg(self, msg, additional_data_key):
 
     def callbackReceivedLabelMsg(self, msg, sensor_key):
         # print('Received labels message for sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL)
+        if self.config['sensors'][sensor_key]['continuous'] == True:
+            self.msg_buffer[sensor_key].append(msg)
+
         self.label_msgs[sensor_key] = msg
 
+
     def callbackReceivedJointStateMsg(self, msg):
         # Add the joint positions to the dictionary
         for name, position in zip(msg.name, msg.position):
@@ -535,6 +562,9 @@ def saveCollection(self):
             all_sensor_labels_dict[pattern_key] = {}
             for sensor_key, sensor in self.sensors.items():
 
+                if sensor['modality'] == 'imu':
+                    continue
+
                 print('Collecting data from sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL +
                       ' for pattern ' + Fore.GREEN + pattern_key + Style.RESET_ALL)
 
@@ -605,6 +635,17 @@ def saveCollection(self):
             all_additional_data_dict[additional_data_key] = message_converter.convert_ros_message_to_dictionary(
                 msg)
 
+        # --------------------------------------
+        # Create continuous_sensor_data_dict
+        # --------------------------------------
+        continuous_sensor_data_dict = {}
+        for sensor_key, msg_list in self.msg_buffer.items():
+            continuous_sensor_data_dict[sensor_key] = []
+            for msg in msg_list:
+                msg_dict = message_converter.convert_ros_message_to_dictionary(msg)
+                continuous_sensor_data_dict[sensor_key].append(msg_dict)
+
+
         # --------------------------------------
         # Create collection_dict
         # --------------------------------------
@@ -621,6 +662,7 @@ def saveCollection(self):
                    'calibration_config': self.config,
                    'collections': self.collections,
                    'additional_sensor_data': self.additional_data,
+                   'continuous_sensor_data': continuous_sensor_data_dict,
                    'sensors': self.sensors,
                    'transforms': self.transforms,
                    'patterns': self.patterns_dict}

atom_calibration/src/atom_calibration/collect/patterns.py

@@ -397,6 +397,10 @@ def estimatePatternPosesForCollection(dataset, collection_key):
 
         for sensor_key, sensor in dataset['sensors'].items():
 
+            # Ignore the sensor if it's an IMU, as no pattern pose estimation can be made
+            if sensor['modality'] == 'imu':
+                continue
+
             # if pattern not detected by sensor in collection
             if not collection['labels'][pattern_key][sensor_key]['detected']:
                 continue

atom_calibration/templates/config.yml

@@ -63,7 +63,9 @@ sensors:
   #           Current options: lidar3d, rgb, depth, lidar2d
   #       throttle:
   #           Set throttle: desired_frame_rate. If you don't use throttle, do not enter a value, i.e. "throttle: "
-  #
+  #       continuous:
+  #           Set to True to save all data from the sensor in the dataset during data collection
+
   # If you are using an image compressed topic such as:
   #   /world_camera/rgb/image_raw/compressed
   # you should not add the "compressed" part, use only:
@@ -77,6 +79,7 @@ sensors:
     topic_name: "/hand_camera/image_raw"
     modality: "rgb"
     throttle: 10
+    continuous: False
 
 
   # You can have more than one sensor:
@@ -87,6 +90,7 @@ sensors:
     topic_name: "/another_camera/image_raw"
     modality: "rgb"
     throttle:
+    continuous: False
 
 # ATOM will calibrate the extrinsic parameters of your links.
 # In this section you should discriminate the additional transformations that will be part of the calibrations.

atom_core/src/atom_core/config_io.py

@@ -13,7 +13,7 @@
 from rospkg.common import ResourceNotFound
 from colorama import Fore, Style
 
-from atom_core.utilities import atomError
+from atom_core.utilities import atomError, atomWarn
 from atom_core.system import resolvePath, expandToLaunchEnv
 
 
@@ -49,15 +49,21 @@ def verifyConfig(config, template_config):
     for sensor_key, sensor in config['sensors'].items():
         same_keys, extra_keys, missing_keys = dictionaries_have_same_keys(sensor, sensor_template)
         if not same_keys:
-            atomError('In config file, sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL + ' does not have the correct keys.\nKeys that should not exist: ' +
+            if missing_keys == ['continuous']:
+                atomWarn('In config file, sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL + ' does not have the "continuous" key. Disbaling continuous data collection.')
+            else:
+                atomError('In config file, sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL + ' does not have the correct keys.\nKeys that should not exist: ' +
                       str(extra_keys) + '\nKeys that are missing : ' + str(missing_keys))
 
     if config['additional_tfs'] is not None:
         additional_tf_template = {'parent_link': None, 'child_link': None}
         for additional_tf_key, additional_tf in config['additional_tfs'].items():
             same_keys, extra_keys, missing_keys = dictionaries_have_same_keys(additional_tf, additional_tf_template)
             if not same_keys:
-                atomError('In config file, additional_tf ' + Fore.BLUE + additional_tf_key + Style.RESET_ALL + ' does not have the correct keys.\nKeys that should not exist: ' +
+                if missing_keys == ['continuous']:
+                    atomWarn('In config file, sensor ' + Fore.BLUE + sensor_key + Style.RESET_ALL + ' does not have the "continuous" key. Disabling continuous data collection.')
+                else:
+                    atomError('In config file, additional_tf ' + Fore.BLUE + additional_tf_key + Style.RESET_ALL + ' does not have the correct keys.\nKeys that should not exist: ' +
                           str(extra_keys) + '\nKeys that are missing : ' + str(missing_keys))
 
     if config['joints'] is not None:

atom_examples/rihibot/rihibot_calibration/calibration/config.yml

@@ -86,6 +86,7 @@ sensors:
     child_link: "imu_link"
     topic_name: "/imu"
     modality: "imu"
+    continuous: True
     throttle:
 
 # ATOM will calibrate the extrinsic parameters of your links.
@@ -105,6 +106,16 @@ additional_tfs:
   #  parent_link: "base_footprint"
   #  child_link: "base_link"
 
+
+# I don't actually think this will work or is even needed in this case -- i can just make the changes in the sensors themselves. However, I will keep this here as a reminder to document this and to add it to the template
+
+# additional_data:
+  # imu:
+    # modality: "imu"
+    # link: "imu_link"
+    # parent_link: "flange"
+    # child_link: "imu_link"
+
 # ATOM can also calibrate several parameters of your joints.
 # In this section you should discriminate the joints that will be part of the calibrations.
 # For each joint you must specify which parameters to optimize.

atom_examples/rihibot/rihibot_calibration/urdf/initial_estimate.urdf.xacro

@@ -95,7 +95,6 @@
     </xacro:rgb_camera>
 
     <!-- IMU -->
-
     <xacro:imu parent="flange" prefix="imu" color_gazebo="Red">
         <origin xyz="0.01 0 0" rpy="0 0 0" />
         <color_rgb_rviz>