~mrol-dev/mrol/trunk

''' File to perform simple scan-to-map alignment and map creation.'''
from __future__ import division
import cProfile
import cPickle
import copy
import sys

import roslib; roslib.load_manifest('mrol_ros')
import rospy
import tf
import message_filters
from sensor_msgs.msg import PointCloud2, CameraInfo, Image

import mrol_mapping.poseutil as poseutil
import mrol_mapping.mrol as mrol
import mrol_mapping.util as util
import mrol_mapping.occupiedlist as occupiedlist
import mrol_mapping.mapper as mapper
import mrol_mapping.pointcloud as pointcloud
import numpy as np
import rosutils_mrol.rosutils_mrol as rosutils
from thread import allocate_lock
from threading import Thread
import Queue

CAMERA_FRAME_ID = "camera_rgb_optical_frame"
RESOLUTION_PARAM_DEFAULT=0.04#0.04#0.05
RESOLUTION_PARAM="~resolution"
"""Resolution for mapping using MROL

Resolution determines how fine the map is and also how accurate the alignment
is, but finer = slower
"""

MINIMAL_REQUIRED_OVERLAP_DEFAULT=0.50
MINIMAL_REQUIRED_OVERLAP="~min_overlap"
"""Minimal overlap required between global map and a new scan so that the new
scan can be included in the global map
"""

MROL_LEVELS_DEFAULT=4
MROL_LEVELS="~mrol_levels"
"""Multi resolution levels for mapping.

If the resolution is small then this should be increased because the overlap
is restricted to an are pow(2, levels) * resolution.  Levels determines region
of convergence, more = larger region of convergences.  smaller is faster, but
may not converge
"""

OUT_SCAN_TOPIC_DEFAULT="/mrol_scan"
OUT_SCAN_TOPIC="~mrol_out_scan_topic"
"""Topic on which MROL scan should be published"""

OUT_MAP_TOPIC_DEFAULT="/mrol"
OUT_MAP_TOPIC="~mrol_out_map_topic"
"""Topic on which MROL map should be published"""

ODOM_FRAME="/odom"
BEST_POSE_FRAME="/mrol_bestpose"
BEST_POSE_GUESS_FRAME="/bestposeguess"
GLOBAL_FRAME="/map"
"""Frame used as a static frame relative to which point clouds are
published.
"""

MAX_RANGE_DEFAULT=10.0
MAX_RANGE="~max_range"
"""The data beyond this range will be discarded"""

class RosMapper:
    def __init__(self, num_base_scans=2, use_ICP=False):
        rospy.loginfo("Initializing")
        # can use mrol's inherent alignment method, or ICP
        self.use_ICP = use_ICP
        if self.use_ICP:
            import icp_mrol.icp_mrol as icp
            rospy.loginfo("Using ICP for alignment")
            # make an ICP instance to do the alignment
            self.icp = icp.ICP(reciprocal=False, visualise=False)

        self.res = rospy.get_param(RESOLUTION_PARAM, RESOLUTION_PARAM_DEFAULT)
        self.levels = rospy.get_param(MROL_LEVELS, MROL_LEVELS_DEFAULT)
        self.min_overlap = rospy.get_param(MINIMAL_REQUIRED_OVERLAP,
                                           MINIMAL_REQUIRED_OVERLAP_DEFAULT)
        self.max_range = rospy.get_param(MAX_RANGE, MAX_RANGE_DEFAULT)

        # variable initialization
        self.listener = tf.TransformListener(
            1,                  # interpolate ?
            rospy.Duration(20)) # max_cache_time
        self.bestpose = None

        self.bestpose = self.bestpose
        self.bestpose_lock = allocate_lock()

        self.bestposeguess = None
        self.bestposeguess_lock = allocate_lock()

        self.odom_wrt_map = None
        self.odom_wrt_map_lock = allocate_lock()

        self.lastodompose = None
        self.use_odom = True
        self.pc_frame_id_lock = allocate_lock()
        self.pc_frame_id = None

        #self.mapper = mapper.VoxelMap(resolution=self.res, levels=self.levels)
        self.lock = allocate_lock() # lock to have thread-safety
        self.num_base_scans = num_base_scans

        self.pointcloud_queue = util.Queue(maxsize=1)

        self.bestpose_queue = util.Queue(maxsize=1)
        self.bestpose_queue.put_force(poseutil.Pose3D())

        self.voxel_queue = util.Queue(maxsize=1)
        self.scan_queue = util.Queue(maxsize=1)

        # Threads
        Thread(target=self.publish_map_thread).start()
        # Transform broadcast should be at a higher frequency then the map
        # scan building, so it's better to have a separate thread for this.
        Thread(target=self.broadcast_transform_thread).start()
        # Let the mapping be in another thread
        Thread(target=self.mapping_thread).start()
        if self.use_odom:
            Thread(target=self.reframing_thread).start()

    def reframing_thread(self):
        # For debugging information only
        rate = rospy.Rate(10.0)
        br = tf.TransformBroadcaster()
        while not rospy.is_shutdown():
            with self.odom_wrt_map_lock:
                if self.odom_wrt_map is None:
                    continue
                robot_wrt_odom = self.get_pose(CAMERA_FRAME_ID, None)
                robot_wrt_map = np.dot(self.odom_wrt_map, robot_wrt_odom)

            # # The /kinect_rgb_optical_frame gets updated a little later than
            # # we get information via lookupTransform via rospy.Time(0).
            # # This is probably because of the PointCloud that rviz receives
            # # over /kinect_rgb_optical_frame. Perhaps rviz tries to get the
            # # frame at the timestamp of PointCloud.
            # translation = tuple(robot_wrt_odom[:3, 3])
            # quat = tf.transformations.quaternion_from_matrix(robot_wrt_odom)
            # br.sendTransform(translation,
            #                  quat,
            #                  rospy.Time.now(),
            #                  "/robot_wrt_odom",
            #                  ODOM_FRAME)

            translation = tuple(robot_wrt_map[:3, 3])
            quat = tf.transformations.quaternion_from_matrix(robot_wrt_map)

            br.sendTransform(translation,
                             quat,
                             rospy.Time.now(),
                             "/robot_wrt_map",
                             GLOBAL_FRAME)
            rate.sleep()

    def publish_map_thread(self):
        counter = 0
        topic_name = rospy.get_param(OUT_MAP_TOPIC, OUT_MAP_TOPIC_DEFAULT)
        pub = rospy.Publisher(topic_name, PointCloud2)
        while not rospy.is_shutdown():
            rospy.loginfo("publisher: Waiting for map to be generated")
            voxel_data, voxel_res = self.voxel_queue.get()
            pointsROS = np.hstack([voxel_data[:,:3]*voxel_res, voxel_data[:,4:]])
            pcROS = rosutils.xyzs2pc2(pointsROS, rospy.Time.now(), counter,
                                      frame_id=GLOBAL_FRAME)
            pub.publish(pcROS)
            rospy.loginfo("pubisher: Map published")
            counter += 1

    def mapping_thread(self):
        def timeout_func():
            return 1 if rospy.is_shutdown() else 30
        self.pointcloud_queue.set_timeout(timeout_func)

        try:
            mapper.buildmap(self.pointcloud_queue,
                            self.res,
                            visualise=False, # visualise using rviz
                            odometry_map=False, # odometry is unreliable
                            bestpose_queue=self.bestpose_queue,
                            voxel_queue=self.voxel_queue
                           )
        except Queue.Empty:
            sys.exit(0)

    def PCcallback(self, msg):
        rospy.loginfo("callback: Recieved point cloud")
        xyzs = pointcloud.PointCloud(
            rosutils.pc22xyzs(msg, max_range=self.max_range))
        # Nick: re-align x fwd, z up
        # 
        # (xyzs.points[:,:3], _) = poseutil.transformPoints(xyzs.points[:,:3],
        #                                             [0,0,0,-np.pi/2.,0,-np.pi/2.])
        # adjustment = poseutil.Pose3D([0,0,0,-np.pi/2.,0,-np.pi/2.]).getMatrix()
        self.process_pointcloud(xyzs, msg.header.frame_id, msg.header.stamp)

    def DepthImageCallback(self, msg):
        rospy.loginfo("callback: Recieved depth image")
        xyzs = rosutils.imgmsg2xyz(msg, max_range=self.max_range)
        assert(len(xyzs.points) > 0)
        if msg.header.frame_id != CAMERA_FRAME_ID:
            if msg.header.frame_id != "kinect_rgb_optical_frame":
                rospy.logerr("frame:{0}".format(msg.header.frame_id))
            msg.header.frame_id = CAMERA_FRAME_ID

        # Nick: re-align x fwd, z up
        # Vikas: align CAMERA_FRAME_ID and kinect_rgb_optical_frame 
        (xyzs.points[:,:3], _) = poseutil.transformPoints(
            xyzs.points[:,:3], [0,0,0,-np.pi/2., -np.pi/2.,-np.pi])
        self.process_pointcloud(xyzs, msg.header.frame_id, msg.header.stamp)


    def process_pointcloud(self, xyzs, frame_id, timestamp):
        robot_wrt_odom = None
        # Get the transformation of PointCloud frame w.r.t. odom global frame.
        with self.pc_frame_id_lock:
            self.pc_frame_id = frame_id
        if self.use_odom:
            rospy.loginfo("Received pointclound on"
                          + " frame_id:{0}".format(frame_id))
            robot_wrt_odom = self.get_pose(frame_id, timestamp)
            rospy.loginfo("{0}: {1}".format(CAMERA_FRAME_ID,
                poseutil.posetuple(robot_wrt_odom)))
            with self.odom_wrt_map_lock:
                if self.odom_wrt_map is None:
                    self.odom_wrt_map = np.linalg.inv(robot_wrt_odom)
        else:
            # Without odometry maporigin_to_coincide with odom
            with self.odom_wrt_map_lock:
                self.odom_wrt_map = tf.transformations.identity_matrix()

        with self.odom_wrt_map_lock:
            odom_wrt_map = copy.deepcopy(self.odom_wrt_map)
        rospy.loginfo("/odom_wrt_map : {0}".format(
            poseutil.posetuple(odom_wrt_map)))

        # Get the bestpose as returned by the mrol nearest neighbour algorithm
        try:
            bestpose = self.bestpose_queue.get_nowait().getMatrix()
            # Creating a copy for publishing frame in a different thread
            with self.bestpose_lock:
                rospy.loginfo("/bestpose : {0}".format(
                    poseutil.posetuple(bestpose)))
                self.bestpose = copy.deepcopy(bestpose)
        except Queue.Empty:
            rospy.logerr("Unable to get bestpose from mrol")
            # If no new result is available, use the last bestpose
            with self.bestpose_lock:
                if self.bestpose is not None:
                    bestpose = self.bestpose

        bestposeguess = self.guess_pose(robot_wrt_odom, bestpose)

        # Creating a copy for publishing frame in a different thread
        with self.bestposeguess_lock:
            rospy.loginfo("/bestposeguess: {0}".format(
                poseutil.posetuple(bestposeguess)))
            self.bestposeguess = copy.deepcopy(bestposeguess)

        # Save this odometry pose for future use
        self.lastodompose = robot_wrt_odom

        xyzs.pose = poseutil.Pose3D()
        xyzs.pose.setMatrix(bestposeguess)
        self.pointcloud_queue.put_force(xyzs)
    
        rospy.loginfo("callback: Finished processing point cloud")

    def guess_pose(self, robot_wrt_odom, bestpose):
        if self.use_odom:
            # Use incremental change in odometry to correct the last bestpose, for
            # a new bestposeguess
            bestposeguess = self.guess_pose_incremental_odom(robot_wrt_odom,
                                                             bestpose)
            # # For absolute dependence on odometery. Just project robot in map
            # # frame
            # bestposeguess = np.dot(self.odom_wrt_map, robot_wrt_odom)
        else:
            if bestpose is not None:
                bestposeguess = bestpose
            else:
                # This is first time we are getting a point cloud.
                # bestposeguess best to be identity 
                bestposeguess = tf.transformations().identity_matrix()
        return bestposeguess

    def guess_pose_incremental_odom(self, robot_wrt_odom, bestpose):
        """
        Use incremental change in odometry to correct the last bestpose, for
        a new bestposeguess
        """
        if self.lastodompose is not None:
            change_in_position = np.dot(np.linalg.inv(self.lastodompose),
                                        robot_wrt_odom)
            bestposeguess = np.dot(bestpose, change_in_position)
        else:
            # This is first time we are getting a point cloud.
            # bestposeguess best to be identity 
            bestposeguess = tf.transformations.identity_matrix()

        return bestposeguess

    def get_pose(self, frame_id, timestamp):
        # get transformations between world and robot at time of message
        # to have a guess for the position
        try:
            # We want the transform at time timestamp, but that time is
            # usually not available. We usually get an exception that cache is
            # emtpy. So we'll look for the latest transform available.
            # if timestamp:
            #     self.listener.waitForTransform(
            #         ODOM_FRAME, frame_id, timestamp, rospy.Duration(1.0))
            #     (trans, quat) = self.listener.lookupTransform(
            #         ODOM_FRAME, frame_id, timestamp)
            (trans, quat) = self.listener.lookupTransform(
                ODOM_FRAME, frame_id, rospy.Time(0))
            now = rospy.Time.now()
            if timestamp and now != timestamp:
                rospy.loginfo("Now:{0}; Timestamp:{1}".format(now, timestamp))
        except (tf.LookupException, tf.ConnectivityException):
            raise
        return quat2mat(quat, trans)

    def broadcast_transform_thread(self):
        """Thread to handle transform broadcast.  """
        rate = rospy.Rate(10.0)
        br = tf.TransformBroadcaster()
        while not rospy.is_shutdown():
            self.broadcast_one_transform(
                self.bestposeguess_lock,
                self.bestposeguess,
                BEST_POSE_GUESS_FRAME,
                br)
            self.broadcast_one_transform(
                self.bestpose_lock,
                self.bestpose,
                BEST_POSE_FRAME,
                br)

            if self.use_odom:
                self.broadcast_one_transform(
                    self.odom_wrt_map_lock,
                    self.odom_wrt_map,
                    ODOM_FRAME,
                    br)
            else:
                with self.pc_frame_id_lock:
                    ref_ros_frame = self.pc_frame_id
                if ref_ros_frame is not None:
                    # If we publish ref_ros_frame wrt GLOBAL_FRAME then we
                    # usually get the following warning and the rviz refuses
                    # to visualize relationship b/w kinect_rgb_optical_frame
                    # and /map
                    # [ WARN] [1331927940.889661623]: TF_OLD_DATA ignoring data from
                    # the past for frame /kinect_rgb_optical_frame at time 1.33039e+09
                    # according to authority /Kinect_mapper_demo
                    #
                    # So we publish map wrt /kinect_rgb_optical_frame
                    with self.bestposeguess_lock:
                        map_wrt_PCframe = np.linalg.inv(self.bestposeguess)
                    self.broadcast_one_transform(
                        self.bestposeguess_lock,
                        map_wrt_PCframe,
                        GLOBAL_FRAME,
                        br,
                        parent_frame_id=ref_ros_frame,
                    )
                    # camera does not show up in rviz
                    self.broadcast_one_transform(
                        self.bestposeguess_lock,
                        map_wrt_PCframe,
                        GLOBAL_FRAME,
                        br,
                        parent_frame_id='/camera_rgb_frame',
                    )

            rate.sleep()

    def broadcast_one_transform(self, lock, matrix, frame_id, broadcaster,
                                parent_frame_id=GLOBAL_FRAME):
        if matrix is None:
            return
        with lock:
            translation = tuple(matrix[:3, 3])
            quat = tf.transformations.quaternion_from_matrix(matrix)

        broadcaster.sendTransform(translation,
                quat,
                rospy.Time.now(),
                frame_id,
                parent_frame_id)

def quat2mat(quat, trans):
    """Converts quaternion and translation to 4x4 transformation matrix"""
    return np.dot(tf.transformations.translation_matrix(trans),
                  tf.transformations.quaternion_matrix(quat))
        
def run(ICP):
    ##########################
    # Initialize ROS interface
    ##########################
    rospy.init_node('Kinect_mapper_demo',log_level=rospy.DEBUG )
    # sleep a little - this avoids some ROS bugs
    rospy.sleep(2)
    orm = RosMapper(num_base_scans=2,use_ICP=ICP)
    # define how many scan should be without minimum overlay
    # this has to be at least 1, otherwise no map can be created
    # rospy.Subscriber("/camera/rgb/points", PointCloud2, orm.PCcallback, queue_size=5) 
    # depth_sub = message_filters.Subscriber("/camera/depth/image", Image)
    # image_sub = message_filters.Subscriber("/camera/rgb/color_image", Image)
    # info_sub = message_filters.Subscriber("/camera/rgb/camera_info",
    #                                     CameraInfo)
    # ts = message_filters.TimeSynchronizer([depth_sub, info_sub], 10)
    # ts.registerCallback(orm.DepthImageCallback)
    rospy.Subscriber("/camera/depth/image_raw", Image, orm.DepthImageCallback, queue_size=5) 
    rospy.spin()
    
def profile_yappi(func):
    try:
        import yappi
    except ImportError:
        rospy.logwarn("Yappi not installed. Not profiling")
        func()
        return

    yappi.start(True)
    func()
    with open("profile_Kinet_mapper_demo.yappi", "w") as f:
        cPickle.dump(yappi.get_stats(yappi.SORTTYPE_TTOTAL,
                                     yappi.SORTORDER_ASCENDING,
                                     yappi.SHOW_ALL),
                     f)

def run_and_catch_interrupt(ICP):
    try:
        run(ICP)
    except rospy.ROSInterruptException: pass


if __name__ == '__main__':
    ICP = '-ICP' in sys.argv
    profile_yappi(lambda:run_and_catch_interrupt(ICP))