ur5demo/demo.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pygame, sys, math, time, config
from ur5controller import DemoController


class UR5Demo(object):

    def __init__(self):
        # Set up pygame, GUI and controls
        pygame.init()
        self.screen = pygame.display.set_mode(config.SCREEN_DIM)
        pygame.display.set_caption("UR5 Demo")
        self.font = pygame.font.SysFont(None, config.FONTSIZE)

        # move vectors are in cylinder coords
        self.kb_map = {
                pygame.K_UP: (-1,0,0),
                pygame.K_DOWN: (1,0,0),
                pygame.K_LEFT: (0,-1,0),
                pygame.K_RIGHT: (0,1,0),
                pygame.K_w: (0,0,1),
                pygame.K_s: (0,0,-1)
        }

        self.clock = pygame.time.Clock()
        self.sig_quit = False
        self.loop_speed = config.DEFAULT_LOOP_SPEED
        self.controller = None

    def _disptxt(self, s_arr):
        self.screen.fill(config.BG_COL)
        for i,s in enumerate(s_arr):
            self.screen.blit(self.font.render(s, True, config.FONT_COL), (10, 10 + i*config.FONTSIZE))
        pygame.display.flip()

    def quit(self, error=None):
        msg = []
        if error:
            msg.append(error)
        msg.append("Disconnecting robot...")
        self._disptxt(msg)
        if self.controller:
            self.controller.cleanup()
        time.sleep(3) # Give the UR5 threads a chance to terminate
        msg.append("Exit...")
        self._disptxt(msg)
        pygame.quit()
        if error:
            return 1
        else:
            return 0

    def go_home_pos(self):
        self._disptxt(["Going to home configuration..."])
        self.controller.movel(config.R_MIN, 0, config.Z_MIN)

    def freedrive(self):
        self.controller.set_freedrive(True)
        self._disptxt(["Freedrive mode"])
        running = True
        while running:
            pressed = pygame.key.get_pressed()
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                   self.sig_quit = True
            if pressed[pygame.K_ESCAPE] or self.sig_quit:
                running = False

            self.clock.tick(config.MENU_FPS)
        self.controller.set_freedrive(False)

    def adjust_loop_speed(self):
        incr = 0.05
        running = True
        while running:
            pressed = pygame.key.get_pressed()
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                    self.sig_quit = True
            if pressed[pygame.K_PLUS] or pressed[pygame.K_KP_PLUS]:
                self.loop_speed += incr
                self.loop_speed = min(self.loop_speed, config.LOOP_SPEED_MAX)
            elif pressed[pygame.K_MINUS] or pressed[pygame.K_KP_MINUS]:
                self.loop_speed -= incr
                self.loop_speed = max(self.loop_speed, config.LOOP_SPEED_MIN)
            elif pressed[pygame.K_ESCAPE] or self.sig_quit:
                running = False

            self._disptxt([
                "Current speed: %.2f m/s" % self.loop_speed,
                "(+) Increase speed",
                "(-) Decrease speed"
            ])
            self.clock.tick(config.MENU_FPS)

    def run_loop1(self):
        loopcount = 0
        self._disptxt(["Running loop 1", "Loopcount: %d" % loopcount])
        time.sleep(1)
        
        # initiate loop, poll for robot move change and keypress
        running = True
        while running:
            pressed = pygame.key.get_pressed()
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                   self.sig_quit = True
            if pressed[pygame.K_ESCAPE] or self.sig_quit:
                running = False
            self.clock.tick(config.MENU_FPS)

        # Set to default pos before return?


    def run_loop2(self):
        loopcount = 0
        self._disptxt(["Running loop 2", "Loopcount: %d" % loopcount])
        time.sleep(1)
        pass


    def loop_init(self, block_init_msg, loopnum):
        msg = [
            "Initital block configuration for loop %d:" % loopnum,
            block_init_msg,
            " ",
            "Current loop speed is %.2f m/s" % self.loop_speed,
            "Start loop? [ENTER]"
        ]
        self._disptxt(msg)
        running = True
        while running:
            pressed = pygame.key.get_pressed()
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                   self.sig_quit = True

            if pressed[pygame.K_RETURN]:
                if loopnum == 1:
                    self.run_loop1()
                elif loopnum == 2:
                    self.run_loop2()
                running = False

            if pressed[pygame.K_ESCAPE] or self.sig_quit:
                running = False

            self.clock.tick(config.MENU_FPS)

    def keyboard_control(self):
        running = True
        while running:
            pressed = pygame.key.get_pressed()
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                   self.sig_quit = True
            if pressed[pygame.K_ESCAPE] or self.sig_quit:
                running = False

            vec = (0,0,0)
            for m in (self.kb_map[key] for key in self.kb_map if pressed[key]):
                vec = map(sum, zip(vec, m))

            self._disptxt([
                "Keyboard control",
                "Use arrow keys, [w], and [s]",
                " ",
                "move vec: %s" % str(vec)
            ])

            dt = 1.0/config.CTR_FPS
            try:
                self.controller.update(vec, dt)
            except Exception, e:
                self._disptxt(["Error: %s" % e])
                time.sleep(2)
                running = False

            self.clock.tick(config.CTR_FPS)


    def main_menu(self):
        menu = [
            "Go to [h]ome position",
            "[f]reedrive mode",
            "Adjust loop [s]peed",
            "Start loop [1]",
            "Start loop [2]",
            "Manual [c]ontrol",
            "[q]uit"
        ]
        self._disptxt(menu)
        running = True
        while running:
            pressed = pygame.key.get_pressed()
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                    self.sig_quit = True

            if pressed[pygame.K_h]:
                self.go_home_pos()
                self._disptxt(menu)

            elif pressed[pygame.K_f]:
                self.freedrive()
                self._disptxt(menu)

            elif pressed[pygame.K_s]:
                self.adjust_loop_speed()
                self._disptxt(menu)

            elif pressed[pygame.K_1]:
                _instr = "2x2 blocks on left side"
                self.loop_init(_instr, 1)
                self._disptxt(menu)

            elif pressed[pygame.K_2]:
                _instr = "2x2 blocks on left side, 2x2 blocks on right side"
                self.loop_init(_instr, 2)
                self._disptxt(menu)

            elif pressed[pygame.K_c]:
                self.keyboard_control()
                self._disptxt(menu)

            elif pressed[pygame.K_q] or self.sig_quit:
                running = False

            self.clock.tick(config.MENU_FPS)

        return self.quit()
        # break loop on quit, or return quit() on error

    def run(self):
        msg = []
        msg.append("Connecting robot...")
        self._disptxt(msg)
        try:
            self.controller = DemoController(config)
            self.controller.set_freedrive(False)
            msg.append("Calibrating cylinder coordinates...")
            self._disptxt(msg)
            self.controller.calibrate_csys()
        except Exception, e:
            return self.quit(error=("Error: %s" % e))
        msg.append("Calibration done.")
        self._disptxt(msg)
        time.sleep(1)

        return self.main_menu()


if __name__ == '__main__':
    demo = UR5Demo()
    sys.exit(demo.run())