class TurtlebotV2

from sensor_msgs.msg import Image
from nav_msgs.msg import Odometry
from math import atan2

import numpy as np
import rospy

class TurtlebotV2:

    def __init__(self, namespace='/tb0', timeout=1.0):
        self.namespace = namespace
        self.timeout= timeout

    def getRobotPose(self):
        rospy.init_node('get_odom')
        try:
            odom = rospy.wait_for_message(self.namespace + '/odom', Odometry, timeout=self.timeout)
            theta = 2*atan2(odom.pose.pose.orientation.z, odom.pose.pose.orientation.w)
            pose = [odom.pose.pose.position.x, odom.pose.pose.position.x, theta]

        except rospy.ROSException:
            # Timeout
            pose = [0.0, 0.0, 0.0]

        return pose

    def getKinectRGBImage(self):
        try:
            rospy.init_node('get_image')
            rgb = rospy.wait_for_message(self.namespace + '/camera/rgb/image_color', Image, timeout=self.timeout)
            img = np.frombuffer(rgb.data, dtype=np.uint8).reshape(rgb.height, rgb.width, -1)

            # Convert to cv2

        except rospy.ROSException:
            # Timeout
            img = None

        return img