generate_traversability_data

#!/usr/bin/env python
"""Generate self-supervised data from ROS bag files with robot trajectories.
Assign traversability values to points travelled by a robot.
Traversability is calculated using vertical acceleration values measured by IMU sensor.

Paramaters:
- Simplified robot model from primitives (bounding boxes, spheres).
- Lookahead time / distance to mark traversed-through points.

Multi-pass processing:
1. Load transforms from bag files into a buffer. There have to be a transform
   from the fixed frame to the robot.
2. Process bags again and annotate points within given horizon contained in
   the model primitives.
"""
from __future__ import absolute_import, division, print_function
from argparse import ArgumentParser
from glob import glob
from matplotlib import cm
import numpy as np
from numpy.lib.recfunctions import merge_arrays, unstructured_to_structured
import os
from ros_numpy import msgify, numpify
from rosbag import Bag, ROSBagException, Compression
import rospy
from sensor_msgs.msg import PointCloud2, Imu
from tf2_ros import BufferCore, TransformException
from traversability_estimation.geometry import affine, Bodies, Body, Box, inverse, Sphere  # needed for eval
from traversability_estimation.segmentation import compute_rigid_support, fit_planes, position, valid_point_mask
from traversability_estimation.utils import show_cloud, slots, normalize
from tqdm import tqdm

DEPTH_FIELD_NAME = 'depth'
LABEL_FIELD_NAME = 'traversability'
LABEL_IGNORE = 255
G_ACC = 9.80665
# LABEL_OBSTACLE = round(G_ACC)
LABEL_OBSTACLE = G_ACC


def str2bool(v):
    return v.lower() in ('1', 'yes', 'true', 't', 'y')


def arg_parser():
    parser = ArgumentParser(epilog="""Path format uses following placeholders:
    {dir} - parent directory of the first bag file,
    {name} - name without extension of the first bag file,
    {topic} - name of the topic as read from bag file,
    {secs}, {nsecs} - timestamp from the header (if available).
    """)
    parser.add_argument('--cloud-topics', type=str, nargs='+')
    parser.add_argument('--imu-topics', type=str, nargs='+')
    parser.add_argument('--imu-rate', type=int, default=100)
    parser.add_argument('--fixed-frame', type=str, default='map')
    parser.add_argument('--robot-frame', type=str, default='base_link')
    parser.add_argument('--exclude-times', type=str, default=None)
    parser.add_argument('--discard-empty', type=bool, default=True)
    parser.add_argument('--input-step', type=int, default=1)
    parser.add_argument('--input-start', type=float, default=-float('inf'), help='Start time in seconds.')
    parser.add_argument('--input-end', type=float, default=float('inf'), help='End time in seconds.')
    parser.add_argument('--discard-model', type=str, default=None,
                        help='Model at current position discarding points.')
    parser.add_argument('--robot-model', type=str, default=None,
                        help='Model to move along robot path marking points traversable.')
    parser.add_argument('--obstacles-model', type=str, default=None,
                        help='Model to move along robot path marking points rigid.')
    parser.add_argument('--obstacles-primitives', type=str2bool, default=False,
                        help='Fit primitives and mark the inliers as rigid.')
    parser.add_argument('--distance-horizon', '-d', type=float, default=10.0)
    parser.add_argument('--time-horizon', '-t', type=float, nargs=2, default=[0.0, 10.0])
    parser.add_argument('--time-step', '-s', type=float, default=0.5)
    parser.add_argument('--output-path', '-o', type=str, default='{dir}/{name}_trav/{topic}/{secs}_{nsecs:09d}.npz')
    parser.add_argument('--output-bag-path', '-O', type=str, default='{dir}/{name}_segmented.bag')
    parser.add_argument('--output-topic', type=str, default='{topic}_segmented')
    parser.add_argument('--output-period', type=float, default=None)
    parser.add_argument('--visualize', type=str2bool, default=False)
    parser.add_argument('--save-data', type=str2bool, default=True)
    parser.add_argument('bag_paths', type=str, nargs='+')
    parser.add_argument('--z-support', type=str, default=None, help='Z support configuration or empty.')
    return parser


def get_topic_types(bag):
    return {k: v.msg_type for k, v in bag.get_type_and_topic_info().topics.items()}


def load_buffer(bag_paths, imu_topics=None, imu_rate=100):
    tf_topics = ['/tf', '/tf_static']
    if imu_topics is None:
        imu_topics = ['/imu/data']

    # tf_buffer = BufferCore(cache_time=rospy.Duration(2**31 - 1))
    # tf_buffer = BufferCore(cache_time=rospy.Duration(24 * 60 * 60))
    tf_buffer = BufferCore(rospy.Duration(24 * 60 * 60))
    
    class IMUBuffer(object):
        def __init__(self, sensor_rate):
            self.stamps = []
            self.msgs = []
            self.sensor_rate = sensor_rate

        def pop(self, msg):
            assert isinstance(msg, Imu)

            self.stamps.append(msg.header.stamp.to_sec())
            self.msgs.append(msg)

        def __len__(self):
            return len(self.stamps)

        @staticmethod
        def estimate_traversability(msgs, G=G_ACC):
            acc_z = np.asarray([msg.linear_acceleration.z for msg in msgs])
            trav = np.mean(np.abs(acc_z + G))
            return trav

        def get_traversability_data(self, t_sec):
            """
            returns traversability data by averaging Z-axis IMU data
            closest to `t_sec` stamp
            """
            t_ind = np.argmin(np.abs(np.asarray(self.stamps) - t_sec))
            # print(np.min(np.abs(np.asarray(self.stamps) - t_sec)))

            window_size = int(self.sensor_rate * 0.5)  # average IMU values received during 0.5 sec interval
            left_ind = max(0, t_ind - window_size // 2)
            right_ind = min(t_ind + window_size // 2, len(self))
            traversability = self.estimate_traversability(self.msgs[left_ind:right_ind])

            return traversability

    imu_buffer = IMUBuffer(sensor_rate=imu_rate)
    
    for path in bag_paths:
        try:
            with Bag(path, 'r') as bag:
                for topic, msg, stamp in tqdm(bag.read_messages(topics=tf_topics + imu_topics),
                                              desc='%s: reading transforms' % path.split('/')[-1],
                                              total=bag.get_message_count(topic_filters=tf_topics)):
                    if topic == '/tf':
                        for tf in msg.transforms:
                            tf_buffer.set_transform(tf, 'bag')
                    elif topic == '/tf_static':
                        for tf in msg.transforms:
                            tf_buffer.set_transform_static(tf, 'bag')
                    elif topic in imu_topics:
                        msg = Imu(*slots(msg))
                        imu_buffer.pop(msg)
                        
        except ROSBagException as ex:
            print('Could not read %s: %s' % (path, ex))

    return tf_buffer, imu_buffer

def fit_primitives(x_valid):
    # Fit obstacle models (trunks or branches as cylinders, ground or walls as plane).
    models = []
    models += fit_planes(x_valid.T, 0.025, normal_z_limits=(0.0, 0.2), max_iterations=200, min_support=25,
                         max_models=10, cluster_eps=0.25, cluster_k=10, verbose=0, visualize=False)
    return models


def segment_cloud(robot_model, arr, input_to_robot_tfs, trav_datas,
                  discard_tf=None, discard_model=None,
                  obstacles_model=None, obstacles_primitives=True,
                  input_to_fixed_tf=None, z_support=None,
                  visualize=False):

    assert robot_model is None or isinstance(robot_model, Body)
    assert isinstance(arr, np.ndarray)
    assert discard_model is None or isinstance(discard_model, Body)

    # Use only valid points for all operations.
    # Assign results to valid points in the end.
    valid = valid_point_mask(arr, discard_tf=discard_tf, discard_model=discard_model)
    arr_valid = arr[valid]
    valid_ind = np.flatnonzero(valid)
    x = position(arr_valid).reshape((-1, 3))
    x = x.T

    # Initialize all labels as unknown.
    labels = np.full(arr.shape, LABEL_IGNORE, dtype=float).ravel()

    # Mark valid points which are contained by future model poses as empty.
    traversability = np.full((x.shape[1],), np.nan, dtype=float)
    assert len(input_to_robot_tfs) == len(trav_datas)
    for input_to_robot_tf, trav_data in zip(input_to_robot_tfs, trav_datas):
        assert isinstance(input_to_robot_tf, np.ndarray)

        x_robot = affine(input_to_robot_tf, x)

        if robot_model:
            robot_points = robot_model.contains_point(x_robot)
            # average currently received traversability values with previous ones in the map
            traversability[robot_points] = np.nanmean([traversability[robot_points],
                                                       trav_data * np.ones_like(traversability[robot_points])], axis=0)

    mask_traversed = ~np.isnan(traversability)
    n_traversed_pts = mask_traversed.sum()
    print('%.3g = %i / %i traversed points' % (mask_traversed.mean(), n_traversed_pts, mask_traversed.size))

    # Fit geometric primitives to points and mark their inliers as obstacles.
    obstacles = np.zeros((x.shape[1],), dtype=bool)
    if obstacles_primitives:
        primitives = fit_primitives(x)
        for primitive, indices in primitives:
            obstacles[indices] = True
        print('%.3g = %i / %i obstacle points' % (obstacles.mean(), obstacles.sum(), obstacles.size))

    if z_support:
        support, tmp_rigid = compute_rigid_support(arr_valid, transform=input_to_fixed_tf, **z_support)
        obstacles[tmp_rigid] = True
        print('Obstacles points: %i' % obstacles.sum())

    if n_traversed_pts > 0:
        print('Traversability range for the current scan: %.3f .. %.3f (mean: %.3f)' %
              (np.min(traversability[mask_traversed]), np.max(traversability[mask_traversed]),
               np.mean(traversability[mask_traversed])))
        # normalize traversability labels as float points in a range between 0 and 1
        # traversability[mask_traversed] = normalize(traversability[mask_traversed])
        labels[valid_ind[mask_traversed]] = traversability[mask_traversed]
        labels[valid_ind[obstacles]] = LABEL_OBSTACLE

    if visualize and (n_traversed_pts > 100):
        show_cloud(x.T, labels[valid_ind],
                   min=labels[valid_ind[mask_traversed]].min(), max=LABEL_OBSTACLE + 1,
                   colormap=cm.jet)

    labels = labels.reshape(arr.shape)
    return labels, n_traversed_pts


def generate_data(bag_paths=None, cloud_topics=None, imu_topics=None, imu_rate=100, fixed_frame=None, robot_frame=None,
                  exclude_times=None, input_step=1, input_start=0.0, input_end=float('inf'),
                  discard_model=None, robot_model=None,
                  obstacles_model=None, obstacles_primitives=False, z_support=None,
                  discard_empty=True, distance_horizon=None, time_horizon=None, time_step=None,
                  output_path=None, output_bag_path=None, output_topic=None, output_period=None,
                  visualize=False, save_data=True):
    assert bag_paths, bag_paths
    assert not exclude_times or all(len(t) == 2 for t in exclude_times), exclude_times
    assert len(time_horizon) == 2, time_horizon
    # TODO: Always look up t=0.
    assert 0.0 <= time_horizon[0] < time_horizon[1], time_horizon
    print('Time horizon:', time_horizon)

    dir = os.path.dirname(bag_paths[0])
    name, _ = os.path.splitext(os.path.basename(bag_paths[0]))

    n = [int(np.floor(h / time_step)) for h in time_horizon]

    last_out = {}

    if output_bag_path:
        output_bag_path = output_bag_path.format(dir=dir, name=name)
        if output_bag_path in bag_paths:
            print('Output %s removed from input bag files.' % output_bag_path)
            del bag_paths[bag_paths.index(output_bag_path)]
        os.makedirs(os.path.dirname(output_bag_path), exist_ok=True)
        output_bag = Bag(output_bag_path, 'w', compression=Compression.LZ4)
    else:
        output_bag = None

    tf_buffer, imu_buffer = load_buffer(bag_paths, imu_topics=imu_topics, imu_rate=imu_rate)

    for bag_path in bag_paths:

        with Bag(bag_path, 'r') as bag:

            topic_types = get_topic_types(bag)
            i = -1
            for topic, msg, stamp in tqdm(bag.read_messages(topics=cloud_topics),
                                          desc='%s: generating data' % bag_path.split('/')[-1],
                                          total=bag.get_message_count(topic_filters=cloud_topics)):
                i += 1
                if i % input_step != 0:
                    continue
                if stamp.to_sec() < input_start or stamp.to_sec() > input_end:
                    print('Skipping %s at %.3f s (outside input interval).' % (topic, stamp.to_sec()))
                    continue
                if exclude_times and any(t[0] <= stamp.to_sec() <= t[1] for t in exclude_times):
                    print('Skipping %s at %.3f s (excluded).' % (topic, stamp.to_sec()))
                    continue

                fmt_kwargs = {'dir': dir, 'name': name, 'topic': topic}
                if hasattr(msg, 'header'):
                    secs, nsecs = msg.header.stamp.secs, msg.header.stamp.nsecs
                    start = msg.header.stamp.to_sec()
                else:
                    secs, nsecs = stamp.secs, stamp.nsecs
                    start = stamp.to_sec()
                fmt_kwargs['secs'], fmt_kwargs['nsecs'] = secs, nsecs

                if output_period and topic in last_out and start - last_out[topic] < output_period:
                    continue

                # Find transform from input cloud to fixed frame.
                try:
                    input_to_fixed = tf_buffer.lookup_transform_core(fixed_frame, msg.header.frame_id, msg.header.stamp)
                except TransformException as ex:
                    print('Could not transform from %s to %s at %.3f s.' %
                          (msg.header.frame_id, fixed_frame, msg.header.stamp.to_sec()))
                    continue
                input_to_fixed_tf = numpify(input_to_fixed.transform)

                # Find transforms from input cloud to robot positions within the horizon.
                input_to_robot_tfs = []
                trav_data_horizon = []
                for t in np.linspace(start - n[0] * time_step, start + n[1] * time_step, sum(n) + 1):
                    try:
                        tf = tf_buffer.lookup_transform_full_core(robot_frame, rospy.Time.from_seconds(t),
                                                                  msg.header.frame_id, msg.header.stamp,
                                                                  fixed_frame)
                    except TransformException as ex:
                        # print('Could not transform from %s to %s at %.3f s.' % (msg.header.frame_id, robot_frame, t))
                        continue
                    tf = numpify(tf.transform)
                    if input_to_robot_tfs:
                        # Check distance horizon.
                        diff = np.matmul(input_to_robot_tfs[n[0]], inverse(tf))
                        distance = np.linalg.norm(diff[:-1, -1])
                        if distance > distance_horizon:
                            print('Distance horizon reached, %.3f m > %.3f m.' % (distance, distance_horizon))
                            break
                    input_to_robot_tfs.append(tf)

                    # compute traversability from IMU data closest to point clouds' stamps
                    trav_data = imu_buffer.get_traversability_data(t)
                    trav_data_horizon.append(trav_data)

                if not input_to_robot_tfs:
                    continue

                if topic_types[topic] == 'sensor_msgs/PointCloud2':
                    msg = PointCloud2(*slots(msg))
                    input_struct = numpify(msg)
                    # print('Input struct:', input_struct.shape)

                    # H x W unstructured depth image.
                    depth = np.linalg.norm(position(input_struct), axis=-1)
                    # H x W structured depth cloud.
                    depth_struct = depth.reshape((input_struct.size, -1))
                    depth_struct = unstructured_to_structured(depth_struct, names=[DEPTH_FIELD_NAME])
                    # depth_struct = depth_struct.reshape(input_struct.shape)
                    # assert depth_struct.shape == input_struct.shape, (depth_struct.shape, input_struct.shape)

                    # H x W unstructured label image.
                    label, n_traversed_pts = segment_cloud(robot_model, input_struct, input_to_robot_tfs, trav_data_horizon,
                                                           discard_tf=input_to_robot_tfs[n[0]], discard_model=discard_model,
                                                           obstacles_model=obstacles_model, obstacles_primitives=obstacles_primitives,
                                                           input_to_fixed_tf=input_to_fixed_tf, z_support=z_support,
                                                           visualize=visualize)
                    # H x W structured label cloud.
                    label_struct = label.reshape((input_struct.size, -1))
                    label_struct = unstructured_to_structured(label_struct, names=[LABEL_FIELD_NAME])

                    cloud_struct = merge_arrays([input_struct, depth_struct, label_struct],
                                                flatten=True)
                    cloud_struct = cloud_struct.reshape(input_struct.shape)
                    # print('Cloud struct:', cloud_struct.shape)
                    assert cloud_struct.shape == input_struct.shape

                    n_valid_labels = (label != LABEL_IGNORE).sum()
                    if discard_empty and (n_valid_labels < 1000):
                        print('Discarding cloud with not enough valid labels.')
                        continue
                    elif n_traversed_pts < 100:
                        print('Discarding cloud with not enough traversed points.')
                        continue

                    last_out[topic] = start

                    if save_data:
                        print('Storing cloud with %i valid labels and %i traversed points.' %
                              (n_valid_labels, n_traversed_pts))

                        # write np arrays
                        if output_path is not None:
                            p = output_path.format(**fmt_kwargs)
                            os.makedirs(os.path.dirname(p), exist_ok=True)
                            np.savez_compressed(p, cloud=cloud_struct)

                        # write segmented point cloud topic to a bag file
                        if output_bag is not None:
                            t = output_topic.format(topic=topic)
                            segmented_msg = msgify(PointCloud2, cloud_struct)
                            segmented_msg.header = msg.header
                            output_bag.write(t, segmented_msg, stamp)

    if output_bag:
        output_bag.close()


def main():
    args = arg_parser().parse_args()
    print(args)

    args.bag_paths = sum((glob(b) for b in args.bag_paths), start=[])
    print('Processing %i bag files:' % len(args.bag_paths), *args.bag_paths, sep='\n')

    if args.exclude_times:
        args.exclude_times = eval(args.exclude_times)
        args.exclude_times = [[t - args.time_horizon[1], t + args.time_horizon[1]]
                              if isinstance(t, (float, int)) else t
                              for t in args.exclude_times]
        print('Excluding times:', *args.exclude_times, sep='\n')
    if args.discard_model:
        args.discard_model = eval(args.discard_model)
        print('Discard model:', args.discard_model)
    if args.robot_model:
        args.robot_model = eval(args.robot_model)
        print('Robot model:', args.robot_model)
    if args.obstacles_model:
        args.obstacles_model = eval(args.obstacles_model)
        print('Obstacles model:', args.obstacles_model)

    if args.z_support is None:
        # TODO: obstacles generation settings for Cemicky les only
        args.z_support = {}
        scale = [1.0, 1.0, 0.05]
        scale = np.asarray(scale).reshape((-1, 3))
        args.z_support['scale'] = scale
        args.z_support['grid'] = 0.05
        args.z_support['radius'] = 0.05
        args.z_support['range'] = [0.6, 8.0]
        print('Z support:', args.z_support)

    generate_data(**vars(args))


if __name__ == '__main__':
    main()