cloud_segmentation

#!/usr/bin/env python

import os
import cv2
from datasets.laserscan import SemLaserScan
from datasets.base_dataset import TRAVERSABILITY_COLOR_MAP, TRAVERSABILITY_LABELS, VOID_VALUE
from datasets.base_dataset import FLEXIBILITY_COLOR_MAP, FLEXIBILITY_LABELS
from traversability_estimation.utils import get_label_map, convert_label, visualize_cloud, visualize_imgs
import rospy
from sensor_msgs.msg import Image, PointCloud2
from ros_numpy import msgify, numpify
import numpy as np
from numpy.lib.recfunctions import structured_to_unstructured, unstructured_to_structured
import torch
import yaml
from threading import RLock
import rospkg


# pkg_path = os.path.realpath(os.path.join(os.path.dirname(__file__), '..', '..'))
pkg_path = rospkg.RosPack().get_path('traversability_estimation')


def msgify_cloud(cloud, frame, stamp, names):
    assert cloud.ndim == 2
    cloud = unstructured_to_structured(cloud, names=names)
    msg = msgify(PointCloud2, cloud)
    msg.header.frame_id = frame
    msg.header.stamp = stamp
    return msg


def prob2entropy(p, axis=0, eps=1e-6):
    p = np.clip(p, eps, 1.0 - eps)
    h = np.sum(-p * np.log2(p), axis=axis)
    return h


class CloudSegmentor:
    def __init__(self, cloud_topic='cloud'):
        self.lidar_frame = None
        self.lidar_channels_H = int(rospy.get_param('~lidar_channels', 128))
        self.lidar_beams_W = int(rospy.get_param('~lidar_beams', 1024))
        self.lidar_fov_up = float(rospy.get_param('~lidar_fov_up', 45.0))
        self.lidar_fov_down = float(rospy.get_param('~lidar_fov_down', -45.0))
        self.range_projection = bool(rospy.get_param('~range_projection', False))

        self.lock = RLock()
        self.device = rospy.get_param('~device', 'cpu')

        self.model_weights = rospy.get_param('~weights')
        self.model_path = os.path.join(pkg_path, "config/weights/", "depth_cloud/%s" % self.model_weights)
        # assert os.path.exists(self.model_path)

        self.model = self.load_model()
        rospy.loginfo('Loaded cloud segmentation model: %s', self.model_weights)

        self.data_fields = [f[1:-1] for f in ['_x_', '_y_', '_z_', '_intensity_', '_depth_'] if f in self.model_weights]
        if not self.data_fields:
            self.data_fields = ['depth']
        rospy.loginfo('Model takes as input: %s' % ','.join(self.data_fields))

        self.model_output = 'labels'
        if 'traversability' in self.model_weights.lower():
            self.model_output = 'traversability'
        elif 'flexibility' in self.model_weights.lower():
            self.model_output = 'flexibility'

        self.input_pc_fields = ['x', 'y', 'z']
        self.output_pc_fields = ['x', 'y', 'z', self.model_output, 'cost', 'entropy']
        self.cost_ind = rospy.get_param('~soft_label_ind', 1)

        assert self.model_output in ['labels', 'traversability', 'flexibility']
        if self.model_output == 'labels':
            self.label_map = None
            cfg = yaml.safe_load(open(os.path.join(pkg_path, "config/rellis.yaml"), 'r'))
            self.color_map = cfg["color_map"]
            self.CLASSES = [v for v in cfg['labels'].values()]
            self.class_values = list(range(len(self.color_map)))
            self.learning_map = cfg["learning_map"]
            self.learning_map_inv = cfg["learning_map_inv"]
            self.ignore_label = 0
        else:
            self.ignore_label = VOID_VALUE
            if self.model_output == 'traversability':
                self.color_map = TRAVERSABILITY_COLOR_MAP
                self.CLASSES = [v for k, v in TRAVERSABILITY_LABELS.items()]
                self.class_values = np.sort([k for k in TRAVERSABILITY_LABELS.keys()]).tolist()
            elif self.model_output == 'flexibility':
                self.color_map = FLEXIBILITY_COLOR_MAP
                self.CLASSES = [v for k, v in FLEXIBILITY_LABELS.items()]
                self.class_values = np.sort([k for k in FLEXIBILITY_LABELS.keys()]).tolist()

            self.label_map = get_label_map(path=os.path.join(pkg_path, "config/rellis_to_%s.yaml" % self.model_output))

        self.n_classes = len(self.CLASSES)
        self.non_bg_classes = np.asarray(self.CLASSES)[np.asarray(self.class_values) != self.ignore_label]

        self.scan = SemLaserScan(nclasses=len(self.non_bg_classes),
                                 sem_color_dict=self.color_map,
                                 project=True,
                                 H=self.lidar_channels_H, W=self.lidar_beams_W,
                                 fov_up=self.lidar_fov_up, fov_down=self.lidar_fov_down)
        self.debug = rospy.get_param('~debug', False)

        # point cloud which time stamp is older is not being processed
        self.max_age = rospy.get_param('~max_age', 0.5)

        self.segm_cloud_pub = rospy.Publisher(rospy.get_param('~cloud_out', 'cloud_out'), PointCloud2, queue_size=1)
        self.resized_cloud_pub = rospy.Publisher('~resized_cloud', PointCloud2, queue_size=1)
        self.depth_pub = rospy.Publisher('~depth', Image, queue_size=1)

        self.cloud_sub = rospy.Subscriber(cloud_topic, PointCloud2, self.segment_cloud_cb)
        rospy.loginfo('Point cloud segmentation node is ready.')

    def label_to_color(self, label):
        if len(label.shape) == 3:
            C, H, W = label.shape
            label = np.argmax(label, axis=0)
            assert label.shape == (H, W)
        if self.model_output == 'labels':
            label = convert_label(label, inverse=False, label_mapping=self.learning_map_inv)
        color = self.scan.sem_color_lut[label]
        return color

    def load_model(self):
        if not os.path.exists(self.model_path):
            url = 'http://subtdata.felk.cvut.cz/robingas/data/traversability_estimation/weights/depth_cloud/'
            path = os.path.abspath(os.path.join(self.model_path, os.pardir))
            print('Downloading model weights by running:')
            print('wget %s/%s -P %s' % (url, self.model_weights, path))
            os.system('wget %s/%s -P %s' % (url, self.model_weights, path))
        model = torch.load(self.model_path, map_location=self.device)
        model = model.eval()
        return model

    def resize_cloud(self, cloud):
        H, W = cloud.shape[:2]
        self.lidar_channels_H = H
        self.lidar_beams_W = W
        resized = []
        for i in range(cloud.shape[2]):
            c = cv2.resize(cloud[..., i], (self.lidar_channels_H, self.lidar_beams_W),
                           interpolation=cv2.INTER_LINEAR)
            resized.append(c)
        resized = np.stack(resized, axis=-1)
        cloud = resized
        rospy.loginfo('Point cloud resized to %s', resized.shape)

        return cloud

    def preprocessing(self, cloud):
        if cloud.ndim == 3:
            H, W, C = cloud.shape
            if self.lidar_channels_H != H or self.lidar_beams_W != W:
                self.lidar_channels_H = H
                self.lidar_beams_W = W
                self.scan = SemLaserScan(nclasses=self.n_classes,
                                         sem_color_dict=self.color_map,
                                         project=True,
                                         H=H, W=W,
                                         fov_up=self.lidar_fov_up, fov_down=self.lidar_fov_down)
                cloud = self.resize_cloud(cloud)
                rospy.loginfo('Point cloud resized to %s', cloud.shape)

        if self.range_projection:
            self.scan.set_points(points=cloud[..., :3].reshape((-1, 3)),
                                 remissions=cloud[..., 3].reshape((-1, 3)) if cloud.shape[-1] >= 4 else None)
            depth = self.scan.proj_range
        else:
            depth = np.linalg.norm(cloud[..., :3], ord=2, axis=-1)

        depth = depth.reshape((self.lidar_channels_H, self.lidar_beams_W))

        # depth_vis = cv2.resize(depth.astype('float'),
        #                        (depth.shape[1] // 2, depth.shape[0] // 2),
        #                        interpolation=cv2.INTER_LINEAR)
        # cv2.imshow('Depth', depth_vis)
        # cv2.waitKey(1)

        depth = depth[None]
        rospy.logdebug('Model input shape: %s', depth.shape)
        assert depth.shape == (1, self.lidar_channels_H, self.lidar_beams_W)

        return depth

    def model_inference(self, depth):
        # Apply inference preprocessing transforms
        batch = torch.from_numpy(depth).unsqueeze(0).to(self.device)
        with torch.no_grad():
            pred = self.model(batch)['out']
        rospy.loginfo('Segmented result shape: %s', pred.shape)
        return pred

    def postprocessing(self, pred_prob, xyz):
        assert isinstance(pred_prob, torch.Tensor)
        assert isinstance(xyz, np.ndarray)

        pred_prob = torch.softmax(pred_prob.squeeze(0), dim=0).cpu().numpy()
        assert pred_prob.ndim == 3

        cost = pred_prob[self.cost_ind]

        # label_soft_vis = cv2.resize(cost.astype('float'), (cost.shape[1] // 2, cost.shape[0] // 2),
        #                             interpolation=cv2.INTER_LINEAR) / cost.max()
        # cv2.imshow('Predicted labels', label_soft_vis)
        # cv2.waitKey(1)

        n_pts = cost.shape[0] * cost.shape[1]

        xyz = xyz.reshape((n_pts, 3))

        pred_class = np.argmax(pred_prob, axis=0)
        entropy = prob2entropy(pred_prob, axis=0)

        cost = cost.reshape((n_pts, 1))
        pred_class = pred_class.reshape((n_pts, 1))
        entropy = entropy.reshape((n_pts, 1))

        output = np.concatenate([xyz, pred_class, cost, entropy], axis=1)
        assert output.shape == (n_pts, 6)

        rospy.logdebug('Output cloud shape: %s', output.shape)
        return np.asarray(output, dtype=np.float32)

    def segment_cloud_cb(self, pc_msg):
        assert isinstance(pc_msg, PointCloud2)
        self.lidar_frame = pc_msg.header.frame_id

        # Discard old messages.
        msg_stamp = rospy.Time.now()
        age = (msg_stamp - pc_msg.header.stamp).to_sec()
        if age > self.max_age:
            rospy.logwarn('Cloud segmentation: Discarding points %.1f s > %.1f s old.', age, self.max_age)
            return

        t0 = rospy.Time.now().to_sec()

        with self.lock:
            # Transform local map to ground truth localization frame
            cloud = numpify(pc_msg)
            if not set(self.input_pc_fields) <= set(cloud.dtype.names):
                rospy.logwarn('Point cloud does not contain all requested input fields. \nUsing only "x y z".')
                self.input_pc_fields = ['x', 'y', 'z']
            cloud = structured_to_unstructured(cloud[self.input_pc_fields])
            rospy.logdebug('Point cloud of shape %s is received', cloud.shape)

            depth = self.preprocessing(cloud)
            t1 = rospy.Time.now().to_sec()
            rospy.logdebug('Preprocessing took: %.3f [sec]' % (t1 - t0))

            pred = self.model_inference(depth)
            t2 = rospy.Time.now().to_sec()
            rospy.logdebug('Model inference took: %.3f [sec]' % (t2 - t1))

            if self.range_projection:
                xyz = self.scan.proj_xyz
            else:
                xyz = cloud[..., :3]
            output = self.postprocessing(pred, xyz=xyz)
            t3 = rospy.Time.now().to_sec()
            rospy.logdebug('Postprocessing took: %.3f [sec]' % (t3 - t2))

            # grass_mask = output[:, 6] == 1
            # bushes_mask = output[:, 6] == 9
            # mean_height = min(np.mean(output[grass_mask, 2]), np.mean(output[bushes_mask, 2]))
            # output[grass_mask, 2] = mean_height
            # output[bushes_mask, 2] = mean_height

            # publish segmented point cloud
            segm_pc_msg = msgify_cloud(output, frame=pc_msg.header.frame_id, stamp=pc_msg.header.stamp,
                                       names=self.output_pc_fields)
            self.segm_cloud_pub.publish(segm_pc_msg)

            if self.debug:
                # publish depth image
                power = 16
                depth_img = np.copy(depth.squeeze())  # depth
                depth_img[depth_img > 0] = depth_img[depth_img > 0] ** (1 / power)
                depth_img[depth_img > 0] = (depth_img[depth_img > 0] - depth_img[depth_img > 0].min()) / \
                                           (depth_img[depth_img > 0].max() - depth_img[depth_img > 0].min())
                depth_msg = msgify(Image, depth_img, '32FC1')
                depth_msg.header.stamp = pc_msg.header.stamp
                depth_msg.header.frame_id = pc_msg.header.frame_id
                self.depth_pub.publish(depth_msg)

            t4 = rospy.Time.now().to_sec()
            rospy.logdebug('Data publishing took: %.3f [sec]' % (t4 - t3))

        rospy.loginfo('Point cloud processing time: %.3f [sec]' % (t4 - t0))


if __name__ == '__main__':
    rospy.init_node('cloud_segmentation', log_level=rospy.INFO)
    proc = CloudSegmentor(cloud_topic=rospy.get_param('~cloud_in', 'cloud_in'))
    rospy.spin()