generate_flexibility_data

#!/usr/bin/env python
"""Generate self-supervised data from ROS bag files with robot trajectories.

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
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
from tqdm import tqdm

DEPTH_FIELD_NAME = 'depth'
# LABEL_FIELD_NAME = 'flexible'
LABEL_FIELD_NAME = 'label'
GROUND_FLAG_FIELD_NAME = 'ground'
GROUND_Z_FIELD_NAME = 'ground_z'
LABEL_RIGID = 0
LABEL_FLEXIBLE = 1
# LABEL_INVALID = 254
LABEL_IGNORE = 255


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('--topics', type=str, nargs='+')
    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('--flexible-model', type=str, default=None,
                        help='Model to move along robot path marking points flexible.')
    parser.add_argument('--rigid-model', type=str, default=None,
                        help='Model to move along robot path marking points rigid.')
    parser.add_argument('--ground-height', type=float, default=0.0,
                        help='Height of the ground relative to the robot frame.')
    parser.add_argument('--rigid-primitives', type=str2bool, default=False,
                        help='Fit primitives and mark the inliers as rigid.')
    parser.add_argument('--z-support', type=str, default=None, help='Z support configuration or empty.')
    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}/{topic}/{secs}_{nsecs:09d}.npz')
    parser.add_argument('--cloud-array', type=str, default='cloud')
    parser.add_argument('--depth-array', type=str, default='depth')
    parser.add_argument('--label-array', type=str, default='target')
    parser.add_argument('--ground-height-array', type=str, default='ground_height')
    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('--ground-topic', type=str, default='{topic}_ground')
    parser.add_argument('--output-period', type=float, default=None)
    parser.add_argument('--visualize', type=str2bool, default=False)
    parser.add_argument('bag_paths', type=str, nargs='+')
    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):
    tf_topics = ['/tf', '/tf_static']
    # buffer = BufferCore(cache_time=rospy.Duration(2**31 - 1))
    # buffer = BufferCore(cache_time=rospy.Duration(24 * 60 * 60))
    buffer = BufferCore(rospy.Duration(24 * 60 * 60))
    for path in bag_paths:
        try:
            with Bag(path, 'r') as bag:
                for topic, msg, stamp in tqdm(bag.read_messages(topics=tf_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:
                            buffer.set_transform(tf, 'bag')
                    elif topic == '/tf_static':
                        for tf in msg.transforms:
                            buffer.set_transform_static(tf, 'bag')
        except ROSBagException as ex:
            print('Could not read %s: %s' % (path, ex))

    return 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 compute_ground_offset(x, robot_to_fixed, ground_height):
    """
    @param x: 3 x N array of points in the robot frame.
    @param robot_to_fixed: 4 x 4 transformation matrix from robot frame to fixed frame.
    @param ground_height: Height of the ground relative to the robot frame.
    """
    # Construct ground plane equation n * x + b = 0 from robot pose in fixed frame.
    n = robot_to_fixed[:3, 2:3]
    o = robot_to_fixed[:3, 3:] + n * ground_height
    b = -np.matmul(n.T, o)
    # Compute ground height offset for points in fixed frame.
    y = affine(robot_to_fixed, x)
    d = np.matmul(n.T, y).ravel() + b.ravel()  # signed distance to plane
    z_offset_fixed = -d / n[2, 0]  # ground offset in fixed frame
    offset_fixed = np.zeros_like(x)
    offset_fixed[2, :] = z_offset_fixed
    # Compute ground height offset for points in robot frame.
    offset = np.matmul(robot_to_fixed[:3, :3].T, offset_fixed)
    return offset, z_offset_fixed


def segment_cloud(flexible_model, arr, input_to_robot_tfs,
                  discard_tf=None, discard_model=None,
                  rigid_model=None, ground_height=None,
                  rigid_primitives=True, input_to_fixed=None, z_support=None,
                  visualize=False):
    assert flexible_model is None or isinstance(flexible_model, Body)
    assert isinstance(arr, np.ndarray)
    assert rigid_model is None or isinstance(rigid_model, Body)
    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=np.uint8).ravel()
    ground_z_offset = np.full(arr.shape, np.nan, dtype=np.float32).ravel()

    # Mark valid points which are contained by future model poses as empty.
    flexible = np.zeros((x.shape[1],), dtype=bool)
    rigid = np.zeros((x.shape[1],), dtype=bool)
    for input_to_robot in input_to_robot_tfs:
        assert isinstance(input_to_robot, np.ndarray)
        x_robot = affine(input_to_robot, x)
        if flexible_model or ground_height is not None:
            tmp_flexible = flexible_model.contains_point(x_robot)
            flexible = np.logical_or(flexible, tmp_flexible)
        if rigid_model:
            tmp_rigid = rigid_model.contains_point(x_robot)
            rigid = np.logical_or(rigid, tmp_rigid)
        if ground_height is not None:
            robot_to_fixed = np.matmul(inverse(input_to_robot), input_to_fixed)
            ground_offset_robot, ground_z_offset_fixed = compute_ground_offset(x_robot, robot_to_fixed, ground_height)
            y_ground = x_robot + ground_offset_robot
            # Offset is valid only for flexible points with their ground
            # location within the rigid model.
            ground_valid = tmp_flexible & rigid_model.contains_point(y_ground)
            # Take minimum offset from all applicable robot positions.
            ground_all = valid_ind[ground_valid]
            ground_z_offset[ground_all] = np.nanmin([ground_z_offset[ground_all],
                                                     ground_z_offset_fixed[ground_valid]], axis=0)

    print('%.3g = %i / %i flexible points' % (flexible.mean(), flexible.sum(), flexible.size))
    print('%.3g = %i / %i rigid points identified.' % (rigid.mean(), rigid.sum(), rigid.size))
    valid_ground = np.logical_not(np.isnan(ground_z_offset))
    print('%.3g = %i / %i ground points identified (mean offset: %.3g m).'
          % (valid_ground.mean(), valid_ground.sum(), valid_ground.size, np.nanmean(ground_z_offset)))

    # Fit geometric primitives to points and mark their inliers as rigid.
    if rigid_primitives:
        primitives = fit_primitives(x)
        for primitive, indices in primitives:
            rigid[indices] = True

    if z_support:
        support, tmp_rigid = compute_rigid_support(arr_valid, transform=input_to_fixed, **z_support)
        # rigid[support > 30] = True
        rigid[tmp_rigid] = True

    labels[valid_ind[rigid]] = LABEL_RIGID
    # Flexible label takes precedence.
    labels[valid_ind[flexible]] = LABEL_FLEXIBLE
    # Don't correct rigid points.
    ground_z_offset[valid_ind[rigid]] = 0.0

    if visualize:
        show_cloud(x.T, labels[valid_ind], min=0, max=2, colormap=cm.jet)

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


def create_ground_cloud(arr, ground_z_offset, input_to_fixed, visualize=False):
    x = position(arr.ravel())

    # Convert ground z offset from fixed to input frame to correct positions.
    ground_z_offset = ground_z_offset.ravel()
    valid = ((x != 0) & np.isfinite(x)).all(axis=1)
    valid = valid & np.isfinite(ground_z_offset) & (ground_z_offset != 0)
    # print('%.3g valid ground points.' % valid)
    print('%.3g = %i / %i ground points identified (mean offset: %.3g m).'
          % (valid.mean(), valid.sum(), valid.size, ground_z_offset[valid].mean()))
    offset_fixed = np.zeros_like(x)
    offset_fixed[:, 2] = ground_z_offset
    # Multiply by inverse of input_to_fixed to convert from fixed to input frame,
    # Inverse/transpose is cancelled by multiplying from right.
    # offset = np.matmul(input_to_fixed[:3, :3].T, offset_fixed.T)
    offset = np.matmul(offset_fixed, input_to_fixed[:3, :3]).astype(dtype=np.float32)
    x_ground = x + offset

    x_all = np.concatenate((x, x_ground[valid]), axis=0)
    x_all_struct = unstructured_to_structured(x_all, names=['x', 'y', 'z'])

    flag_all = np.concatenate((np.zeros(x.shape[:1], dtype=np.uint8),
                               np.ones(x_ground[valid].shape[:1], dtype=np.uint8)), axis=0).reshape((-1, 1))
    flag_all_struct = unstructured_to_structured(flag_all, names=[GROUND_FLAG_FIELD_NAME])

    cloud = merge_arrays([x_all_struct, flag_all_struct], flatten=True)
    # cloud = cloud.reshape(arr.shape)

    if visualize and valid.sum() > 0:
        # viz_x = np.concatenate((x[valid], ground[valid]), axis=0)
        # viz_v = np.concatenate((np.zeros(x[valid].shape[:1]), np.ones(ground[valid].shape[:1])))
        # viz_x = np.concatenate((x, x_ground[valid]), axis=0)
        # viz_v = np.concatenate((np.zeros(x.shape[:1]), np.ones(x_ground[valid].shape[:1])))
        # show_cloud(x[valid], ground_z_offset.ravel()[valid], colormap=cm.jet)
        # show_cloud(viz_x, viz_v, colormap=cm.jet)
        show_cloud(x_all, flag_all.ravel(), colormap=cm.jet)

    return cloud


def generate_data(bag_paths=None, topics=None, fixed_frame=None, robot_frame=None,
                  exclude_times=None, input_step=1, input_start=0.0, input_end=float('inf'),
                  discard_model=None, flexible_model=None, rigid_model=None, ground_height=None,
                  rigid_primitives=False, z_support=None,
                  discard_empty=True, distance_horizon=None, time_horizon=None, time_step=None,
                  output_path=None, cloud_array=None, depth_array=None, label_array=None, ground_height_array=None,
                  output_bag_path=None, output_topic=None, ground_topic=None, output_period=None, visualize=False):
    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

    buffer = load_buffer(bag_paths)
    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=topics),
                                          desc='%s: generating data' % bag_path.split('/')[-1],
                                          total=bag.get_message_count(topic_filters=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 = 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 = numpify(input_to_fixed.transform)

                # Find transforms from input cloud to robot positions within the horizon.
                input_to_robot_tfs = []
                for t in np.linspace(start - n[0] * time_step, start + n[1] * time_step, sum(n) + 1):
                    try:
                        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)

                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, ground_z_offset = segment_cloud(flexible_model, input_struct, input_to_robot_tfs,
                                                           rigid_model=rigid_model, ground_height=ground_height,
                                                           discard_tf=input_to_robot_tfs[n[0]], discard_model=discard_model,
                                                           rigid_primitives=rigid_primitives,
                                                           input_to_fixed=input_to_fixed, 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])

                    ground_z_offset_struct = ground_z_offset.reshape((input_struct.size, -1))
                    ground_z_offset_struct = unstructured_to_structured(ground_z_offset_struct, names=[GROUND_Z_FIELD_NAME])

                    cloud_struct = merge_arrays([input_struct, depth_struct, label_struct, ground_z_offset_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()
                    n_flexible = (label == LABEL_FLEXIBLE).sum()
                    if discard_empty and (n_valid_labels < 1000 or n_flexible < 100):
                        print('Discarding cloud with no valid labels.')
                        continue
                    else:
                        print('Storing cloud with %i valid labels.' % n_valid_labels)

                    # Create cloud with ground points for visualization and debugging.
                    # ground = create_ground_cloud(input_struct, ground_z_offset_struct, input_to_fixed)
                    ground_struct = create_ground_cloud(input_struct, ground_z_offset, input_to_fixed,
                                                        visualize=visualize)
                    # ground_struct = ground.reshape((input_struct.size, -1))
                    # ground_struct = unstructured_to_structured(ground_struct, names=['x', 'y', 'z'])

                    last_out[topic] = start

                    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': segmented_arr})
                        arrays = {}
                        if cloud_array:
                            arrays[cloud_array] = cloud_struct
                        if depth_array:
                            arrays[depth_array] = depth
                        if ground_height_array:
                            arrays[ground_height_array] = ground_z_offset
                        if label_array:
                            arrays[label_array] = label
                        np.savez_compressed(p, **arrays)

                    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)

                        t = ground_topic.format(topic=topic)
                        ground_msg = msgify(PointCloud2, ground_struct)
                        ground_msg.header = msg.header
                        output_bag.write(t, ground_msg, stamp)

    if output_bag:
        output_bag.close()


def main():
    args = arg_parser().parse_args()
    print(args)
    if args.z_support is not None:
        args.z_support = eval(args.z_support)
        if 'scale' in args.z_support:
            scale = args.z_support['scale']
            if isinstance(scale, (float, int)):
                scale = [1.0, 1.0, scale]
            scale = np.asarray(scale).reshape((-1, 3))
            args.z_support['scale'] = scale
        print('Z support:', args.z_support)

    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.flexible_model:
        args.flexible_model = eval(args.flexible_model)
        print('Flexible model:', args.flexible_model)
    if args.rigid_model:
        args.rigid_model = eval(args.rigid_model)
        print('Rigid model:', args.rigid_model)

    generate_data(**vars(args))


if __name__ == '__main__':
    main()