automatic_data_collector

#!/usr/bin/env python3

import numpy as np
import yaml
import argparse
import sys
import rospy

# 3rd-party
from scipy.spatial.transform import Rotation as R
import atom_msgs.srv
from tf.listener import TransformListener

# local packages
from atom_core.config_io import loadConfig
from atom_core.ros_utils import filterLaunchArguments
from atom_core.transformations import compareTransforms


def convertToMatrix(t):
    matrix = np.zeros((4, 4))
    matrix[:3, :3] = R.from_quat(t[1]).as_matrix()
    matrix[:3, 3] = t[0]
    matrix[3, 3] = 1
    return matrix


def compareTransformsWithThresholds(t1, t2, translation_threshold, rotation_threshold):
    """
    Compare two 3D transforms composed of translation and rotation components.

    Parameters:
    - t1 (list): First transform list containing a 3-element translation vector and a 4-element quaternion.
    - t2 (list): Second transform list with the same structure as t1.
    - translation_threshold (float): Maximum allowable Euclidean distance between translation vectors.
    - rotation_threshold (float): Maximum allowable difference between rotation matrices.

    Returns:
    - bool: True if the translation and rotation differences are below the specified thresholds, False otherwise.
    """

    # Converting the transformation to the 4x4 matrix
    t1_matrix = convertToMatrix(t1)
    t2_matrix = convertToMatrix(t2)

    # Extract translation and rotation components
    translation_error, rotation_error, _, _, _, _, _, _  = compareTransforms(t1_matrix, t2_matrix)

    # Check if differences are below the specified thresholds
    return translation_error < translation_threshold and rotation_error < rotation_threshold


class AutomaticDataCollector:

    def __init__(self, config, rate=10, time_threshold=4.0, translation_threshold=0.1, rotation_threshold=0.1):
        """
        Initialize the AutomaticDataCollector.

        Parameters:
        - config (dict): The ATOM config.
        - rate (float): ROS Rate for the data collection loop.
        - time_threshold (float): Maximum time difference for considering frames as static.
        - translation_threshold (float): Maximum allowable Euclidean distance between translation vectors.
        - rotation_threshold (float): Maximum allowable difference between rotation matrices, measured using Frobenius norm.
        """
        self.world_link = config['world_link']
        self.rate = rospy.Rate(rate) #hz
        self.time_threshold = time_threshold
        self.listener = TransformListener()
        self.automatic_data_collector_on = False
        self.frame_times = {}
        self.translation_threshold = translation_threshold
        self.rotation_threshold = rotation_threshold
        self.robot_static = False
        self.collection_taken = False


        # Add service to toggle the automatic data collector
        self.service_get_dataset = rospy.Service('~toggle_automatic_data_collector',
                                                 atom_msgs.srv.ToggleAutomaticDataCollector,
                                                 self.callbackToggleAutomaticDataCollector)

    def callbackToggleAutomaticDataCollector(self, request):
        """
        Callback function for toggling the automatic data collector service.

        Parameters:
        - request: Service request.

        Returns:
        - atom_msgs.srv.ToggleAutomaticDataCollectorResponse: Service response.
        """
        print('callbackAutomaticDataCollector service called')
        self.automatic_data_collector_on = not self.automatic_data_collector_on

        response = atom_msgs.srv.ToggleAutomaticDataCollectorResponse()
        response.success = True
        response.message = f"Automatic Data Collector on: {self.automatic_data_collector_on}"
        return response
    
    def run(self):
        """
        Main loop for automatic data collection and checking if the robot is static.
        """
        frames_dict = {}
        while not rospy.is_shutdown():
            if not self.automatic_data_collector_on:
                continue
            
            # Load full frames dictionary only once
            if frames_dict == {}:
                frames_dict = yaml.safe_load(self.listener._buffer.all_frames_as_yaml())

            for frame in frames_dict.keys():
                t = self.listener.getLatestCommonTime(frame, self.world_link)
                (trans, quat) = self.listener.lookupTransform(frame, self.world_link, t)
                if frame not in self.frame_times.keys():
                    self.frame_times[frame]={'times': [t.to_sec()],
                                            'values': [[trans, quat]],
                                            'static': False}
                elif t.to_sec() not in self.frame_times[frame]['times']:
                    self.frame_times[frame]['times'].append(t.to_sec())
                    self.frame_times[frame]['values'].append([trans, quat])

                    while self.frame_times[frame]['times'][-1] - self.frame_times[frame]['times'][0] > self.time_threshold:
                        self.frame_times[frame]['static'] = compareTransformsWithThresholds(self.frame_times[frame]['values'][0], \
                            self.frame_times[frame]['values'][-1], self.translation_threshold, self.rotation_threshold)
                            
                        
                        self.frame_times[frame]['times'].pop(0)
                        self.frame_times[frame]['values'].pop(0)
                elif t.to_sec() == rospy.Time(0).to_sec():
                    self.frame_times[frame]['static'] = True

                
            # Check if all frames are static
            static_frames = [self.frame_times[frame]['static'] for frame in self.frame_times if 'static' in self.frame_times[frame]]
            self.robot_static = all(static_frames)

            # If the robot is static and collection hasn't been taken yet, attempt to save the collection
            if self.robot_static and not self.collection_taken:
                rospy.wait_for_service('/collect_data/save_collection')
                save_collection = rospy.ServiceProxy('/collect_data/save_collection', atom_msgs.srv.SaveCollection)
                response = save_collection()
                self.collection_taken = response.success
            elif not self.robot_static:
                self.collection_taken = False

            self.rate.sleep()


    

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-c", "--calibration_file", help='full path to calibration file.', type=str, required=True)
    args = vars(ap.parse_args(args=filterLaunchArguments(sys.argv)))
    config = loadConfig(args['calibration_file'])

    # Initialize ROS 
    rospy.init_node("automatic_data_collector")

    automatic_data_collector = AutomaticDataCollector(config)
    automatic_data_collector.run()

if __name__ == '__main__':
    main()