MotorControlGui.py

download
#!/usr/bin/pythonw

'''
File: MotorControlGui.py
Author: Aaron M. Hoover
Date: 2009-12-27 15:27:32 PST
Description: A wxPython GUI for interfacing to a robot and sending commands. 
'''

import wx
from serial import Serial, SerialException
import time
import os
import struct
import sys
from numpy import array, savetxt, arange
import threading
from multiprocessing import Process, Queue, Value
from Queue import Empty
import wx.lib.newevent

import stxetx
from util import cvtStrideFreqToEMF
from consts import BAUD_RATE, STATUS_UNUSED, PCKT_CMD, \
                   MAX_RECEIVE_BUFFER_LEN, SET_MOTOR_SPEED, \
                   CONFIG_PID, SENSE_DATA, DISCONNECT, \
                   CONNECT, MAX_EMF, Y_MIN, Y_MAX

import matplotlib
matplotlib.use('WXAgg')
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import \
    FigureCanvasWxAgg as FigCanvas
import pylab

(UpdatePlotEvent, EVT_UPDATE_PLOT) = wx.lib.newevent.NewEvent()
(SerialStatusEvent, EVT_UPDATE_SERIAL_STATUS) = wx.lib.newevent.NewEvent()

class QueueData(object):
    """An object for holding queue data consisting of a command and payload."""
    def __init__(self, cmd_id, payload):
        self.cmd_id = cmd_id
        self.payload = payload

def send_receive(tx_queue, rx_queue, connected, num_packets):
    conn = None
    rx_buffer = []
    ser = stxetx.StxEtx()

    while True:
        try:
            cmd = tx_queue.get_nowait()
            if cmd.cmd_id == CONNECT:
                if conn == None:
                    port_name = cmd.payload[0]
                    conn = Serial(port_name, BAUD_RATE, timeout=0,
                                  rtscts=1)
                    connected.value = conn.isOpen()
                    ser.set_data_out(conn.write)
                else:
                    conn.open()
                    connected.value = conn.isOpen()
                success = connected.value
                if success:
                    conn.flushInput()
                    conn.flushOutput()
            elif cmd.cmd_id == DISCONNECT:
                if connected.value:
                    conn.close()
                    connected.value = False
                success = not connected.value
            else:
                if connected.value:
                    ser.send(STATUS_UNUSED, cmd.cmd_id, cmd.payload)
                else:
                    print "Did not attempt to send command because " + \
                        "there is no connection."
                    success = False
        except Empty:
            success = None
        except SerialException, ex:
            print ex.message
        finally:
            if success != None:
                rx_queue.put_nowait(QueueData(cmd.cmd_id,
                                              [success]))
            success = None

        if connected.value:
            if conn.inWaiting():
                rx_buffer.extend(list(conn.read(conn.inWaiting())))

        if len(rx_buffer) > stxetx.HEADER_LENGTH + stxetx.TRAILER_LENGTH:
            found_packet, next_idx = ser.process(
                                     rx_buffer[0:MAX_RECEIVE_BUFFER_LEN])
            if found_packet:
                num_packets.value += 1
                packet_type = ser.get_rx_packet_type()
                payload = ser.get_rx_packet_data()
                format = 'h' * int(len(payload)/2)
                data = struct.unpack(format, ''.join(payload))
                print 'Packet number %d  |  Type: %d  |  Data  |  %s ' % \
                    (num_packets.value, packet_type, str(data))
                rx_queue.put(QueueData(packet_type, data))
                del(rx_buffer[0:next_idx])

def make_cmd_packet(data):
    payload = []
    for datum in data:
        payload.append(struct.pack('<h', datum)[0])
        payload.append(struct.pack('<h', datum)[1])
    return payload


class MyFrame(wx.Frame):
    """ The main frame of the application
    """
    title = 'Robot control with realtime plotting'

    def __init__(self):
        wx.Frame.__init__(self, None, -1, self.title)

        self.recv_data = []
        self.emf_data = []
        self.rx_queue = Queue()
        self.tx_queue = Queue()
        self.num_packets = Value('i', 0)
        self.connected = Value('i', 0)

        self.process = Process(target=send_receive,
                               args=(self.tx_queue, self.rx_queue,
                               self.connected, self.num_packets))
        self.process.start()
        self.data_thread = threading.Thread(target = self.handle_rx_queue)
        self.data_thread.setDaemon(True)
        self.data_thread.start()
        self.Bind(EVT_UPDATE_SERIAL_STATUS, self.on_serial_status_update)
        self.paused = False
        self.rx_buffer = []

        font = wx.SystemSettings_GetFont(wx.SYS_SYSTEM_FONT)
        font.SetPointSize(10)

        self.create_menu()
        self.create_status_bar()
        self.create_main_panel()

        self.redraw_timer = wx.Timer(self)
        self.Bind(wx.EVT_TIMER, self.on_redraw_timer, self.redraw_timer)
        self.redraw_timer.Start(100)

    def create_main_panel(self):
        self.panel = wx.Panel(self)
        self.init_plot()
        self.canvas = FigCanvas(self.panel, -1, self.fig)

        self.vbox = wx.BoxSizer(wx.VERTICAL)
        self.create_plot_control_panel()
        self.create_robot_control_panel()

    def create_menu(self):
        self.menubar = wx.MenuBar()

        menu_file = wx.Menu()
        m_connect = menu_file.Append(-1, "&Connect\tCtrl-C", "Connect")
        self.Bind(wx.EVT_MENU, self.on_click_connect, m_connect)
        m_disconnect = menu_file.Append(-1, "&Disconnect\tCtrl-D", "Disconnect")
        self.Bind(wx.EVT_MENU, self.on_click_disconnect, m_disconnect)
        m_expt = menu_file.Append(-1, "&Save plot\tCtrl-S", "Save plot to file")
        self.Bind(wx.EVT_MENU, self.on_save_plot, m_expt)
        m_save_data = menu_file.Append(-1, "&Save data\tCtrl-Shift-S",
                                       "Save data to file")
        self.Bind(wx.EVT_MENU, self.on_save_data, m_save_data)
        menu_file.AppendSeparator()
        m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit")
        self.Bind(wx.EVT_MENU, self.on_exit, m_exit)

        self.menubar.Append(menu_file, "&File")
        self.SetMenuBar(self.menubar)

    def create_status_bar(self):
        self.statusbar = self.CreateStatusBar()

    def create_plot_control_panel(self):
        self.ico_connected = wx.StaticBitmap(self.panel)
        self.ico_connected.SetBitmap(wx.Bitmap('./discon.png'))
        self.cb_grid = wx.CheckBox(self.panel, -1,
            "Show Grid",
            style=wx.ALIGN_RIGHT)
        self.Bind(wx.EVT_CHECKBOX, self.on_cb_grid, self.cb_grid)
        self.cb_grid.SetValue(True)

        self.cb_xlab = wx.CheckBox(self.panel, -1,
            "Show X labels",
            style=wx.ALIGN_RIGHT)
        self.Bind(wx.EVT_CHECKBOX, self.on_cb_xlab, self.cb_xlab)
        self.cb_xlab.SetValue(True)

        self.btn_clear = wx.Button(self.panel, -1, 'Clear Data',
                                   style=wx.ALIGN_RIGHT)
        self.Bind(wx.EVT_BUTTON, self.on_btn_clear, self.btn_clear)

        self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox1.Add(self.cb_grid, border=5,
                       flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
        self.hbox1.AddSpacer(10)
        self.hbox1.Add(self.cb_xlab, border=5,
                       flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
        self.hbox1.Add(self.btn_clear, border=5,
                       flag=wx.ALL | wx.ALIGN_CENTER_VERTICAL)
        self.hbox1.AddSpacer(wx.Size(180, 10))
        self.hbox1.Add(self.ico_connected, border =5, flag=wx.ALL)

        self.vbox.Add(self.canvas, 1, flag=wx.LEFT | wx.TOP | wx.GROW)
        self.vbox.Add(self.hbox1, 0, flag=wx.ALIGN_LEFT | wx.TOP)

    def create_robot_control_panel(self):
        # Gains panel
        self.panel2 = wx.Panel(self.panel, -1)
        hbox3 = wx.BoxSizer(wx.HORIZONTAL)
        hbox4 = wx.BoxSizer(wx.HORIZONTAL)
        hbox5 = wx.BoxSizer(wx.HORIZONTAL)
        vsbox2 = wx.StaticBoxSizer(wx.StaticBox(self.panel2, -1, 'Gains'),
                                   orient=wx.VERTICAL)

        self.lblPGain = wx.StaticText(self.panel2, -1, 'P:', (0, 0))
        self.pGain = wx.TextCtrl(self.panel2, -1, size=(40, -1))
        self.lblIGain = wx.StaticText(self.panel2, -1, 'I: ', (0, 0))
        self.iGain = wx.TextCtrl(self.panel2, -1, size=(40, -1))
        self.lblDGain = wx.StaticText(self.panel2, -1, 'D: ', (0, 0))
        self.dGain = wx.TextCtrl(self.panel2, -1, size=(40, -1))

        hbox3.Add(self.lblPGain, 0, wx.EXPAND | wx.ALL, 5)
        hbox3.Add(self.pGain, 0)
        hbox3.Add((30, -1), 1)
        hbox3.Add(self.lblIGain, 0, wx.EXPAND | wx.ALL, 5)
        hbox3.Add(self.iGain, 0)
        hbox3.Add((30, -1), 1)
        hbox3.Add(self.lblDGain, 0, wx.EXPAND | wx.ALL, 5)
        hbox3.Add(self.dGain, 0)
        vsbox2.Add(hbox3, 1, wx.EXPAND | wx.ALL, border = 5)

        self.lblAWGain = wx.StaticText(self.panel2, -1, 'Anti-windup Gain: ',
                                       (0, 0))
        self.awGain = wx.TextCtrl(self.panel2, -1, size=(40, -1))

        hbox4.Add(self.lblAWGain, 0, wx.EXPAND | wx.ALL, 5)
        hbox4.Add(self.awGain, 0)
        vsbox2.Add(hbox4, 1, wx.EXPAND | wx.ALL, border = 5)

        self.btnSetGains = wx.Button(self.panel2, -1, "Set Gains", (0, 0))
        self.Bind(wx.EVT_BUTTON, self.on_click_gains, self.btnSetGains)
        self.btnSetGains.SetToolTipString("This button sets the gains" +
                                          " of the PID controller")

        hbox5.Add(self.btnSetGains, 0, wx.ALIGN_LEFT)

        vsbox2.Add(hbox5, 1, wx.EXPAND | wx.ALL, border=5)

        self.panel2.SetSizer(vsbox2)

        # Speed setting and Run panel
        self.panel3 = wx.Panel(self.panel, -1)
        hbox6 = wx.BoxSizer(wx.HORIZONTAL)
        vsbox3 = wx.StaticBoxSizer(
                    wx.StaticBox(self.panel3, -1, 'Run Settings'),
                    orient=wx.VERTICAL)

        self.lblStrideFreq = wx.StaticText(self.panel3, -1, 'Stride Freq: ',
                                           (0, 0))
        self.strideFreq = wx.TextCtrl(self.panel3, -1, size=(40, -1))
        self.lblRuntime = wx.StaticText(self.panel3, -1, 'Time: ', (0, 0))
        self.runTime = wx.TextCtrl(self.panel3, -1, size=(40, -1))

        hbox6.Add(self.lblStrideFreq, 0, wx.EXPAND | wx.ALL, 5)
        hbox6.Add(self.strideFreq, 0)
        hbox6.Add(self.lblRuntime, 0, wx.EXPAND | wx.ALL, 5)
        hbox6.Add(self.runTime, 0)
        vsbox3.Add(hbox6)

        hbox7 = wx.BoxSizer(wx.HORIZONTAL)

        self.btnRun = wx.Button(self.panel3, -1, 'Go!', (0, 0))
        self.Bind(wx.EVT_BUTTON, self.on_click_run, self.btnRun)
        self.btnRun.SetToolTipString('Run the robot at the' +
                                     ' frequency entered in the box')
        self.cbClosedLoop = wx.CheckBox(self.panel3, -1, "Closed Loop",
                                        style=wx.ALIGN_RIGHT)
        self.cbClosedLoop.SetValue(False)

        hbox7.Add(self.btnRun)
        hbox7.AddSpacer(wx.Size(30, 10))
        hbox7.Add(self.cbClosedLoop)
        vsbox3.Add(hbox7, 1, wx.EXPAND | wx.ALL, border=5)

        self.panel3.SetSizer(vsbox3)

        self.hbox8 = wx.BoxSizer(wx.HORIZONTAL)

        # Put it all together into the vbox holding all panels
        self.hbox8.Add(self.panel2, 0, wx.EXPAND | wx.ALL, 5)
        self.hbox8.Add(self.panel3, 0, wx.EXPAND | wx.ALL, 5)

        self.vbox.Add(self.hbox8, 0, flag=wx.ALIGN_CENTER | wx.TOP)

        self.panel.SetSizer(self.vbox)
        self.vbox.Fit(self)
        self.enable_controls(self.connected.value)

    def on_click_connect(self, event):
        dlg = wx.TextEntryDialog(self,
        'Enter the serial port name for the Bluetooth device to connect to.',
        'Bluetooth Serial Port Name', 'RoachCtrl1-SPP-1')

        dlg.SetValue('/dev/tty.RoachCtrl2-SPP-1')

        if dlg.ShowModal() == wx.ID_OK:
            port_name = dlg.GetValue()
            self.tx_queue.put_nowait(QueueData(CONNECT, [port_name]))
        else:
            pass
        dlg.Destroy()

    def on_click_disconnect(self, event):
        self.tx_queue.put_nowait(QueueData(DISCONNECT, []))

    def on_click_gains(self, event):
        gains = []
        try:
            gains.append(int(100 * float(self.pGain.GetValue())))
            gains.append(int(100 * float(self.iGain.GetValue())))
            gains.append(int(100 * float(self.dGain.GetValue())))
            gains.append(int(100 * float(self.awGain.GetValue())))
        except ValueError:
            self.pGain.SetValue('')
            self.iGain.SetValue('')
            self.dGain.SetValue('')
            self.awGain.SetValue('')
            self.create_modal_dialog('At least one gain value was not' +
                                     ' recognized as a decimal value.')
            return

        packet = make_cmd_packet(gains)
        self.tx_queue.put_nowait(QueueData(CONFIG_PID, packet))

    def on_click_run(self, event):
        try:
            time = int(self.runTime.GetValue())
        except ValueError:
            self.create_modal_dialog('Run time must be an integer' +
                                     ' value from 0 to 255')
            return
        try:
            closed_loop = int(self.cbClosedLoop.IsChecked())
            ref_emf = cvtStrideFreqToEMF(float(self.strideFreq.GetValue()))
            packet = make_cmd_packet([ref_emf, time, closed_loop])
            self.tx_queue.put_nowait(QueueData(SET_MOTOR_SPEED, packet))
        except ValueError:
            # Okay, the following error is somewhat inaccurate. 
            #We need to be able to allow the robot to back up in the future.
            self.create_modal_dialog('Stride frequency must be a ' +
                                     'number between 0 and 40')
            return

    def on_serial_status_update(self, event):
        status = event.success
        if event.cmd_id == CONNECT:
            if not status:
                self.create_modal_dialog('Could not connect.' +
                                         ' Is your Bluetooth enabled?')
                self.ico_connected.SetBitmap(wx.Bitmap('./discon.png'))
            else:
                self.ico_connected.SetBitmap(wx.Bitmap('./con.png'))
            self.enable_controls(status)
            self.flash_status_message("Connection success status: " +
                                      str(status))
        elif event.cmd_id == DISCONNECT:
            self.enable_controls(not status)
            if status:
                self.ico_connected.SetBitmap(wx.Bitmap('./discon.png'))
            self.flash_status_message("Disconnection success status: " +
                                      str(status))
        elif event.cmd_id == CONFIG_PID:
            self.flash_status_message("Gain setting success status: " +
                                      str(status))
            if not status:
                self.create_modal_dialog('Controller gains could not be' +
                                         ' successfully set. You may need' +
                                         ' to reset the connection.')
        else:
            print "Unknown command ID: " + str(event.cmd_id)

    def enable_controls(self, enabled):
        if enabled:
            txt_color = wx.Color(0, 0, 0)
        else:
            txt_color = wx.Color(128, 128, 128)

        self.pGain.SetEditable(enabled)
        self.lblPGain.SetForegroundColour(txt_color)
        self.iGain.SetEditable(enabled)
        self.lblIGain.SetForegroundColour(txt_color)
        self.dGain.SetEditable(enabled)
        self.lblDGain.SetForegroundColour(txt_color)
        self.awGain.SetEditable(enabled)
        self.lblAWGain.SetForegroundColour(txt_color)
        self.btnSetGains.Enable(enabled)

        self.runTime.SetEditable(enabled)
        self.lblRuntime.SetForegroundColour(txt_color)
        self.strideFreq.SetEditable(enabled)
        self.lblStrideFreq.SetForegroundColour(txt_color)
        self.cbClosedLoop.Enable(enabled)
        self.btnRun.Enable(enabled)


    def handle_rx_queue(self):
        while True:
            try:
                data = self.rx_queue.get()
                if data.cmd_id != SENSE_DATA and \
                   data.cmd_id != SET_MOTOR_SPEED:
                    wx.PostEvent(self, SerialStatusEvent(
                                cmd_id = data.cmd_id,
                                success = data.payload[0]))
                elif data.cmd_id == SENSE_DATA:
                    self.recv_data.append(data.payload)
                    self.emf_data.append(data.payload[0])
            except Empty:
                pass

    def create_modal_dialog(self, msg):
        dlg = wx.MessageDialog(self, msg, 'Error Encountered', wx.OK |\
        wx.ICON_INFORMATION)
        dlg.ShowModal()
        dlg.Destroy()

    def init_plot(self):
        """Initialize the plot for the realtime data."""
        self.dpi = 100
        self.fig = Figure((3.0, 3.0), dpi=self.dpi)

        self.axes = self.fig.add_subplot(111)
        self.axes.set_axis_bgcolor('black')
        self.axes.set_title('Realtime Robot Data', size=12)

        pylab.setp(self.axes.get_xticklabels(), fontsize=8)
        pylab.setp(self.axes.get_yticklabels(), fontsize=8)

        # plot the data as a line series, and save the reference 
        # to the plotted line series
        #

        self.plot_data = self.axes.plot(
            self.emf_data,
            linewidth=1,
            color=(1, 1, 0),
            )[0]

    def draw_plot(self):
        """ Redraws the plot
        """
        # when xmin is on auto, it "follows" xmax to produce a 
        # sliding window effect. therefore, xmin is assigned after
        # xmax.
        #

        xmax = len(self.emf_data) if len(self.emf_data) > 5000 else 5000
        xmin = xmax - 5000

        # for ymin and ymax, find the minimal and maximal values
        # in the data set and add a mininal margin.
        # 
        # note that it's easy to change this scheme to the 
        # minimal/maximal value in the current display, and not
        # the whole data set.
        # 

        self.axes.set_xbound(lower=xmin, upper=xmax)
        self.axes.set_ybound(lower=Y_MIN, upper=Y_MAX)

        # anecdote: axes.grid assumes b=True if any other flag is
        # given even if b is set to False.
        # so just passing the flag into the first statement won't
        # work.
        #
        if self.cb_grid.IsChecked():
            self.axes.grid(True, color='gray')
        else:
            self.axes.grid(False)

        #Using setp here is convenient, because get_xticklabels
        #returns a list over which one needs to explicitly 
        #iterate, and setp already handles this.

        pylab.setp(self.axes.get_xticklabels(),
            visible=self.cb_xlab.IsChecked())

        self.plot_data.set_xdata(arange(len(self.emf_data)))
        self.plot_data.set_ydata(array(self.emf_data))
        tic = time.time()
        try:
            self.canvas.draw()
        except RuntimeError, err:
            # Something weird happening here. Data lengths appear to be the 
            # same,but we're getting a runtime error complaining 
            # they are different
            print err.message

    def on_cb_grid(self, event):
        self.draw_plot()

    def on_cb_xlab(self, event):
        self.draw_plot()

    def on_btn_clear(self, event):
        del self.recv_data[:]
        del self.emf_data[:]

    def on_save_plot(self, event):
        file_choices = "PNG (*.png)|*.png"

        dlg = wx.FileDialog(
            self,
            message="Save plot as...",
            defaultDir=os.getcwd(),
            defaultFile="plot.png",
            wildcard=file_choices,
            style=wx.SAVE)

        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            self.canvas.print_figure(path, dpi=self.dpi)
            self.flash_status_message("Plot saved to %s" % path)

    def on_save_data(self, event):
        file_choices = "CSV (*.csv)|*.csv"

        dlg = wx.FileDialog(
            self,
            message="Save data as...",
            defaultFile = "EMF_Data.csv",
            wildcard=file_choices,
            style=wx.SAVE)

        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            #datafile = open(path, 'w')
            savetxt(path, array(self.recv_data), '%4.3f', delimiter=',')
            self.flash_status_message("Data saved to %s" % path)

    def on_redraw_timer(self, event):
        self.draw_plot()

    def on_exit(self, event):
        print "Found a total of %d packets" % self.num_packets.value
        self.process.terminate()
        self.Destroy()
        sys.exit(0)

    def flash_status_message(self, msg, flash_len_ms=1500):
        self.statusbar.SetStatusText(msg)
        self.timeroff = wx.Timer(self)
        self.Bind(
            wx.EVT_TIMER,
            self.on_flash_status_off,
            self.timeroff)
        self.timeroff.Start(flash_len_ms, oneShot=True)

    def on_flash_status_off(self, event):
        self.statusbar.SetStatusText('')


if __name__ == '__main__':
    app = wx.PySimpleApp()
    app.frame = MyFrame()
    app.frame.Show()
    app.MainLoop()