pose_elevation_srv

#!/usr/bin/env python

"""
This node provides a service for obtaining the elevation of a point given by x,y (point_elevation_get)
or lat,lon (geopoints_elevation_get) and a coordinate frame.
"""

import rospy
import tf2_ros


from geometry_msgs.msg import Point
from geographic_msgs.msg import GeoPoint
from std_msgs.msg import String

from cuzk_tools.srv import PointElevationGet,PointElevationGetResponse
from cuzk_tools.srv import GeoPointElevationGet,GeoPointElevationGetResponse

import numpy as np
import os
from scipy.spatial import cKDTree
from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator


from cuzk_tools.dmr5g import Dmr5gParser, WGS_TO_SJTSK, get_utm_to_sjtsk_trans
from cuzk_tools.elevation_class import Elevation,UnsupportedFrameError


class PoseElevation(Elevation):
    def __init__(self, default_utm_zone):
        rospy.init_node('pose_elevation_srv')

        self.utm_frame = "utm"
        self.utm_local_frame = "utm_local"

        self.cache_dir = os.environ['HOME'] + "/.ros/cache/cuzk_tools/elevation/"

        if not os.path.exists(self.cache_dir):
            os.mkdir(self.cache_dir)
            
        self.elev_data_parser = Dmr5gParser(self.cache_dir)

        self.utm_zone = default_utm_zone
        self.utm_zone_sub = rospy.Subscriber("utm_zone", String, self.update_utm_zone, queue_size=10)

        self.utm_local_trans = None
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        self.tf_broadcaster = tf2_ros.TransformBroadcaster()
        
        rospy.Service('point_elevation_get', PointElevationGet, self.handle_point_elevation)
        rospy.Service('geopoint_elevation_get', GeoPointElevationGet, self.handle_geopoint_elevation)

        self.tile_id = None
        self.tree = None
        self.tile_data = None

    def coords2sjtsk(self, position, frame, response):
        sjtsk_coords = [0.,0.]

        if frame == "sjtsk":
            sjtsk_coords = position

        elif frame == "utm":
            if self.utm_zone is None:
                raise ValueError("utm_zone has not been set.")
            sjtsk_coords = get_utm_to_sjtsk_trans(self.utm_zone[2], self.utm_zone[:2]).transform(position[0],position[1])

        elif frame == "utm_local":
            try:
                self.utm_local_trans = self.tf_buffer.lookup_transform(self.utm_frame, self.utm_local_frame, rospy.Time())
            except:
                rospy.logwarn("pose_elevation_srv (utm_local): Cannot obtain transform (utm, utm_local).")
                if self.utm_local_trans is not None:
                    rospy.logwarn("pose_elevation_srv (utm_local): Using last known transform.")
                else:
                    rospy.logwarn("pose_elevation_srv (utm_local): Returning default pose msg.")
                    return response
                
            position[0] += self.utm_local_trans.transform.translation.x
            position[1] += self.utm_local_trans.transform.translation.y

            if self.utm_zone is None:
                raise ValueError("utm_zone has not been set.")
            
            sjtsk_coords = get_utm_to_sjtsk_trans(self.utm_zone[2], self.utm_zone[:2]).transform(position[0],position[1])
        
        elif frame == "wgs":
            sjtsk_coords = WGS_TO_SJTSK.transform(position[1],position[0])

        else:
            raise UnsupportedFrameError("Frame {} is not one of ('sjtsk','utm','utm_local','wgs').".format(frame))
        
        return sjtsk_coords
    
    def handle_point_elevation(self, req):
        position = [req.point.x, req.point.y]
        frame = req.frame.data
        response = PointElevationGetResponse()

        if frame == "wgs":
            rospy.loginfo("Consider using 'geopoint_elevation_get' service for 'wgs' frame. If you want to use this one make sure that lat-->y, lon-->x.")

        # Transform coordinates to S-JTSK.
        sjtsk_coords = self.coords2sjtsk(position, frame, response)
        
        # Get tile containing robot position.
        if not len(self.elev_data_parser.get_tile_ids(sjtsk_coords, 0)):
            rospy.logwarn("No elevation data for the given coordinates exists. Make sure the coordinates lie in Czehia and that they are correctly formatted and ordered.")
            rospy.logwarn("An example of coordinates that should work is x: -744572.45, y: -1042164.50")
            rospy.logwarn("Returning default msg.")
            return response
        
        tile_id = self.elev_data_parser.get_tile_ids(sjtsk_coords, 0)[0]

        # If tile is new, open it, create cKDTree (x,y only), save it.
        if tile_id == self.tile_id:
            pass
        else:
            self.tile_data = self.get_data(sjtsk_coords,None)

            if len(self.tile_data) == 0:
                rospy.logwarn("No data available. Returning default pose msg.")
                return response

            self.tile_data = self.tile_data.view((np.float64,len(self.tile_data.dtype.names)))
            self.tree = cKDTree(self.tile_data[:,:2])
            self.tile_id = tile_id

        # Use (saved) cKDTree to find nearest point.
        _,inds = self.tree.query(sjtsk_coords, k=10)
        nearest_points = self.tile_data[inds,:]
            
        # Interpolate to obtain elevation.
        elevation = self.interpolate_elevation(sjtsk_coords, nearest_points)

        # Create response with x,y,z. 
        response.point.x = req.point.x
        response.point.y = req.point.y
        response.point.z = elevation

        return response
    
    def handle_geopoint_elevation(self, req):
        if not req.frame.data == "wgs":
            rospy.logwarn("Frame '{}' is not 'wgs'. This will probably lead to wrong/unexpected results. If you want to use another frame then 'wgs' try 'point_elevation_get' service.".format(req.frame.data))
        p = Point(req.point.longitude, req.point.latitude, req.point.altitude)
        req.point = p

        response = self.handle_point_elevation(req)

        gp = GeoPoint()
        gp.latitude, gp.longitude, gp.altitude = response.point.y, response.point.x, response.point.z
        geo_response = GeoPointElevationGetResponse()
        geo_response.point = gp

        return geo_response
    
    def interpolate_elevation(self, point, near_points):
        try:
            interpolator = LinearNDInterpolator(near_points[:,:2], near_points[:,2])
        except:
            rospy.logwarn("Using nearest interpolation instead of linear due to error.")
            interpolator = NearestNDInterpolator(near_points[:,:2], near_points[:,2])

        elevation = interpolator(point)

        if np.isnan(elevation):
            rospy.logwarn("Using nearest interpolation instead of linear due to nan.")
            interpolator = NearestNDInterpolator(near_points[:,:2], near_points[:,2])
            elevation = interpolator(point)
        
        return elevation

if __name__ == "__main__":
    default_utm_zone = rospy.get_param('/pose_elevation_srv/default_utm_zone', None)
    node = PoseElevation(default_utm_zone)
    rospy.spin()