Tools that facilitate processing rosbag data as dictionaries or storing them in RLBench friendly format.

import yaml
import logging
import pathlib as pl
import sys
from ros_tf2_wrapper import get_populated_tf2_wrapper
from rosbag_to_rlbench import (_setup_rosbag_dataset, 
                               extract_list_from_rosbag_dataset, 
                               add_tfs_to_episodes, 
                               keypoints_from_frame,
                               save_keypoint)
from rlbench.backend.observation import Observation
import numpy as np

class Config:
    def __init__(self, cfg: dict = {}):
        self.__dict__.update(cfg)

logging.basicConfig(level=logging.INFO, format="%(message)s",
                    handlers=[logging.StreamHandler(sys.stdout)])

CONFIG_DIR = './configs/peract.yaml'



# TODO: What about ignore collision and low_dim_state and time?????

def get_rlbench_obs(eps, misc):
     # pose format -> x, y, z, qx, qy, qz, qw
    dummy_pose = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0 ,1.0])

    gPO_ratio = eps['gripper_input']['gPO'] / 255.0
    obs = Observation(
        left_shoulder_rgb=None,          left_shoulder_depth=None,
        left_shoulder_point_cloud=None,  right_shoulder_rgb=None,
        right_shoulder_depth=None,       right_shoulder_point_cloud=None,
        overhead_rgb=None,               overhead_depth=None,
        overhead_point_cloud=None,       wrist_rgb=None,
        wrist_depth=None,                wrist_point_cloud=None,
        front_rgb=None,                  front_depth=None,
        front_point_cloud=None,          left_shoulder_mask=None,
        right_shoulder_mask=None,        overhead_mask=None,
        wrist_mask=None,                 front_mask=None,
        joint_velocities            = eps['joint_states']['velocities'],
        joint_positions             = eps['joint_states']['positions'],
        joint_forces                = eps['joint_states']['efforts'],
        gripper_open                = gPO_ratio > 0.9,
        gripper_pose                = dummy_pose,
        gripper_matrix              = eps['end_effector'],
        gripper_touch_forces        = None,
        gripper_joint_positions     = [gPO_ratio/2, gPO_ratio/2],
        task_low_dim_state          = None,
        ignore_collisions           = True, # TODO: fix
        misc=misc,)
    return obs


# Load yaml config
assert pl.Path(CONFIG_DIR).exists(), f'Config file not found: {CONFIG_DIR}'
cfg = Config(yaml.safe_load(open(CONFIG_DIR, 'r')))

data_dir = pl.Path(cfg.data_dir)
out_dir = pl.Path(cfg.save_path)

# For every task folder and its description
for task_folder, description in zip(cfg.task_folders, cfg.task_descriptions):
    task_dir = data_dir/task_folder
    logging.info('-'*30)
    logging.info('.'*5 + f'Processing task: {task_folder}' + '.'*5)
    # For every rosbag/demo in folder
    for demo_idx, demo_path in enumerate(task_dir.glob('*.bag')):
        logging.info('\n'+'.'*3 + f'processing: {demo_path}' + '.'*3)
        logging.info(f'depth scale: {cfg.depth_scale}')
        
        tf_data = get_populated_tf2_wrapper(demo_path, slack_sec=cfg.topic_time_slack)
        tf_data.register_getter_hook(lambda x: cfg.reference_frame_T @ x, cfg.reference_frame)
        dataset = _setup_rosbag_dataset(demo_path, cfg.topics_and_names,
                                        time_slack=cfg.topic_time_slack,
                                        reference_topic=cfg.reference_topic)
        dataset.register_topic_postprocess('/right/Robotiq2FGripperRobotInput', 
                                           lambda a: {'gACT': a.gACT, 'gGTO': a.gGTO, 'gSTA': a.gSTA,
                                                      'gOBJ': a.gOBJ, 'gFLT': a.gFLT,'gPR': a.gPR,
                                                      'gPO': a.gPO,'gCU': a.gCU,} 
                                           )
        episodes = extract_list_from_rosbag_dataset(dataset)
        
        skipped = add_tfs_to_episodes(tf_data, episodes, cfg.names_and_frames, cfg.reference_frame,)
        episodes = [e for i, e in enumerate(episodes) if i not in skipped]  # Remove points with no tfs
        logging.info(f'Saving {len(episodes)} episodes')
        # keypoints_from_frame(episodes, tf_data, cfg.reference_frame)
        
        save_keypoint(episodes, 
                      out_dir/pl.Path(task_folder).name, 
                      description, 
                      demo_idx, 
                      depth_scale=cfg.depth_scale,
                      cameras_used=cfg.cameras_used,
                      episode_to_rlbench_obs=get_rlbench_obs,
                      near_far_planes=cfg.near_far_planes,)
        logging.info('.'*3 + f'Completed demo {demo_idx}' + '.'*3)
    logging.info('Done with task ')


bag_to_rlbench.py
baldeeb/rosbag_data_utils
#####################################################################################################################
#####################################################################################################################
# NOTE:
# Derived from the Google Brain robot_setup repo: 
# https://github.com/ahundt/robotics_setup/blob/master/datasets/google_brain_robot_data/depth_image_encoding.py
# Might contain changes.
#####################################################################################################################
#####################################################################################################################

"""Creates an image from a numpy array of floating point depth data.

Examples:

  depth_array is a 2D numpy array of floating point depth data in meters.

  depth_rgb = FloatArrayToRgbImage(depth_array)
  depth_rgb is a PIL Image object containing the same data as 24-bit
  integers encoded in the RGB bytes.
  depth_rgb.save('image_file.png') - to save to a file.

  depth_array2 = ImageToFloatArray(depth_rgb)
  depth_array2 is a 2D numpy array containing the same data as
  depth_array up to the precision of the RGB image (1/256 mm).

  depth_gray = FloatArrayToGrayImage(depth_array)
  depth_gray is a PIL Image object containing the same data rounded to
  8-bit integers.
  depth_gray.save('image_file.jpg', quality=95) - to save to a file.

  depth_array3 = ImageToFloatArray(depth_gray)
  depth_array3 is a 2D numpy array containing the same data as
  depth_array up to the precision of the grayscale image (1 cm).

The image conversions first scale and round the values and then pack
them into the desired type in a numpy array before converting the
array to a PIL Image object.  The Image can be saved in any format.
We are using PNG for RGB and high quality JPEG for grayscale images.

You can use different numeric types (e.g. np.uint16, np.int32), but
not all combinations of numeric type and image format are supported by
PIL or standard image viewers.

"""

import numpy as np
from PIL import Image
import logging

def ClipFloatValues(float_array, min_value, max_value):
  """Clips values to the range [min_value, max_value].

  First checks if any values are out of range and prints a message.
  Then clips all values to the given range.

  Args:
    float_array: 2D array of floating point values to be clipped.
    min_value: Minimum value of clip range.
    max_value: Maximum value of clip range.

  Returns:
    The clipped array.

  """
  if float_array.min() < min_value or float_array.max() > max_value:
    logging.debug('Image has out-of-range values [%f,%f] not in [%d,%d]',
                  float_array.min(), float_array.max(), min_value, max_value)
    float_array = np.clip(float_array, min_value, max_value)
  return float_array


DEFAULT_RGB_SCALE_FACTOR = 256000.0


def FloatArrayToRgbImage(float_array,
                         scale_factor=DEFAULT_RGB_SCALE_FACTOR,
                         drop_blue=False):
  """Convert a floating point array of values to an RGB image.

  Convert floating point values to a fixed point representation where
  the RGB bytes represent a 24-bit integer.
  R is the high order byte.
  B is the low order byte.
  The precision of the depth image is 1/256 mm.

  Floating point values are scaled so that the integer values cover
  the representable range of depths.

  This image representation should only use lossless compression.

  Args:
    float_array: Input array of floating point depth values in meters.
    scale_factor: Scale value applied to all float values.
    drop_blue: Zero out the blue channel to improve compression, results in 1mm
      precision depth values.

  Returns:
    24-bit RGB PIL Image object representing depth values.
  """
  # Scale the floating point array.
  scaled_array = np.floor(float_array * scale_factor + 0.5)
  # Convert the array to integer type and clip to representable range.
  min_inttype = 0
  max_inttype = 2**24 - 1
  scaled_array = ClipFloatValues(scaled_array, min_inttype, max_inttype)
  int_array = scaled_array.astype(np.uint32)
  # Calculate:
  #   r = (f / 256) / 256  high byte
  #   g = (f / 256) % 256  middle byte
  #   b = f % 256          low byte
  rg = np.divide(int_array, 256)
  r = np.divide(rg, 256)
  g = np.mod(rg, 256)
  image_shape = int_array.shape
  rgb_array = np.zeros((image_shape[0], image_shape[1], 3), dtype=np.uint8)
  rgb_array[..., 0] = r
  rgb_array[..., 1] = g
  if not drop_blue:
    # Calculate the blue channel and add it to the array.
    b = np.mod(int_array, 256)
    rgb_array[..., 2] = b
  image_mode = 'RGB'
  image = Image.fromarray(rgb_array, mode=image_mode)
  return image


DEFAULT_GRAY_SCALE_FACTOR = {np.uint8: 100.0,
                             np.uint16: 1000.0,
                             np.int32: DEFAULT_RGB_SCALE_FACTOR}


def FloatArrayToGrayImage(float_array, scale_factor=None, image_dtype=np.uint8):
  """Convert a floating point array of values to an RGB image.

  Convert floating point values to a fixed point representation with
  the given bit depth.

  The precision of the depth image with default scale_factor is:
    uint8: 1cm, with a range of [0, 2.55m]
    uint16: 1mm, with a range of [0, 65.5m]
    int32: 1/256mm, with a range of [0, 8388m]

  Right now, PIL turns uint16 images into a very strange format and
  does not decode int32 images properly.  Only uint8 works correctly.

  Args:
    float_array: Input array of floating point depth values in meters.
    scale_factor: Scale value applied to all float values.
    image_dtype: Image datatype, which controls the bit depth of the grayscale
      image.

  Returns:
    Grayscale PIL Image object representing depth values.

  """
  # Ensure that we have a valid numeric type for the image.
  if image_dtype == np.uint16:
    image_mode = 'I;16'
  elif image_dtype == np.int32:
    image_mode = 'I'
  else:
    image_dtype = np.uint8
    image_mode = 'L'
  if scale_factor is None:
    scale_factor = DEFAULT_GRAY_SCALE_FACTOR[image_dtype]
  # Scale the floating point array.
  scaled_array = np.floor(float_array * scale_factor + 0.5)
  # Convert the array to integer type and clip to representable range.
  min_dtype = np.iinfo(image_dtype).min
  max_dtype = np.iinfo(image_dtype).max
  scaled_array = ClipFloatValues(scaled_array, min_dtype, max_dtype)

  image_array = scaled_array.astype(image_dtype)
  image = Image.fromarray(image_array, mode=image_mode)
  return image


def ImageToFloatArray(image, scale_factor=None):
  """Recovers the depth values from an image.

  Reverses the depth to image conversion performed by FloatArrayToRgbImage or
  FloatArrayToGrayImage.

  The image is treated as an array of fixed point depth values.  Each
  value is converted to float and scaled by the inverse of the factor
  that was used to generate the Image object from depth values.  If
  scale_factor is specified, it should be the same value that was
  specified in the original conversion.

  The result of this function should be equal to the original input
  within the precision of the conversion.

  Args:
    image: Depth image output of FloatArrayTo[Format]Image.
    scale_factor: Fixed point scale factor.

  Returns:
    A 2D floating point numpy array representing a depth image.

  """
  image_array = np.array(image)
  image_dtype = image_array.dtype
  image_shape = image_array.shape

  channels = image_shape[2] if len(image_shape) > 2 else 1
  assert 2 <= len(image_shape) <= 3
  if channels == 3:
    # RGB image needs to be converted to 24 bit integer.
    float_array = np.sum(image_array * [65536, 256, 1], axis=2)
    if scale_factor is None:
      scale_factor = DEFAULT_RGB_SCALE_FACTOR
  else:
    if scale_factor is None:
      scale_factor = DEFAULT_GRAY_SCALE_FACTOR[image_dtype.type]
    float_array = image_array.astype(np.float32)
  scaled_array = float_array / scale_factor
  return scaled_array
build/lib/rosbag_utils/depth_image_encoding.py
baldeeb/rosbag_data_utils
from geometry_msgs.msg import TransformStamped 
from std_msgs.msg import Time
from tf2_ros import Buffer, TransformException
from rospy import Duration
from typing import List, Optional
from rosbag_utils.rosbag_as_dataset import RosbagReader
from rosbag_utils.ros_to_numpy_helpers import ros_transform_to_np_se3_matrix
import yaml

class Tf2Wrapper:
    def __init__(self, 
                 authority:str='tf2_wrapper', 
                 cache_time:Optional[Duration]=None
                 ):
        self._buffer = Buffer(cache_time, debug=False)
        self._authority = authority
        self._hooks = {}

    def register_getter_hook(self, hook, ref_frame:str=None, frame:str=None):
        if ref_frame is not None and frame is not None:
            self._hooks[(ref_frame, frame)] = hook
        elif ref_frame is not None:
            self._hooks[ref_frame] = hook
        elif frame is not None:
            self._hooks[frame] = hook
        else:
            raise ValueError('Either ref_frame or frame must be specified')

    def set(self, transform:TransformStamped, static:bool):
        if static:
            self._buffer.set_transform_static(transform, self._authority)
        else:
            self._buffer.set_transform(transform, self._authority)

    def _post_process(self, T, ref_frame, frame):
        if (ref_frame, frame) in self._hooks:
            return self._hooks[(ref_frame, frame)](T)
        elif ref_frame in self._hooks:
            return self._hooks[ref_frame](T)
        elif frame in self._hooks:
            return self._hooks[frame](T)
        else:
            return T

    def get(self, reference:str, frame_name:str, timestamp:Time):
        t = self._buffer.lookup_transform(reference, frame_name, timestamp)
        T = ros_transform_to_np_se3_matrix(t.transform)
        T = self._post_process(T, reference, frame_name)
        return T
    
    def __str__(self):
        return self._buffer.all_frames_as_yaml()
    
    def get_as_dict(self):
        return yaml.safe_load(self._buffer.all_frames_as_yaml())


def get_populated_tf2_wrapper(bag_path:str, 
                              static_topics:List[str]= ['/tf_static'], 
                              dynamic_topics:List[str]=['/tf'],
                              cache_time_sec:float=60*2,
                              slack_sec:float=0.1,
                              ):

    def add_topics(topics, is_static):
        for topic in topics:
            dataset = RosbagReader(bag_path, {topic: 't'}, 
                                   permissible_asynchronisity_sec=slack_sec)
            for d in dataset:
                for transfrom in d['t'].transforms:
                    wrapper.set(transfrom, is_static)
    
    wrapper = Tf2Wrapper(cache_time=Duration(cache_time_sec))
    add_topics(static_topics, True)
    add_topics(dynamic_topics, False)

    return wrapper

build/lib/rosbag_utils/ros_tf2_wrapper.py
baldeeb/rosbag_data_utils
from sensor_msgs.msg import (
    Image, 
    CameraInfo,
    )
from geometry_msgs.msg import (
    Pose,
    Point, 
    Quaternion,
    PoseWithCovarianceStamped,
    Transform
    )
import numpy as np
from quaternion import as_rotation_matrix, quaternion, from_rotation_matrix


def ros_camera_info_to_np_intrinsic(info:CameraInfo):
    return np.array([
        [info.K[0], info.K[1], info.K[2]],
        [info.K[3], info.K[4], info.K[5]],
        [info.K[6], info.K[7], info.K[8]]
    ])

def ros_image_to_np(image:Image, depth_to_meters=1e-3):
    H, W = image.height, image.width
    if image.encoding == 'rgb8':
        rgb = np.frombuffer(image.data, dtype=np.byte)
        img = rgb.reshape(H, W, 3).astype(np.uint8)
    elif image.encoding == 'rgba8':
        rgb = np.frombuffer(image.data, dtype=np.byte)
        img = rgb.reshape(H, W, 4).astype(np.uint8)
    elif image.encoding == 'bgra8':
        rgb = np.frombuffer(image.data, dtype=np.byte)
        img = rgb.reshape(H, W, 4)[:, :, (2,1,0)].astype(np.uint8)
    elif image.encoding == '16UC1':
        d = np.frombuffer(image.data, dtype=np.uint16).reshape(H, W)
        img = d.astype(np.float32) * depth_to_meters
    elif image.encoding == 'bgra8':
        rgbd = np.frombuffer(image.data, dtype=np.byte)
        rgbd = rgbd.reshape(H, W, 4)[:, :, (2,1,0)].astype(np.uint8)
        img = rgbd[:,:,3].astype(np.uint16).astype(np.float32) * depth_to_meters
    else: 
        raise RuntimeError(f'Image to Numpy is not setup to handle {image.encoding}.')
    return img

def _ros_point_to_np(point:Point):
    return np.array([point.x, point.y, point.z])

def _ros_quaternion_to_np(quat:Quaternion):
    q = quaternion(quat.w, quat.x, quat.y, quat.z)
    return as_rotation_matrix(q)

def ros_pose_to_np_se3_matrix(pose:Pose):
    mat = np.eye(4)
    mat[:3,:3] = _ros_quaternion_to_np(pose.orientation)
    mat[:3,3] = _ros_point_to_np(pose.position)
    return mat

def ros_pose_to_7d_position_quaternion(pose:Pose):
    vec = np.empty(7)
    vec[:3] = _ros_point_to_np(pose.position)
    vec[3:] = np.array([pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w])
    return vec
    
def ros_pose_with_cov_stamped_to_np_se3_matrix(pose:PoseWithCovarianceStamped):
    return ros_pose_to_np_se3_matrix(pose.pose.pose)

def ros_transform_to_np_se3_matrix(transform:Transform):
    mat = np.eye(4)
    mat[:3,:3] = _ros_quaternion_to_np(transform.rotation)
    mat[:3,3] = _ros_point_to_np(transform.translation)
    return mat

def se3_mat_to_position_and_quaterion_vec(mat):
    vec = np.empty(7)
    vec[:3] = mat[:3,3]
    q = from_rotation_matrix(mat[:3,:3])
    vec[3:] = np.array([q.x, q.y, q.z, q.w]) 
    return vec

def ros_joint_states_to_numpy(msg):
    return {
        'names': msg.name,
        'positions': np.array(msg.position),
        'velocities': np.array(msg.velocity),
        'efforts': np.array(msg.effort),
    }


build/lib/rosbag_utils/ros_to_numpy_helpers.py
baldeeb/rosbag_data_utils
import rospy
import rosbag
from rosbag_utils.ros_to_numpy_helpers import (
    ros_camera_info_to_np_intrinsic,
    ros_image_to_np,
    ros_pose_to_np_se3_matrix,
    ros_joint_states_to_numpy,
)
import logging

def get_topics_in_path(path:str):
    bag = rosbag.Bag(path)
    return list(bag.get_type_and_topic_info()[1])

class MissingTopicError(Exception):
    def __init__(self, message, topics):
        super().__init__(message)
        self.topics = topics

class RosbagReader:
    ''' 
    Takes a ros bag and returns one item at a time.
    Assumes that the bag topics are logged in a normally distributed fashion. 
    ''' 
    def __init__(self, 
                 path, 
                 topics_and_names=None, 
                 default_topic_process=True, 
                 anytime_topics=[],
                 permissible_asynchronisity_sec=None,
                 reference_topic=None):
        '''
        Args:
            path: path to the rosbag file
            topics_and_names: a dictionary of topic names to be used in the output
            default_topic_process: if True, the reader will use default postprocess functions
        '''
        self.bag = rosbag.Bag(path)

        if topics_and_names is not None:
            available_topics = self.get_topics()
            desired_topics = list(topics_and_names.keys())
            self._validate_requested_topics(desired_topics, available_topics)
            self._topics = desired_topics
        else:
            self._topics = self.get_topics()

        # Select reference topic
        if reference_topic is not None:
            self._len, self._time_slots = self._get_reference_topic_info(reference_topic)
        else:
            reference_topic, self._len, self._time_slots = self._get_topic_with_least_messages()
        
        slack_sec = self._resolve_timing_details(permissible_asynchronisity_sec)
        self._resolve_naming_details(topics_and_names)

        self._anytime_topics = anytime_topics

        self._postprocess_type = {}
        self._postprocess_topic = {}
        self._postprocess_output_name = {}

        if default_topic_process:
            self._postprocess_type = {
                'sensor_msgs/Image': ros_image_to_np,
                'geometry_msgs/PoseWithCovarianceStamped': ros_pose_to_np_se3_matrix,
                'sensor_msgs/CameraInfo': ros_camera_info_to_np_intrinsic,
                'sensor_msgs/JointState': ros_joint_states_to_numpy,
            }

        logging.info(f'Initialized BagReader - ref topic: {reference_topic} - time slack {slack_sec}.')


    def _resolve_timing_details(self, slack_sec=None):
        '''Determines time window around said topic's messages to search for other topics'''
        self._start, self._end = self.bag.get_start_time(), self.bag.get_end_time()
        if slack_sec is not None: 
            self._time_eps = rospy.Time.from_sec(slack_sec)
        else: 
            slack_sec = (self._end - self._start) / self.__len__()
            self._time_eps = rospy.Time.from_sec(slack_sec)
        return slack_sec

    def _resolve_naming_details(self, topics_and_names):
        '''Figures out what names to use for what topics in output.'''
        if topics_and_names is not None:
            self._name_of_topic = topics_and_names
        else: 
            self._name_of_topic = {}
            for t in self._topics:
                if t not in self._name_of_topic:
                    self._name_of_topic[t] = t

    def _validate_requested_topics(self, desired, available):
        if not all([t in available for t in desired]):
            not_found = [t for t in desired if t not in available]
            raise RuntimeError(f"The topics {not_found} are not available in the bag.\n" +
                               f'available topics: {available}')    
    
    def __iter__(self):
        for i in range(self.__len__()):
            yield self[i]


    def clear_postprocess(self):
        self._postprocess_type = {}
    def unregister_type_postprocess(self, msg_type:str):
        del self._postprocess_type[msg_type]
    def register_type_postprocess(self, msg_type:str, func:callable):
        self._postprocess_type[msg_type] = func


    def register_topic_postprocess(self, topic:str, func:callable):
        self._postprocess_topic[topic] = func
    def unregister_topic_postprocess(self, topic:str):
        del self._postprocess_topic[topic]


    def register_output_postprocess(self, msg_name:str, func:callable):
        self._postprocess_output_name[msg_name] = func
    def unregister_output_postprocess(self, msg_name:str):
        del self._postprocess_output_name[msg_name]


    def __len__(self): return self._len


    def _get_topic_with_least_messages(self):
        topic = self._topics[0]
        count = self.bag.get_message_count(topic)
        for t in self._topics[1:]:
            c = self.bag.get_message_count(t)
            if c < count: count, topic = c, t
        time_sec = [m[2] for m in self.bag.read_messages(topics=[topic])]
        return topic, count, time_sec  

    def _get_reference_topic_info(self, ref):
        count = self.bag.get_message_count(ref)
        time_sec = [m[2] for m in self.bag.read_messages(topics=[ref])]
        return count, time_sec  


    def get_topics(self):
        return list(self.bag.get_type_and_topic_info()[1])


    def _msgs_at(self, time:rospy.Time, rigourous=True, anytime=False, topics=None):
        '''Ensures that messages are returned by proximity from intended time.'''
        if topics is None: topics = self._topics
        if anytime: 
            rigourous = True
            s = e = None
        else:
            t_time, elapsed = time.to_time(), self._time_eps.to_time()
            s = rospy.Time.from_sec(t_time - elapsed)
            e = rospy.Time.from_sec(t_time + elapsed)

        msgs = self.bag.read_messages(topics = topics, 
                                      start_time = s, 
                                      end_time = e)
        if rigourous:
            msgs = list(msgs)
            dt = [abs(m[2] - time) for m in msgs]
            idxs = sorted(range(len(dt)), key=lambda k: dt[k])
            for i in idxs:
                yield msgs[i]
        else:
            for m in msgs:
                yield m


    def _process(self, topic, name, msg):
        data = {}
        if topic in self._postprocess_topic:           # if topic has a postprocess function
            data[name] = self._postprocess_topic[topic](msg)
        elif msg._type in self._postprocess_type:    # if message type has a postprocess function
            data[name] = self._postprocess_type[msg._type](msg)
        else:                                           # no postprocess function
            data[name] = msg

        if name in self._postprocess_output_name:   # if generated data has a postprocess function
            data[name] = self._postprocess_output_name[name](data[name])
        return data


    def __getitem__(self, idx):
        time = self._time_slots[idx]
        data = {}
        for msg in self._msgs_at(time):
            msg_name = self._name_of_topic[msg[0]]
            if msg_name in data: continue                           # Continue, if topic already obtained
            data.update(self._process(msg[0], msg_name, msg[1]))    # Add msg to data
            if len(data) == len(self._topics): break                # All topics collected
        
        # If some topics are set to time agnostic, try to get them from anytime
        if len(data) != len(self._topics):
            if len(self._anytime_topics) > 0:
                missing = []
                for topic, name in self._name_of_topic.items():
                    if name not in data:
                        if name in self._anytime_topics:
                            msg = next(iter(self._msgs_at(time, anytime=True, 
                                                        topics=[topic])))[1]
                            data.update(self._process(topic, name, msg))
                        else:
                            missing.append(name)
                if len(missing) > 0:
                    raise MissingTopicError(f'Not all topics were found at time {time}.', 
                                            topics=missing)
            else:
                raise MissingTopicError(f'Not all topics were found at time {time}.', 
                                        topics=[k for k in data if k not in self._topics])
        return data

    def __del__(self):
        self.bag.close()

build/lib/rosbag_utils/rosbag_as_dataset.py
baldeeb/rosbag_data_utils
import os
import sys
import logging
import pathlib as pl
from copy import deepcopy
from typing import List, Dict, Optional

import cv2
import pickle
import numpy as np
from PIL import Image

from rosbag_utils.rosbag_as_dataset import RosbagReader, MissingTopicError
from rosbag_utils.ros_to_numpy_helpers import (ros_pose_to_np_se3_matrix,
                                                     ros_image_to_np,
                                                     se3_mat_to_position_and_quaterion_vec)
from rosbag_utils.ros_tf2_wrapper import (get_populated_tf2_wrapper,
                                                Tf2Wrapper)

from rlbench.backend.observation import Observation
from rlbench.demo import Demo

from rosbag_utils.depth_image_encoding import FloatArrayToRgbImage, DEFAULT_RGB_SCALE_FACTOR
from rospy import Time, Duration
from tf2_ros import tf2

ALL_CAMERAS = [
    'left_shoulder',
    'right_shoulder',
    'front',
    'overhead',
    'wrist'
]


def check_and_mkdirs(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path, exist_ok=True)


def save_img(img, path):
    if isinstance(img, np.ndarray):
        cv2.imwrite(path, img[:, :, :3])
    elif isinstance(img, Image.Image):
        img.save(path)


def extract_list_from_rosbag_dataset(dataset:RosbagReader) -> List[Dict]:
    '''This is where the assumptions about the dataset lie.
    We are assuming that there is a series of datapoints where all topics are present.
    This function finds the start and iterates until it hits the first missing topic.
    
    Args:
        dataset: RosbagReader object
    Returns: 
        List[Dict] where each dict is a datapoint
    '''

    starting_idx = 0
    datalist = []
    for i in range(starting_idx, len(dataset)):
        try:
            datalist.append(dataset[i])
        except MissingTopicError as e:
            logging.info(f'skipped episode {i} - missing: {e.topics}.')
    return datalist


def get_dummy_observation(frame, misc):

    # pose format -> x, y, z, qx, qy, qz, qw
    dummy_pose = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0 ,1.0])

    obs = Observation(
        left_shoulder_rgb=None,
        left_shoulder_depth=None,
        left_shoulder_point_cloud=None,
        right_shoulder_rgb=None,
        right_shoulder_depth=None,
        right_shoulder_point_cloud=None,
        overhead_rgb=None,
        overhead_depth=None,
        overhead_point_cloud=None,
        wrist_rgb=None,
        wrist_depth=None,
        wrist_point_cloud=None,
        front_rgb=None,
        front_depth=None,
        front_point_cloud=None,
        left_shoulder_mask=None,
        right_shoulder_mask=None,
        overhead_mask=None,
        wrist_mask=None,
        front_mask=None,
        joint_velocities            = dummy_pose,
        joint_positions             = dummy_pose,
        joint_forces                = dummy_pose,
        gripper_open                = 0.9,
        gripper_pose                = dummy_pose,
        gripper_matrix              = dummy_pose,
        gripper_touch_forces        = None,
        gripper_joint_positions     = np.ones(2),
        task_low_dim_state          =None,
        ignore_collisions           =True, # TODO: fix
        misc=misc,
    )

    return obs


def save_keypoint(episodes:List[Dict], 
                  save_path:str, 
                  description:str,
                  episode_idx:int=0, 
                  depth_scale=DEFAULT_RGB_SCALE_FACTOR,
                  cameras_used=ALL_CAMERAS,
                  episode_to_rlbench_obs=get_dummy_observation,
                  near_far_planes = [0, 1]
                  ):
    '''TODO: add docstring'''

     # make directories
    variation_idx = 0

    episode_path = os.path.join(save_path, 'all_variations', 'episodes', f"episode{episode_idx}")
    check_and_mkdirs(episode_path)

    base_misc = dict()

    for cam in cameras_used:
        for info in ['rgb', 'depth', 'camera_intrinsics',]:
            assert any([f'{cam}_{info}' in k for k in episodes[0].keys()]),\
                   f"Camera {cam}_{info} not found in frames!"

        image_dir = pl.Path(episode_path)/f'{cam}_rgb'
        image_dir.mkdir(parents=True, exist_ok=True)

        image_dir = pl.Path(episode_path)/f'{cam}_depth'
        image_dir.mkdir(parents=True, exist_ok=True)

        K_key = f'{cam}_camera_intrinsics'
        base_misc[K_key] = np.array(episodes[0][K_key]).reshape(3,3)
        base_misc[cam + '_camera_near'] = near_far_planes[0]                           # TODO: scrutinize these values
        base_misc[cam + '_camera_far'] = near_far_planes[1]                            # TODO: scrutinize these values

    base_misc['keypoint_idxs'] = [sum([1 for e in episodes if 'keypoint' in e]) - 1]

    observations = []
    image_keys = [f'{cam}_rgb' for cam in cameras_used]
    depth_keys = [f'{cam}_depth' for cam in cameras_used]
    for eps_idx, eps in enumerate(episodes):
        misc = deepcopy(base_misc)  
        found_image = False      
        for key, val in eps.items():
            image_dir = str(pl.Path(episode_path) / key / f'{eps_idx}.png')
            if key in image_keys: 
                save_img(cv2.cvtColor(val['image'], cv2.COLOR_RGB2BGR),
                         image_dir)
                misc[f'{key[:-4]}_camera_extrinsics'] = val['extrinsics']
                found_image = True
            elif key in depth_keys: 
                depth_img = (val['image'] - near_far_planes[0]) / (near_far_planes[1] - near_far_planes[0])
                v = FloatArrayToRgbImage(depth_img, scale_factor=depth_scale)
                save_img(v, image_dir)
                found_image = True
        assert found_image, f"Could not find image in episode {eps_idx}."
        if 'keypoint' in eps:
            misc['keypoint'] = eps['keypoint']
        observations.append(episode_to_rlbench_obs(eps, misc))

    demo = Demo(observations, random_seed=0)
    demo.variation_number = variation_idx

    low_dim_obs_path = os.path.join(episode_path, 'low_dim_obs.pkl')
    with open(low_dim_obs_path, 'wb') as f:
        pickle.dump(demo, f)

    variation_number_path = os.path.join(episode_path, 'variation_number.pkl')
    with open(variation_number_path, 'wb') as f:
        pickle.dump(variation_idx, f)

    descriptions_path = os.path.join(episode_path, 'variation_descriptions.pkl')
    with open(descriptions_path, 'wb') as f:
        pickle.dump([description], f)

    logging.info(f"Saved {len(episodes)} frames to {save_path}")


def _setup_rosbag_dataset(data_dir, 
                          topics_and_names:dict, 
                          time_slack:Optional[float]=None, 
                          reference_topic:str=None
                          ):
    dataset = RosbagReader(data_dir, topics_and_names, 
                           permissible_asynchronisity_sec=time_slack,
                           reference_topic=reference_topic)
    def task_status_postprocessing(msg):
        return {'name': msg.task_name,
                'pose': ros_pose_to_np_se3_matrix(msg.target_pose.pose),
                'reference_frame': msg.target_pose.header.frame_id,
                'time': msg.target_pose.header.stamp,
                'status': msg.status,}
    def image_and_time(msg):
        return {'image': ros_image_to_np(msg),
                'time': msg.header.stamp}
    # instead of just getting image get its timestamp as well
    dataset.unregister_type_postprocess('sensor_msgs/Image')
    dataset.register_type_postprocess('sensor_msgs/Image', image_and_time)
    # package task status into a keypoint dict
    dataset.register_topic_postprocess('/spot/task_state', task_status_postprocessing)
    logging.info(f'Dataset initial length: {len(dataset)}')
    
    return dataset

def add_tfs_to_episodes(tf_wrapper:Tf2Wrapper, data:List[Dict], 
                    names_and_frames:Dict, ref_frame:str,
                    try_handle_extrapolation:bool=True,
                    offset_t:Duration=Duration(0,0),
                    ):
    '''Adds the extrinsics to the data. Optionally ignores missing data.
    Args:
        tf_wrapper: Tf2Wrapper object containing buffer of all tfs.
        data: List of dicts containing the data into which the extrinsics will be added.
        names_and_frames: dictionary of the data keys (or names) that are to be augmented
            with the extrinsics (or frames).
        ref_frame: The frame of reference for all the extrinsics.
    '''
    skipped = []
    for data_i, d in enumerate(data):
        try:
            for name, frame in names_and_frames.items():
                if name not in d.keys(): 
                    d_time = [val['time'] for val in d.values() 
                              if isinstance(val, dict) and 'time' in val][0]
                    time = Time(secs=d_time.secs + offset_t.secs, 
                                nsecs=d_time.nsecs + offset_t.nsecs)
                    d[name] = tf_wrapper.get(ref_frame, frame, time)
                else:
                    d_time = d[name]['time']
                    time = Time(secs=d_time.secs + offset_t.secs, 
                                nsecs=d_time.nsecs + offset_t.nsecs)
                    d[name]['extrinsics'] = tf_wrapper.get(ref_frame, frame, time)
        except tf2.ExtrapolationException as e:
            if try_handle_extrapolation:
                if data_i == 0 and offset_t == Duration(0,0):
                    tf_dict = tf_wrapper.get_as_dict()
                    earliest = max([tf_dict[f]['oldest_transform'] 
                                    for f in [frame, ref_frame]])
                    return add_tfs_to_episodes(tf_wrapper, data, names_and_frames, ref_frame, 
                                               try_handle_extrapolation, 
                                               offset_t=Time.from_sec(earliest) - time + Duration(0.001,0))
                elif data_i > int(0.8 * len(data)):
                    skipped.append(data_i)
                    logging.warning(f"Skipped {data_i} frame due to mismatched timestamps between tfs and data.")
                else:
                    logging.error(f'Failed to handle extrapolation exception at data index: {data_i}.')
                    raise e
            else: 
                raise e
        
    if len(skipped) > 0: logging.warning(f"Skipped frames {skipped}.")
    
    return skipped


def keypoints_from_frame(data:List[Dict], tf_wrapper:Tf2Wrapper, ref_frame:str):
    '''Transforms all keypoint poses to the reference frame.'''
    for d in data:
        k = d['keypoint']
        #### NOTE: TEMPORARY FIX!!!
        ## This should be the data time.
        time = Time(secs=d['left_shoulder_rgb']['time'].secs, nsecs=d['left_shoulder_rgb']['time'].nsecs)
        # time = Time(secs=k['time'].secs, nsecs=k['time'].nsecs)
        T = tf_wrapper.get(ref_frame, k['reference_frame'], time)
        k['pose'] = se3_mat_to_position_and_quaterion_vec(T @ k['pose'])
        k['reference_frame'] = ref_frame






###################################################################
###################################################################
###################################################################
###################################################################
# TODO: remove below functions
###################################################################
###################################################################


def rosbag_to_rlbench(data_dir:pl.Path, 
                      out_dir:pl.Path, 
                      topics_and_names:Dict, 
                      names_and_frames:Dict,
                      reference_frame:str, 
                      task_description:List[str]=["no lang goal defined"],
                      episode_idx:int = 0,
                      depth_scale = ( ( 2**22 ) - 1.0 ),
                      cameras_used = ALL_CAMERAS,
                      time_slack:Optional[float]=None,
                      reference_topic:Optional[str]=None,
                      reference_frame_T:np.ndarray=np.eye(4),
                      ) -> None:
    '''
    Args:
        names_and_frames: dictionary of all the datapoints' names 
            and tf frame names. This list indicates which datapoints
            we care to associate with a pose. Mainly intended for
            camera data extrinsics.
        reference: The frame of reference for all the extrinsics.
    '''
    logging.info('.'*3 + f'processing: {data_dir}' + '.'*3)
    tf_data = get_populated_tf2_wrapper(data_dir)
    tf_data.register_getter_hook(lambda x: reference_frame_T @ x, reference_frame)
    dataset = _setup_rosbag_dataset(data_dir, topics_and_names,
                                    time_slack=time_slack,
                                    reference_topic=reference_topic)

    logging.info(f'depth scale: {depth_scale}')
    episodes = extract_list_from_rosbag_dataset(dataset)
    add_tfs_to_episodes(tf_data, 
                        episodes, 
                        names_and_frames, 
                        reference_frame)
    keypoints_from_frame(episodes, 
                         tf_data, 
                         reference_frame)
    save_keypoint(episodes, 
                  out_dir, 
                  task_description, 
                  episode_idx, 
                  depth_scale=depth_scale,
                  cameras_used=cameras_used)
    logging.info('.'*5 + 'Done!' + '.'*5)


def rosbags_to_rlbench_episodes(data_dir:pl.Path, 
                                out_dir:pl.Path, 
                                topics_and_names:Dict,
                                names_and_frames, 
                                reference_frame,
                                task_description:List[str]=["no lang goal defined"],
                                depth_scale = ( ( 2**22 ) - 1.0 ),
                                cameras_used:List[str]=ALL_CAMERAS,
                                time_slack:Optional[float]=None,
                                reference_topic:Optional[str]=None,
                                reference_frame_T:np.ndarray=np.eye(4),
                                ) -> None:

    for i, data_path in enumerate(data_dir.glob('*.bag')):
        rosbag_to_rlbench(data_path, 
                          out_dir/data_dir.name, 
                          topics_and_names, 
                          names_and_frames, 
                          reference_frame,
                          task_description,
                          i,
                          depth_scale=depth_scale,
                          cameras_used=cameras_used,
                          time_slack=time_slack,
                          reference_topic=reference_topic,
                          reference_frame_T=reference_frame_T,
                          )
    logging.getLogger(__name__).info('Done!')

build/lib/rosbag_utils/rosbag_to_rlbench.py
baldeeb/rosbag_data_utils
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import logging
import numpy as np
import pathlib as pl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from data_preprocessing.ros_tf2_wrapper import (get_populated_tf2_wrapper)
from data_preprocessing.rosbag_to_rlbench import (keypoints_from_frame, 
                                                  extract_list_from_rosbag_dataset, 
                                                  add_tfs_to_episodes, 
                                                  _setup_rosbag_dataset)

from quaternion import quaternion, as_rotation_matrix, from_rotation_matrix

###### Settings ##########################################
SHOW_2D = True
SHOW_3D = False
REFERENCE_FRAME = 'body' # 'map' or 'body'
PERSPECTIVE = 'other'  # 'waypoint' or other
                            # FOR VISUALIZATION - using map frame will yield static waypoints and meaningless visuals.
NUMBER_POINTS = False
##########################################################

if __name__ == '__main__':


    logging.basicConfig(level=logging.INFO)

    # data_dir = pl.Path('/media/baldeeb/ssd2/Data/SpotBags/2023_08_01')
    # topics_and_names = {
    #     '/spot/camera/frontleft/image':                         'left_shoulder_rgb',
    #     '/spot/depth/frontleft/depth_in_visual':                'left_shoulder_depth',
    #     '/spot/camera/frontleft/camera_info':                   'left_shoulder_camera_intrinsics',
    #     # '/spot/depth/frontleft/depth_in_visual/camera_info':    '',

    #     '/spot/camera/frontright/image':                        'right_shoulder_rgb',
    #     '/spot/depth/frontright/depth_in_visual':               'right_shoulder_depth',
    #     '/spot/camera/frontright/camera_info':                  'right_shoulder_camera_intrinsics',
    #     # '/spot/depth/frontright/depth_in_visual/camera_info':   '',

    #     '/spot/camera/right/image':                             'front_rgb' ,
    #     '/spot/depth/right/image':                              'front_depth',
    #     '/spot/camera/right/camera_info':                       'front_camera_intrinsics',

    #     '/spot/camera/left/image':                              'overhead_rgb' ,
    #     '/spot/depth/left/image':                               'overhead_depth',
    #     '/spot/camera/left/camera_info':                        'overhead_camera_intrinsics',

    #     # '/spot/camera/hand_color/image':                        'wrist_rgb' ,
    #     # '/spot/depth/hand/depth_in_color':                      'wrist_depth',
    #     # '/spot/camera/hand_color/camera_info':                  'wrist_camera_intrinsics',

    #     '/spot/task_state':                                     'keypoint',
    #     }



    # data_dir = pl.Path('/media/baldeeb/ssd2/Data/SpotBags/2023_08_08')

    # topics_and_names = {
    #     '/spot/camera/frontleft/image':                         'left_shoulder_rgb',
    #     '/spot/depth/frontleft_in_visual/image':                'left_shoulder_depth',
    #     '/spot/camera/frontleft/camera_info':                   'left_shoulder_camera_intrinsics',

    #     '/spot/camera/frontright/image':                        'right_shoulder_rgb',
    #     '/spot/depth/frontright_in_visual/image':               'right_shoulder_depth',
    #     '/spot/camera/frontright/camera_info':                  'right_shoulder_camera_intrinsics',

    #     '/spot/camera/right/image':                             'front_rgb' ,
    #     '/spot/depth/right/image':                              'front_depth',
    #     '/spot/camera/right/camera_info':                       'front_camera_intrinsics',

    #     '/spot/camera/left/image':                              'overhead_rgb' ,
    #     '/spot/depth/left/image':                               'overhead_depth',
    #     '/spot/camera/left/camera_info':                        'overhead_camera_intrinsics',

    #     # '/spot/camera/hand_color/image':                        'wrist_rgb' ,
    #     # '/spot/depth/hand_in_visual/image':                      'wrist_depth',
    #     # '/spot/camera/hand_color/camera_info':                  'wrist_camera_intrinsics',

    #     '/spot/task_state':                                     'keypoint',
    #     }

    # reference_frame = REFERENCE_FRAME
    # names_and_frames = {
    #     'left_shoulder_rgb':    'frontleft_fisheye',
    #     'right_shoulder_rgb':   'frontright_fisheye',
    #     'front_rgb':            'right_fisheye',
    #     'overhead_rgb':         'left_fisheye',
    #     # 'wrist_rgb':            'hand_color_image_sensor',
    #     'body_T_vision':        'vision',  # The waypoint is in the vision frame
    #     }
    


    data_dir = pl.Path('/media/baldeeb/ssd2/Data/SpotBags/2023_08_25_eval')

    topics_and_names = {
        '/camera1/color/image_raw':                              'left_shoulder_rgb',
        # '/camera1/depth/image_rect_raw':                         'left_shoulder_depth',
        '/camera1/aligned_depth_to_color/image_raw':             'left_shoulder_depth',
        '/camera1/color/camera_info':                            'left_shoulder_camera_intrinsics',

        '/camera2/color/image_raw':                              'right_shoulder_rgb',
        # '/camera2/depth/image_rect_raw':                         'right_shoulder_depth',
        '/camera2/aligned_depth_to_color/image_raw':             'right_shoulder_depth',
        '/camera2/color/camera_info':                            'right_shoulder_camera_intrinsics',

        '/spot/task_state':                                     'keypoint',
    }
    time_slack = 0.1
    reference_frame = REFERENCE_FRAME
    names_and_frames = {
        'left_shoulder_rgb':    'camera1_color_optical_frame',
        'right_shoulder_rgb':   'camera2_color_optical_frame',
        }


    def to_Rt(qt, pt):
        Rt = np.eye(4)
        Rt[:3, :3] = as_rotation_matrix(qt)
        Rt[:3, 3] = pt
        Rt_inv = np.linalg.inv(Rt)
        return Rt, Rt_inv

    def get_pose_details(pose,
                         perspective='body'  # 'body' or 'waypoint'
                         ):
        x,y,z,qx,qy,qz,qw = pose
        offset = np.array([0.05, 0.0, 0.0])
        q = quaternion(qw, qx, qy, qz)
        p = np.array([x,y,z])

        waypointTperspective, perspectiveTwaypoint = to_Rt(q, p)
        
        def _pt_and_dir(T):
            return T[:3, 3], T[:3, :3] @ offset

        if perspective == 'waypoint':
            ## To display the body in the waypoint frame
            p, dir_pt =  _pt_and_dir(perspectiveTwaypoint)
        else:
            ## To display the waypoint in the body frame
            p, dir_pt =  _pt_and_dir(waypointTperspective)

        # bodyTwaypoint = as_rotation_matrix(q)
        # dir_pt = bodyTwaypoint @ offset
        return p, np.array(dir_pt)

    group = []
    bag_paths = list(data_dir.glob('*.bag'))
    for i, data_path in enumerate(bag_paths):
        logging.info('*'*20 + f'\nprocessing: {data_path}')
        # Accumulates all tfs from a bag in one buffer
        tf_data = get_populated_tf2_wrapper(data_path, slack_sec=0.1)

        # Extracts the data from the bag as a list of dicts.
        #   The dict is described by topics_and_names.
        #   The result discards any leading or training frames that had missing data.
        dataset = _setup_rosbag_dataset(data_path, topics_and_names, time_slack=time_slack, reference_topic='/spot/task_state')
        episode = extract_list_from_rosbag_dataset(dataset)
        skipped = add_tfs_to_episodes(tf_data, episode, names_and_frames, reference_frame, handle_excepts=True)

        episode = [e for i, e in enumerate(episode) if i not in skipped]
        if len(episode) < 2:
            logging.warning(f"Skipping {data_path} due to insufficient data.")
            continue
        keypoints_from_frame(episode, tf_data, reference_frame)

        # Validate all waypoint states
        # assert episode[-1]['keypoint']['status'] == 'success', f"Last waypoint is not success: {episode[-1]['keypoint']['status']}"
        # for i in range(1, len(episode)-1):
        #     assert episode[i]['keypoint']['status'] == 'running', f"Waypoint {i} is not progress: {episode[i]['keypoint']['status']}"
         
        # accumulate data to plot
        points, dir_pts = [], []
        for f in episode:
            # Note that the waypoints currently are represented in the body frame
            #   but we want body relative to the waypoint frame.
            pt, dir_pt = get_pose_details(f['keypoint']['pose'], perspective=PERSPECTIVE)
            points.append(pt); dir_pts.append(dir_pt)

        group.append([np.array(points), np.array(dir_pts), episode[-1]['keypoint']['status']])

    #### Plotting #######################################
    if SHOW_2D:
        # 2D
        fig = plt.figure()
        # plt.scatter(0, 0, c='g', alpha=0.25)
        # plt.plot([0, 0.2], [0, 0], c='g', alpha=0.25)
        legend = [p.name for p in bag_paths]
        for gi, (pts, dpts, status) in enumerate(group):
            plt.scatter(pts[:,0], pts[:,1], alpha=0.25)
            for pi, (p, frame) in enumerate(zip(pts, dpts)):
                x, y = [p[0], p[0]+frame[0]], [p[1], p[1]+frame[1]]
                plt.plot(x, y, c='r', alpha=0.25)                    
                if NUMBER_POINTS:
                    plt.annotate(f'{pi}', (p[0], p[1]))
            plt.plot(pts[:,0], pts[:,1], alpha=0.15, c='blue', label=legend[gi])
        plt.show()


    if SHOW_3D:
        # 3D  -  buggy
        for i, (pts, dpts, status) in enumerate(group):
            path = bag_paths[i]
            fig = plt.figure()
            points = np.array(pts)
            dir_pts = np.array(dpts)

            ax = fig.add_subplot(111, projection='3d')
            ax.title.set_text(f"{bag_paths[i]} \n {status}")        
            ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
            
            # Robot
            ax.scatter(0, 0, 0, c='g', alpha=0.25)
            ax.plot([0, 0.2], [0, 0], [0, 0], c='g', alpha=0.25)

            # Points
            ax.scatter(points[:,0], points[:,1], points[:,2], alpha=0.25)
            
            # Directions
            for p, frame in zip(points, dir_pts):
                ax.plot([p[0], p[0]+frame[0]], [p[1], p[1]+frame[1]], [p[2], p[2]+frame[2]], c='r', alpha=0.25)
        plt.show()

    logging.getLogger(__name__).info('Done!')

examples/visualize_waypoint_distribution.py
baldeeb/rosbag_data_utils
#####################################################################################################################
#####################################################################################################################
# NOTE:
# Derived from the Google Brain robot_setup repo: 
# https://github.com/ahundt/robotics_setup/blob/master/datasets/google_brain_robot_data/depth_image_encoding.py
# Might contain changes.
#####################################################################################################################
#####################################################################################################################

"""Creates an image from a numpy array of floating point depth data.

Examples:

  depth_array is a 2D numpy array of floating point depth data in meters.

  depth_rgb = FloatArrayToRgbImage(depth_array)
  depth_rgb is a PIL Image object containing the same data as 24-bit
  integers encoded in the RGB bytes.
  depth_rgb.save('image_file.png') - to save to a file.

  depth_array2 = ImageToFloatArray(depth_rgb)
  depth_array2 is a 2D numpy array containing the same data as
  depth_array up to the precision of the RGB image (1/256 mm).

  depth_gray = FloatArrayToGrayImage(depth_array)
  depth_gray is a PIL Image object containing the same data rounded to
  8-bit integers.
  depth_gray.save('image_file.jpg', quality=95) - to save to a file.

  depth_array3 = ImageToFloatArray(depth_gray)
  depth_array3 is a 2D numpy array containing the same data as
  depth_array up to the precision of the grayscale image (1 cm).

The image conversions first scale and round the values and then pack
them into the desired type in a numpy array before converting the
array to a PIL Image object.  The Image can be saved in any format.
We are using PNG for RGB and high quality JPEG for grayscale images.

You can use different numeric types (e.g. np.uint16, np.int32), but
not all combinations of numeric type and image format are supported by
PIL or standard image viewers.

"""

import numpy as np
from PIL import Image
import logging

def ClipFloatValues(float_array, min_value, max_value):
  """Clips values to the range [min_value, max_value].

  First checks if any values are out of range and prints a message.
  Then clips all values to the given range.

  Args:
    float_array: 2D array of floating point values to be clipped.
    min_value: Minimum value of clip range.
    max_value: Maximum value of clip range.

  Returns:
    The clipped array.

  """
  if float_array.min() < min_value or float_array.max() > max_value:
    logging.debug('Image has out-of-range values [%f,%f] not in [%d,%d]',
                  float_array.min(), float_array.max(), min_value, max_value)
    float_array = np.clip(float_array, min_value, max_value)
  return float_array


DEFAULT_RGB_SCALE_FACTOR = 256000.0


def FloatArrayToRgbImage(float_array,
                         scale_factor=DEFAULT_RGB_SCALE_FACTOR,
                         drop_blue=False):
  """Convert a floating point array of values to an RGB image.

  Convert floating point values to a fixed point representation where
  the RGB bytes represent a 24-bit integer.
  R is the high order byte.
  B is the low order byte.
  The precision of the depth image is 1/256 mm.

  Floating point values are scaled so that the integer values cover
  the representable range of depths.

  This image representation should only use lossless compression.

  Args:
    float_array: Input array of floating point depth values in meters.
    scale_factor: Scale value applied to all float values.
    drop_blue: Zero out the blue channel to improve compression, results in 1mm
      precision depth values.

  Returns:
    24-bit RGB PIL Image object representing depth values.
  """
  # Scale the floating point array.
  scaled_array = np.floor(float_array * scale_factor + 0.5)
  # Convert the array to integer type and clip to representable range.
  min_inttype = 0
  max_inttype = 2**24 - 1
  scaled_array = ClipFloatValues(scaled_array, min_inttype, max_inttype)
  int_array = scaled_array.astype(np.uint32)
  # Calculate:
  #   r = (f / 256) / 256  high byte
  #   g = (f / 256) % 256  middle byte
  #   b = f % 256          low byte
  rg = np.divide(int_array, 256)
  r = np.divide(rg, 256)
  g = np.mod(rg, 256)
  image_shape = int_array.shape
  rgb_array = np.zeros((image_shape[0], image_shape[1], 3), dtype=np.uint8)
  rgb_array[..., 0] = r
  rgb_array[..., 1] = g
  if not drop_blue:
    # Calculate the blue channel and add it to the array.
    b = np.mod(int_array, 256)
    rgb_array[..., 2] = b
  image_mode = 'RGB'
  image = Image.fromarray(rgb_array, mode=image_mode)
  return image


DEFAULT_GRAY_SCALE_FACTOR = {np.uint8: 100.0,
                             np.uint16: 1000.0,
                             np.int32: DEFAULT_RGB_SCALE_FACTOR}


def FloatArrayToGrayImage(float_array, scale_factor=None, image_dtype=np.uint8):
  """Convert a floating point array of values to an RGB image.

  Convert floating point values to a fixed point representation with
  the given bit depth.

  The precision of the depth image with default scale_factor is:
    uint8: 1cm, with a range of [0, 2.55m]
    uint16: 1mm, with a range of [0, 65.5m]
    int32: 1/256mm, with a range of [0, 8388m]

  Right now, PIL turns uint16 images into a very strange format and
  does not decode int32 images properly.  Only uint8 works correctly.

  Args:
    float_array: Input array of floating point depth values in meters.
    scale_factor: Scale value applied to all float values.
    image_dtype: Image datatype, which controls the bit depth of the grayscale
      image.

  Returns:
    Grayscale PIL Image object representing depth values.

  """
  # Ensure that we have a valid numeric type for the image.
  if image_dtype == np.uint16:
    image_mode = 'I;16'
  elif image_dtype == np.int32:
    image_mode = 'I'
  else:
    image_dtype = np.uint8
    image_mode = 'L'
  if scale_factor is None:
    scale_factor = DEFAULT_GRAY_SCALE_FACTOR[image_dtype]
  # Scale the floating point array.
  scaled_array = np.floor(float_array * scale_factor + 0.5)
  # Convert the array to integer type and clip to representable range.
  min_dtype = np.iinfo(image_dtype).min
  max_dtype = np.iinfo(image_dtype).max
  scaled_array = ClipFloatValues(scaled_array, min_dtype, max_dtype)

  image_array = scaled_array.astype(image_dtype)
  image = Image.fromarray(image_array, mode=image_mode)
  return image


def ImageToFloatArray(image, scale_factor=None):
  """Recovers the depth values from an image.

  Reverses the depth to image conversion performed by FloatArrayToRgbImage or
  FloatArrayToGrayImage.

  The image is treated as an array of fixed point depth values.  Each
  value is converted to float and scaled by the inverse of the factor
  that was used to generate the Image object from depth values.  If
  scale_factor is specified, it should be the same value that was
  specified in the original conversion.

  The result of this function should be equal to the original input
  within the precision of the conversion.

  Args:
    image: Depth image output of FloatArrayTo[Format]Image.
    scale_factor: Fixed point scale factor.

  Returns:
    A 2D floating point numpy array representing a depth image.

  """
  image_array = np.array(image)
  image_dtype = image_array.dtype
  image_shape = image_array.shape

  channels = image_shape[2] if len(image_shape) > 2 else 1
  assert 2 <= len(image_shape) <= 3
  if channels == 3:
    # RGB image needs to be converted to 24 bit integer.
    float_array = np.sum(image_array * [65536, 256, 1], axis=2)
    if scale_factor is None:
      scale_factor = DEFAULT_RGB_SCALE_FACTOR
  else:
    if scale_factor is None:
      scale_factor = DEFAULT_GRAY_SCALE_FACTOR[image_dtype.type]
    float_array = image_array.astype(np.float32)
  scaled_array = float_array / scale_factor
  return scaled_array
rosbag_utils/depth_image_encoding.py
baldeeb/rosbag_data_utils
from geometry_msgs.msg import TransformStamped 
from std_msgs.msg import Time
from tf2_ros import Buffer, TransformException
from rospy import Duration
from typing import List, Optional
from rosbag_utils.rosbag_as_dataset import RosbagReader
from rosbag_utils.ros_to_numpy_helpers import ros_transform_to_np_se3_matrix
import yaml

class Tf2Wrapper:
    def __init__(self, 
                 authority:str='tf2_wrapper', 
                 cache_time:Optional[Duration]=None
                 ):
        self._buffer = Buffer(cache_time, debug=False)
        self._authority = authority
        self._hooks = {}

    def register_getter_hook(self, hook, ref_frame:str=None, frame:str=None):
        if ref_frame is not None and frame is not None:
            self._hooks[(ref_frame, frame)] = hook
        elif ref_frame is not None:
            self._hooks[ref_frame] = hook
        elif frame is not None:
            self._hooks[frame] = hook
        else:
            raise ValueError('Either ref_frame or frame must be specified')

    def set(self, transform:TransformStamped, static:bool):
        if static:
            self._buffer.set_transform_static(transform, self._authority)
        else:
            self._buffer.set_transform(transform, self._authority)

    def _post_process(self, T, ref_frame, frame):
        if (ref_frame, frame) in self._hooks:
            return self._hooks[(ref_frame, frame)](T)
        elif ref_frame in self._hooks:
            return self._hooks[ref_frame](T)
        elif frame in self._hooks:
            return self._hooks[frame](T)
        else:
            return T

    def get(self, reference:str, frame_name:str, timestamp:Time):
        t = self._buffer.lookup_transform(reference, frame_name, timestamp)
        T = ros_transform_to_np_se3_matrix(t.transform)
        T = self._post_process(T, reference, frame_name)
        return T
    
    def __str__(self):
        return self._buffer.all_frames_as_yaml()
    
    def get_as_dict(self):
        return yaml.safe_load(self._buffer.all_frames_as_yaml())


def get_populated_tf2_wrapper(bag_path:str, 
                              static_topics:List[str]= ['/tf_static'], 
                              dynamic_topics:List[str]=['/tf'],
                              cache_time_sec:float=60*2,
                              slack_sec:float=0.1,
                              ):

    def add_topics(topics, is_static):
        for topic in topics:
            dataset = RosbagReader(bag_path, {topic: 't'}, 
                                   permissible_asynchronisity_sec=slack_sec)
            for d in dataset:
                for transfrom in d['t'].transforms:
                    wrapper.set(transfrom, is_static)
    
    wrapper = Tf2Wrapper(cache_time=Duration(cache_time_sec))
    add_topics(static_topics, True)
    add_topics(dynamic_topics, False)

    return wrapper

rosbag_utils/ros_tf2_wrapper.py
baldeeb/rosbag_data_utils
from sensor_msgs.msg import (
    Image, 
    CameraInfo,
    )
from geometry_msgs.msg import (
    Pose,
    Point, 
    Quaternion,
    PoseWithCovarianceStamped,
    Transform
    )
import numpy as np
from quaternion import as_rotation_matrix, quaternion, from_rotation_matrix


def ros_camera_info_to_np_intrinsic(info:CameraInfo):
    return np.array([
        [info.K[0], info.K[1], info.K[2]],
        [info.K[3], info.K[4], info.K[5]],
        [info.K[6], info.K[7], info.K[8]]
    ])

def ros_image_to_np(image:Image, depth_to_meters=1e-3):
    H, W = image.height, image.width
    if image.encoding == 'rgb8':
        rgb = np.frombuffer(image.data, dtype=np.byte)
        img = rgb.reshape(H, W, 3).astype(np.uint8)
    elif image.encoding == 'rgba8':
        rgb = np.frombuffer(image.data, dtype=np.byte)
        img = rgb.reshape(H, W, 4).astype(np.uint8)
    elif image.encoding == 'bgra8':
        rgb = np.frombuffer(image.data, dtype=np.byte)
        img = rgb.reshape(H, W, 4)[:, :, (2,1,0)].astype(np.uint8)
    elif image.encoding == '16UC1':
        d = np.frombuffer(image.data, dtype=np.uint16).reshape(H, W)
        img = d.astype(np.float32) * depth_to_meters
    elif image.encoding == 'bgra8':
        rgbd = np.frombuffer(image.data, dtype=np.byte)
        rgbd = rgbd.reshape(H, W, 4)[:, :, (2,1,0)].astype(np.uint8)
        img = rgbd[:,:,3].astype(np.uint16).astype(np.float32) * depth_to_meters
    else: 
        raise RuntimeError(f'Image to Numpy is not setup to handle {image.encoding}.')
    return img

def _ros_point_to_np(point:Point):
    return np.array([point.x, point.y, point.z])

def _ros_quaternion_to_np(quat:Quaternion):
    q = quaternion(quat.w, quat.x, quat.y, quat.z)
    return as_rotation_matrix(q)

def ros_pose_to_np_se3_matrix(pose:Pose):
    mat = np.eye(4)
    mat[:3,:3] = _ros_quaternion_to_np(pose.orientation)
    mat[:3,3] = _ros_point_to_np(pose.position)
    return mat

def ros_pose_to_7d_position_quaternion(pose:Pose):
    vec = np.empty(7)
    vec[:3] = _ros_point_to_np(pose.position)
    vec[3:] = np.array([pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w])
    return vec
    
def ros_pose_with_cov_stamped_to_np_se3_matrix(pose:PoseWithCovarianceStamped):
    return ros_pose_to_np_se3_matrix(pose.pose.pose)

def ros_transform_to_np_se3_matrix(transform:Transform):
    mat = np.eye(4)
    mat[:3,:3] = _ros_quaternion_to_np(transform.rotation)
    mat[:3,3] = _ros_point_to_np(transform.translation)
    return mat

def se3_mat_to_position_and_quaterion_vec(mat):
    vec = np.empty(7)
    vec[:3] = mat[:3,3]
    q = from_rotation_matrix(mat[:3,:3])
    vec[3:] = np.array([q.x, q.y, q.z, q.w]) 
    return vec

def ros_joint_states_to_numpy(msg):
    return {
        'names': msg.name,
        'positions': np.array(msg.position),
        'velocities': np.array(msg.velocity),
        'efforts': np.array(msg.effort),
    }


rosbag_utils/ros_to_numpy_helpers.py
baldeeb/rosbag_data_utils
import rospy
import rosbag
from rosbag_utils.ros_to_numpy_helpers import (
    ros_camera_info_to_np_intrinsic,
    ros_image_to_np,
    ros_pose_to_np_se3_matrix,
    ros_joint_states_to_numpy,
)
import logging

def get_topics_in_path(path:str):
    bag = rosbag.Bag(path)
    return list(bag.get_type_and_topic_info()[1])

class MissingTopicError(Exception):
    def __init__(self, message, topics):
        super().__init__(message)
        self.topics = topics

class RosbagReader:
    ''' 
    Takes a ros bag and returns one item at a time.
    Assumes that the bag topics are logged in a normally distributed fashion. 
    ''' 
    def __init__(self, 
                 path, 
                 topics_and_names=None, 
                 default_topic_process=True, 
                 anytime_topics=[],
                 permissible_asynchronisity_sec=None,
                 reference_topic=None):
        '''
        Args:
            path: path to the rosbag file
            topics_and_names: a dictionary of topic names to be used in the output
            default_topic_process: if True, the reader will use default postprocess functions
        '''
        self.bag = rosbag.Bag(path)

        if topics_and_names is not None:
            available_topics = self.get_topics()
            desired_topics = list(topics_and_names.keys())
            self._validate_requested_topics(desired_topics, available_topics)
            self._topics = desired_topics
        else:
            self._topics = self.get_topics()

        # Select reference topic
        if reference_topic is not None:
            self._len, self._time_slots = self._get_reference_topic_info(reference_topic)
        else:
            reference_topic, self._len, self._time_slots = self._get_topic_with_least_messages()
        
        slack_sec = self._resolve_timing_details(permissible_asynchronisity_sec)
        self._resolve_naming_details(topics_and_names)

        self._anytime_topics = anytime_topics

        self._postprocess_type = {}
        self._postprocess_topic = {}
        self._postprocess_output_name = {}

        if default_topic_process:
            self._postprocess_type = {
                'sensor_msgs/Image': ros_image_to_np,
                'geometry_msgs/PoseWithCovarianceStamped': ros_pose_to_np_se3_matrix,
                'sensor_msgs/CameraInfo': ros_camera_info_to_np_intrinsic,
                'sensor_msgs/JointState': ros_joint_states_to_numpy,
            }

        logging.info(f'Initialized BagReader - ref topic: {reference_topic} - time slack {slack_sec}.')


    def _resolve_timing_details(self, slack_sec=None):
        '''Determines time window around said topic's messages to search for other topics'''
        self._start, self._end = self.bag.get_start_time(), self.bag.get_end_time()
        if slack_sec is not None: 
            self._time_eps = rospy.Time.from_sec(slack_sec)
        else: 
            slack_sec = (self._end - self._start) / self.__len__()
            self._time_eps = rospy.Time.from_sec(slack_sec)
        return slack_sec

    def _resolve_naming_details(self, topics_and_names):
        '''Figures out what names to use for what topics in output.'''
        if topics_and_names is not None:
            self._name_of_topic = topics_and_names
        else: 
            self._name_of_topic = {}
            for t in self._topics:
                if t not in self._name_of_topic:
                    self._name_of_topic[t] = t

    def _validate_requested_topics(self, desired, available):
        if not all([t in available for t in desired]):
            not_found = [t for t in desired if t not in available]
            raise RuntimeError(f"The topics {not_found} are not available in the bag.\n" +
                               f'available topics: {available}')    
    
    def __iter__(self):
        for i in range(self.__len__()):
            yield self[i]


    def clear_postprocess(self):
        self._postprocess_type = {}
    def unregister_type_postprocess(self, msg_type:str):
        del self._postprocess_type[msg_type]
    def register_type_postprocess(self, msg_type:str, func:callable):
        self._postprocess_type[msg_type] = func


    def register_topic_postprocess(self, topic:str, func:callable):
        self._postprocess_topic[topic] = func
    def unregister_topic_postprocess(self, topic:str):
        del self._postprocess_topic[topic]


    def register_output_postprocess(self, msg_name:str, func:callable):
        self._postprocess_output_name[msg_name] = func
    def unregister_output_postprocess(self, msg_name:str):
        del self._postprocess_output_name[msg_name]


    def __len__(self): return self._len


    def _get_topic_with_least_messages(self):
        topic = self._topics[0]
        count = self.bag.get_message_count(topic)
        for t in self._topics[1:]:
            c = self.bag.get_message_count(t)
            if c < count: count, topic = c, t
        time_sec = [m[2] for m in self.bag.read_messages(topics=[topic])]
        return topic, count, time_sec  

    def _get_reference_topic_info(self, ref):
        count = self.bag.get_message_count(ref)
        time_sec = [m[2] for m in self.bag.read_messages(topics=[ref])]
        return count, time_sec  


    def get_topics(self):
        return list(self.bag.get_type_and_topic_info()[1])


    def _msgs_at(self, time:rospy.Time, rigourous=True, anytime=False, topics=None):
        '''Ensures that messages are returned by proximity from intended time.'''
        if topics is None: topics = self._topics
        if anytime: 
            rigourous = True
            s = e = None
        else:
            t_time, elapsed = time.to_time(), self._time_eps.to_time()
            s = rospy.Time.from_sec(t_time - elapsed)
            e = rospy.Time.from_sec(t_time + elapsed)

        msgs = self.bag.read_messages(topics = topics, 
                                      start_time = s, 
                                      end_time = e)
        if rigourous:
            msgs = list(msgs)
            dt = [abs(m[2] - time) for m in msgs]
            idxs = sorted(range(len(dt)), key=lambda k: dt[k])
            for i in idxs:
                yield msgs[i]
        else:
            for m in msgs:
                yield m


    def _process(self, topic, name, msg):
        data = {}
        if topic in self._postprocess_topic:           # if topic has a postprocess function
            data[name] = self._postprocess_topic[topic](msg)
        elif msg._type in self._postprocess_type:    # if message type has a postprocess function
            data[name] = self._postprocess_type[msg._type](msg)
        else:                                           # no postprocess function
            data[name] = msg

        if name in self._postprocess_output_name:   # if generated data has a postprocess function
            data[name] = self._postprocess_output_name[name](data[name])
        return data


    def __getitem__(self, idx):
        time = self._time_slots[idx]
        data = {}
        for msg in self._msgs_at(time):
            msg_name = self._name_of_topic[msg[0]]
            if msg_name in data: continue                           # Continue, if topic already obtained
            data.update(self._process(msg[0], msg_name, msg[1]))    # Add msg to data
            if len(data) == len(self._topics): break                # All topics collected
        
        # If some topics are set to time agnostic, try to get them from anytime
        if len(data) != len(self._topics):
            if len(self._anytime_topics) > 0:
                missing = []
                for topic, name in self._name_of_topic.items():
                    if name not in data:
                        if name in self._anytime_topics:
                            msg = next(iter(self._msgs_at(time, anytime=True, 
                                                        topics=[topic])))[1]
                            data.update(self._process(topic, name, msg))
                        else:
                            missing.append(name)
                if len(missing) > 0:
                    raise MissingTopicError(f'Not all topics were found at time {time}.', 
                                            topics=missing)
            else:
                raise MissingTopicError(f'Not all topics were found at time {time}.', 
                                        topics=[k for k in data if k not in self._topics])
        return data

    def __del__(self):
        self.bag.close()

rosbag_utils/rosbag_as_dataset.py
baldeeb/rosbag_data_utils
import os
import sys
import logging
import pathlib as pl
from copy import deepcopy
from typing import List, Dict, Optional

import cv2
import pickle
import numpy as np
from PIL import Image

from rosbag_utils.rosbag_as_dataset import RosbagReader, MissingTopicError
from rosbag_utils.ros_to_numpy_helpers import (ros_pose_to_np_se3_matrix,
                                                     ros_image_to_np,
                                                     se3_mat_to_position_and_quaterion_vec)
from rosbag_utils.ros_tf2_wrapper import (get_populated_tf2_wrapper,
                                                Tf2Wrapper)

from rlbench.backend.observation import Observation
from rlbench.demo import Demo

from rosbag_utils.depth_image_encoding import FloatArrayToRgbImage, DEFAULT_RGB_SCALE_FACTOR
from rospy import Time, Duration
from tf2_ros import tf2

ALL_CAMERAS = [
    'left_shoulder',
    'right_shoulder',
    'front',
    'overhead',
    'wrist'
]


def check_and_mkdirs(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path, exist_ok=True)


def save_img(img, path):
    if isinstance(img, np.ndarray):
        cv2.imwrite(path, img[:, :, :3])
    elif isinstance(img, Image.Image):
        img.save(path)


def extract_list_from_rosbag_dataset(dataset:RosbagReader) -> List[Dict]:
    '''This is where the assumptions about the dataset lie.
    We are assuming that there is a series of datapoints where all topics are present.
    This function finds the start and iterates until it hits the first missing topic.
    
    Args:
        dataset: RosbagReader object
    Returns: 
        List[Dict] where each dict is a datapoint
    '''

    starting_idx = 0
    datalist = []
    for i in range(starting_idx, len(dataset)):
        try:
            datalist.append(dataset[i])
        except MissingTopicError as e:
            logging.info(f'skipped episode {i} - missing: {e.topics}.')
    return datalist


def get_dummy_observation(frame, misc):

    # pose format -> x, y, z, qx, qy, qz, qw
    dummy_pose = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0 ,1.0])

    obs = Observation(
        left_shoulder_rgb=None,
        left_shoulder_depth=None,
        left_shoulder_point_cloud=None,
        right_shoulder_rgb=None,
        right_shoulder_depth=None,
        right_shoulder_point_cloud=None,
        overhead_rgb=None,
        overhead_depth=None,
        overhead_point_cloud=None,
        wrist_rgb=None,
        wrist_depth=None,
        wrist_point_cloud=None,
        front_rgb=None,
        front_depth=None,
        front_point_cloud=None,
        left_shoulder_mask=None,
        right_shoulder_mask=None,
        overhead_mask=None,
        wrist_mask=None,
        front_mask=None,
        joint_velocities            = dummy_pose,
        joint_positions             = dummy_pose,
        joint_forces                = dummy_pose,
        gripper_open                = 0.9,
        gripper_pose                = dummy_pose,
        gripper_matrix              = dummy_pose,
        gripper_touch_forces        = None,
        gripper_joint_positions     = np.ones(2),
        task_low_dim_state          =None,
        ignore_collisions           =True, # TODO: fix
        misc=misc,
    )

    return obs


def save_keypoint(episodes:List[Dict], 
                  save_path:str, 
                  description:str,
                  episode_idx:int=0, 
                  depth_scale=DEFAULT_RGB_SCALE_FACTOR,
                  cameras_used=ALL_CAMERAS,
                  episode_to_rlbench_obs=get_dummy_observation,
                  near_far_planes = [0, 1]
                  ):
    '''TODO: add docstring'''

     # make directories
    variation_idx = 0

    episode_path = os.path.join(save_path, 'all_variations', 'episodes', f"episode{episode_idx}")
    check_and_mkdirs(episode_path)

    base_misc = dict()

    for cam in cameras_used:
        for info in ['rgb', 'depth', 'camera_intrinsics',]:
            assert any([f'{cam}_{info}' in k for k in episodes[0].keys()]),\
                   f"Camera {cam}_{info} not found in frames!"

        image_dir = pl.Path(episode_path)/f'{cam}_rgb'
        image_dir.mkdir(parents=True, exist_ok=True)

        image_dir = pl.Path(episode_path)/f'{cam}_depth'
        image_dir.mkdir(parents=True, exist_ok=True)

        K_key = f'{cam}_camera_intrinsics'
        base_misc[K_key] = np.array(episodes[0][K_key]).reshape(3,3)
        base_misc[cam + '_camera_near'] = near_far_planes[0]                           # TODO: scrutinize these values
        base_misc[cam + '_camera_far'] = near_far_planes[1]                            # TODO: scrutinize these values

    base_misc['keypoint_idxs'] = [sum([1 for e in episodes if 'keypoint' in e]) - 1]

    observations = []
    image_keys = [f'{cam}_rgb' for cam in cameras_used]
    depth_keys = [f'{cam}_depth' for cam in cameras_used]
    for eps_idx, eps in enumerate(episodes):
        misc = deepcopy(base_misc)  
        found_image = False      
        for key, val in eps.items():
            image_dir = str(pl.Path(episode_path) / key / f'{eps_idx}.png')
            if key in image_keys: 
                save_img(cv2.cvtColor(val['image'], cv2.COLOR_RGB2BGR),
                         image_dir)
                misc[f'{key[:-4]}_camera_extrinsics'] = val['extrinsics']
                found_image = True
            elif key in depth_keys: 
                depth_img = (val['image'] - near_far_planes[0]) / (near_far_planes[1] - near_far_planes[0])
                v = FloatArrayToRgbImage(depth_img, scale_factor=depth_scale)
                save_img(v, image_dir)
                found_image = True
        assert found_image, f"Could not find image in episode {eps_idx}."
        if 'keypoint' in eps:
            misc['keypoint'] = eps['keypoint']
        observations.append(episode_to_rlbench_obs(eps, misc))

    demo = Demo(observations, random_seed=0)
    demo.variation_number = variation_idx

    low_dim_obs_path = os.path.join(episode_path, 'low_dim_obs.pkl')
    with open(low_dim_obs_path, 'wb') as f:
        pickle.dump(demo, f)

    variation_number_path = os.path.join(episode_path, 'variation_number.pkl')
    with open(variation_number_path, 'wb') as f:
        pickle.dump(variation_idx, f)

    descriptions_path = os.path.join(episode_path, 'variation_descriptions.pkl')
    with open(descriptions_path, 'wb') as f:
        pickle.dump([description], f)

    logging.info(f"Saved {len(episodes)} frames to {save_path}")


def _setup_rosbag_dataset(data_dir, 
                          topics_and_names:dict, 
                          time_slack:Optional[float]=None, 
                          reference_topic:str=None
                          ):
    dataset = RosbagReader(data_dir, topics_and_names, 
                           permissible_asynchronisity_sec=time_slack,
                           reference_topic=reference_topic)
    def task_status_postprocessing(msg):
        return {'name': msg.task_name,
                'pose': ros_pose_to_np_se3_matrix(msg.target_pose.pose),
                'reference_frame': msg.target_pose.header.frame_id,
                'time': msg.target_pose.header.stamp,
                'status': msg.status,}
    def image_and_time(msg):
        return {'image': ros_image_to_np(msg),
                'time': msg.header.stamp}
    # instead of just getting image get its timestamp as well
    dataset.unregister_type_postprocess('sensor_msgs/Image')
    dataset.register_type_postprocess('sensor_msgs/Image', image_and_time)
    # package task status into a keypoint dict
    dataset.register_topic_postprocess('/spot/task_state', task_status_postprocessing)
    logging.info(f'Dataset initial length: {len(dataset)}')
    
    return dataset

def add_tfs_to_episodes(tf_wrapper:Tf2Wrapper, data:List[Dict], 
                    names_and_frames:Dict, ref_frame:str,
                    try_handle_extrapolation:bool=True,
                    offset_t:Duration=Duration(0,0),
                    ):
    '''Adds the extrinsics to the data. Optionally ignores missing data.
    Args:
        tf_wrapper: Tf2Wrapper object containing buffer of all tfs.
        data: List of dicts containing the data into which the extrinsics will be added.
        names_and_frames: dictionary of the data keys (or names) that are to be augmented
            with the extrinsics (or frames).
        ref_frame: The frame of reference for all the extrinsics.
    '''
    skipped = []
    for data_i, d in enumerate(data):
        try:
            for name, frame in names_and_frames.items():
                if name not in d.keys(): 
                    d_time = [val['time'] for val in d.values() 
                              if isinstance(val, dict) and 'time' in val][0]
                    time = Time(secs=d_time.secs + offset_t.secs, 
                                nsecs=d_time.nsecs + offset_t.nsecs)
                    d[name] = tf_wrapper.get(ref_frame, frame, time)
                else:
                    d_time = d[name]['time']
                    time = Time(secs=d_time.secs + offset_t.secs, 
                                nsecs=d_time.nsecs + offset_t.nsecs)
                    d[name]['extrinsics'] = tf_wrapper.get(ref_frame, frame, time)
        except tf2.ExtrapolationException as e:
            if try_handle_extrapolation:
                if data_i == 0 and offset_t == Duration(0,0):
                    tf_dict = tf_wrapper.get_as_dict()
                    earliest = max([tf_dict[f]['oldest_transform'] 
                                    for f in [frame, ref_frame]])
                    return add_tfs_to_episodes(tf_wrapper, data, names_and_frames, ref_frame, 
                                               try_handle_extrapolation, 
                                               offset_t=Time.from_sec(earliest) - time + Duration(0.001,0))
                elif data_i > int(0.8 * len(data)):
                    skipped.append(data_i)
                    logging.warning(f"Skipped {data_i} frame due to mismatched timestamps between tfs and data.")
                else:
                    logging.error(f'Failed to handle extrapolation exception at data index: {data_i}.')
                    raise e
            else: 
                raise e
        
    if len(skipped) > 0: logging.warning(f"Skipped frames {skipped}.")
    
    return skipped


def keypoints_from_frame(data:List[Dict], tf_wrapper:Tf2Wrapper, ref_frame:str):
    '''Transforms all keypoint poses to the reference frame.'''
    for d in data:
        k = d['keypoint']
        #### NOTE: TEMPORARY FIX!!!
        ## This should be the data time.
        time = Time(secs=d['left_shoulder_rgb']['time'].secs, nsecs=d['left_shoulder_rgb']['time'].nsecs)
        # time = Time(secs=k['time'].secs, nsecs=k['time'].nsecs)
        T = tf_wrapper.get(ref_frame, k['reference_frame'], time)
        k['pose'] = se3_mat_to_position_and_quaterion_vec(T @ k['pose'])
        k['reference_frame'] = ref_frame






###################################################################
###################################################################
###################################################################
###################################################################
# TODO: remove below functions
###################################################################
###################################################################


def rosbag_to_rlbench(data_dir:pl.Path, 
                      out_dir:pl.Path, 
                      topics_and_names:Dict, 
                      names_and_frames:Dict,
                      reference_frame:str, 
                      task_description:List[str]=["no lang goal defined"],
                      episode_idx:int = 0,
                      depth_scale = ( ( 2**22 ) - 1.0 ),
                      cameras_used = ALL_CAMERAS,
                      time_slack:Optional[float]=None,
                      reference_topic:Optional[str]=None,
                      reference_frame_T:np.ndarray=np.eye(4),
                      ) -> None:
    '''
    Args:
        names_and_frames: dictionary of all the datapoints' names 
            and tf frame names. This list indicates which datapoints
            we care to associate with a pose. Mainly intended for
            camera data extrinsics.
        reference: The frame of reference for all the extrinsics.
    '''
    logging.info('.'*3 + f'processing: {data_dir}' + '.'*3)
    tf_data = get_populated_tf2_wrapper(data_dir)
    tf_data.register_getter_hook(lambda x: reference_frame_T @ x, reference_frame)
    dataset = _setup_rosbag_dataset(data_dir, topics_and_names,
                                    time_slack=time_slack,
                                    reference_topic=reference_topic)

    logging.info(f'depth scale: {depth_scale}')
    episodes = extract_list_from_rosbag_dataset(dataset)
    add_tfs_to_episodes(tf_data, 
                        episodes, 
                        names_and_frames, 
                        reference_frame)
    keypoints_from_frame(episodes, 
                         tf_data, 
                         reference_frame)
    save_keypoint(episodes, 
                  out_dir, 
                  task_description, 
                  episode_idx, 
                  depth_scale=depth_scale,
                  cameras_used=cameras_used)
    logging.info('.'*5 + 'Done!' + '.'*5)


def rosbags_to_rlbench_episodes(data_dir:pl.Path, 
                                out_dir:pl.Path, 
                                topics_and_names:Dict,
                                names_and_frames, 
                                reference_frame,
                                task_description:List[str]=["no lang goal defined"],
                                depth_scale = ( ( 2**22 ) - 1.0 ),
                                cameras_used:List[str]=ALL_CAMERAS,
                                time_slack:Optional[float]=None,
                                reference_topic:Optional[str]=None,
                                reference_frame_T:np.ndarray=np.eye(4),
                                ) -> None:

    for i, data_path in enumerate(data_dir.glob('*.bag')):
        rosbag_to_rlbench(data_path, 
                          out_dir/data_dir.name, 
                          topics_and_names, 
                          names_and_frames, 
                          reference_frame,
                          task_description,
                          i,
                          depth_scale=depth_scale,
                          cameras_used=cameras_used,
                          time_slack=time_slack,
                          reference_topic=reference_topic,
                          reference_frame_T=reference_frame_T,
                          )
    logging.getLogger(__name__).info('Done!')

rosbag_utils/rosbag_to_rlbench.py
baldeeb/rosbag_data_utils
# Setup file for installing the package

from distutils.core import setup

setup(name='rosbag_utils',
        version='0.1',
        description='Utilities for translating rosbags to numpy arrays',
        author='Bahaa Aldeeb',
        author_email='aldeeb.bahaa@gmail.com',
        packages=['rosbag_utils'],
        requires=['rlbench', 'numpy', 'rospy', 'rosbag', 'tf2_ros', 
                  'geometry_msgs', 'sensor_msgs', 'std_msgs', 
                  'rospkg', 'pycryptodomex', 'gnupg', 'pyyaml',]
)


setup.py
baldeeb/rosbag_data_utils