AirSim-ROS RL environment

#!/usr/bin/env python3
import rospy
import time
from sensor_msgs.msg import Image
from airsim_ros_pkgs.msg import VelCmd

from cv_bridge import CvBridge
import numpy as np
import airsim

import torch
from torchvision import transforms as ttf

rospy.init_node("copy_node")

class copy_class:
    def __init__(self):
        self.seg_sub = rospy.Subscriber("/airsim_node/SimpleFlight/segmentation/Segmentation", Image, self.seg_cb)
        self.color_sub = rospy.Subscriber("/airsim_node/SimpleFlight/color/Scene", Image, self.color_cb)
        self.img_pub = rospy.Publisher("/easytopic", Image, queue_size = 5)
        self.cmd_pub = rospy.Publisher("/airsim_node/SimpleFlight/vel_cmd_body_frame", VelCmd, queue_size = 10)

        self.seg_in = Image()
        self.road_mask = None
        self.color_in = Image()
        self.cmd_stop = VelCmd()

        self.bridge = CvBridge()
        self.client = airsim.MultirotorClient()
        self.client.confirmConnection()

    def seg_cb(self, msg):
        self.seg_in = msg
        my_img = self.bridge.imgmsg_to_cv2(self.seg_in, desired_encoding = "passthrough")
        my_img = my_img[:,:,0].copy()
        my_img[my_img != 246] = 0
        my_img[my_img == 246] = 255

        my_img = self.bridge.cv2_to_imgmsg(my_img, encoding = "mono8")
        self.img_pub.publish(my_img)

    def color_cb(self, msg):
        self.color_in = msg

    def env_recurse_reset(self):
        while not rospy.is_shutdown():
            self.cmd_stop = VelCmd()
            self.cmd_pub.publish(self.cmd_stop)
            self.client.reset()
            self.client.enableApiControl(True)
            self.client.armDisarm(True)
            self.cmd_stop.twist.linear.z = -5
            self.cmd_stop.twist.linear.x = 0
            prev_time = rospy.get_time()
            next_time = rospy.get_time()
            while next_time - prev_time < 5.0:
                self.cmd_pub.publish(self.cmd_stop)

                ### Put training code here ###

                next_time = rospy.get_time()
            time.sleep(5)


copy_class = copy_class()
copy_class.env_recurse_reset()
rospy.spin()