RoboRally/remote_control/position_controller.py

# startup:
# roscore
# rosrun cv_camera cv_camera_node

import sys
import rospy
import pygame
import numpy as np
import cv2
import cv2.aruco as aruco
import socket

from sensor_msgs.msg import Image
from sensor_msgs.msg import CompressedImage
from cv_bridge import CvBridge, CvBridgeError

pygame.init()
pygame.font.init()
#pygame.joystick.init()

myfont = pygame.font.SysFont('Comic Sans MS', 30)
pygame.display.set_caption("ROS camera stream on Pygame")
screenheight = 480
screenwidth = 640 #4*screenheight//3
screen = pygame.display.set_mode([screenwidth, screenheight])
red = (255, 0, 0)
teal = (0, 255, 255)

# ros setup
camera_stream = "/cv_camera/image_raw"
#camera_stream = "/image_raw"
compression = False

class Robot:
    def __init__(self, id, ip=None):
        self.pos = None

        self.id = id

        self.ip = None

class RemoteController:
    def __init__(self):
        self.cam = cv2.VideoCapture(2)

        #self.image_pub = rospy.Publisher("pygame_image", Image)

        #self.bridge = CvBridge()

        #if compression:
        #    self.image_sub = rospy.Subscriber(camera_stream + "/compressed", CompressedImage, self.callback)
        #else:
        #    self.image_sub = rospy.Subscriber(camera_stream, Image, self.callback)

        #self.cv_image = np.zeros((1, 1, 3), np.uint8)

        self.robots = [Robot(1), Robot(2), Robot(512)]

        screen.fill([0, 0, 0])

        self.aruco_dict = aruco.Dictionary_get(aruco.DICT_ARUCO_ORIGINAL)
        self.parameters = aruco.DetectorParameters_create()

        self.rc_socket = socket.socket()
        try:
            pass
            self.rc_socket.connect(('192.168.4.1', 1234))  # connect to robot
        except socket.error:
            print("could not connect to socket")

    def callback(self, data):
        try:
            if compression:
                self.cv_image = self.bridge.compressed_imgmsg_to_cv2(data)
            else:
                self.cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print(e)
        #self.cv_image_small = np.fliplr(self.cv_image_small) # why is this necessary?

        # marker detection
        #gray = cv2.cvtColor(self.cv_image, cv2.COLOR_BGR2GRAY)


                        #print("robot {} pos = {}".format(r.id, r.pos))

    def capture(self):
        ret_val, self.cv_image = self.cam.read()

        corners, ids, rejectedImgPoints = aruco.detectMarkers(self.cv_image, self.aruco_dict, parameters=self.parameters)
        rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners[0], 0.05, mtx, dist)
        marker_found = len(corners) > 0

        if marker_found:
            for i in range(len(corners)):
                #print("id = {}".format(ids[i]))
                pos = np.mean(corners[i][0], axis=0)

                for r in self.robots:
                    if r.id == ids[i]:
                        pos_flipped = (self.cv_image.shape[1] - int(pos[0]), int(pos[1]))
                        r.pos = pos_flipped

    def show_display(self):
        while True:
            self.capture()

            # show ros camera image on the pygame screen
            #image = cv2.resize(self.cv_image,(screenwidth,screenheight))
            frame = cv2.cvtColor(self.cv_image, cv2.COLOR_BGR2RGB)
            frame = np.rot90(frame)
            frame = pygame.surfarray.make_surface(frame)
            screen.blit(frame, (0, 0))

            # plot robot positions
            for r in self.robots:
                if r.pos is not None:
                    #print("drawing at {}".format(r.pos))
                    pygame.draw.circle(screen, (255, 0, 0), r.pos, 10)

            pygame.display.update()

            events = pygame.event.get()
            for event in events:
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT:
                        u1 = 1.0
                        u2 = -1.0
                        #print("turn left: ({},{})".format(u1, u2))
                    elif event.key == pygame.K_RIGHT:
                        u1 = -1.0
                        u2 = 1.0
                        #print("turn right: ({},{})".format(u1, u2))
                    elif event.key == pygame.K_UP:
                        u1 = -1.0
                        u2 = -1.0
                        #print("forward: ({},{})".format(u1, u2))
                    elif event.key == pygame.K_DOWN:
                        u1 = 1.0
                        u2 = 1.0
                        #print("forward: ({},{})".format(u1, u2))
                    self.rc_socket.send('({},{})\n'.format(u1, u2))
                elif event.type == pygame.KEYUP:
                    u1 = 0
                    u2 = 0
                    #print("key released, resetting: ({},{})".format(u1, u2))
                    self.rc_socket.send('({},{})\n'.format(u1, u2))


def main(args):
    rospy.init_node('controller_node', anonymous=True)

    rc = RemoteController()

    pygame.init()
    pygame.display.set_mode((640, 480))

    rc.show_display()
        #game.loop()

    # try:
    #     rospy.spin()
    # except KeyboardInterrupt:
    #     print("Shutting down")
    # cv2.destroyAllWindows()


if __name__ == '__main__':
    main(sys.argv)