Reading Wii Nunchuck data in Python

Brief description: Use the Wii Nunchuck to draw in the terminal and turn the drawing into a SVG, or in a broader sense: read values sensed by the Arduino into Python

Since I installed Linux I use the terminal a lot. I really like how you can use something text based and really basic to perform all kinds of complex tasks.

After messing with SVG for a bit I thought it might be interesting to use the text based terminal as the interface to a SVG generating script. Also, I wanted to see how I could use the terminal together with an Arduino, so I used one to connect my Python script too a Wii controller (Nunchuck), which is used as input for the script.

The Nunchuck's joystick controls the cursor, the Z button lets jou draw lines (hold down the button and then move around) and the c button creates and shows the SVG. Using the WiiChuck class the arduino reads the data which is then send to Python though a serial connection. To connect the Nunchuck to the Arduino without cutting the cable I used the Wiichuck adapter. Both the class and the adapter are available here.

[Disclaimer] Since I wrote the code from scratch and added 'features' as I went along I bet it's far from optimized. But it seems to do the job.

Arduino code

Using the WiiChuck adapter the Nunchuck is connected to analog pins 2, 3, 4 and 5. Pins 4 and 5 are the I2C communication pins and Pins 2 and 3 are used to supply power. You could also connect the Nunchuck by cutting the cables, but then you have to figure out which cable does what.

This script gets several values from the Nunchuck: the x and y acceleration (e.g. if the controller is physically rotated), the state of the Z and C buttons (wether they are up or down) and the x and y position of the joystick, where ±128 should be the center position, but these values tend to be slightly different per device.

   #include <Wire.h>
   #include "nunchuck_funcs.h" 

   int loop_cnt=0;

   byte accx,accy,zbut,cbut,joyx,joyy;
   int ledPin = 13;

   void setup() {
       Serial.begin(19200); //baudrate of the serial connection, be sure to set the same number in the Python script
       nunchuck_setpowerpins(); // initiates pin 2 and 3 so they can power the nunchuck
       nunchuck_init(); // send the initilization handshake
   }

   void loop() {
       if(loop_cnt > 100 ) { // every 100 msecs get new data
           loop_cnt = 0;

           nunchuck_get_data(); // receive the data from the device
           //nunchuck_print_data(); //if you want to show the data in the Serial monitor

           accx  = nunchuck_accelx(); // ranges from approx 70 - 182
           accy  = nunchuck_accely(); // ranges from approx 65 - 173
           zbut = nunchuck_zbutton(); // either 0 or 1
           cbut = nunchuck_cbutton();  // dito
           joyx = nunchuck_joyx();
           joyy = nunchuck_joyy(); 
   
           //print each of the values as a decimal value, followed by a comma
           Serial.print((byte)accx,DEC);
           Serial.print(",");
           Serial.print((byte)accy,DEC);
           Serial.print(",");
           Serial.print((byte)zbut,DEC);
           Serial.print(",");
           Serial.print((byte)cbut,DEC);
           Serial.print(",");
           Serial.print((byte)joyx,DEC);
           Serial.print(",");
           Serial.print((byte)joyy,DEC);

           // when all values have been printed, print a newline to end the transmission
           Serial.println();
       }
       loop_cnt++;
       delay(1);
   }

Python code

#=========================================================
# Import libs and set basic vars
#=========================================================
import serial, os, math

width, height = 50, 20 # drawing area size
xpos, ypos = 0, 0 # cursor location
sx, sy = -1, -1 # stores the location to draw a line from
zbut, prevZbut, cbut = 0, 0, 0 # wether the buttons are pressed
cursorChar, prevCursorChar = "o", "" # the character used for the cursor
svgScale = 10.0 # the svg is svgScale-times bigger than the terminal
svgHeightCorrection = 2 # compensate for the line-height of the text

#theScren will hold all of the 'pixels'
theScreen = [] 
# fullRow is the line of #s for use at the top and bottom
fullRow = "" 
for x in range(0, width+2):
    fullRow += "#"

#initiate the screen arrays, fill them with spaces
for y in range(0, height):
    xlist = []
    for x in range(0, width):
        xlist.append(" ")
    theScreen.append(xlist)

#variable to hold the svg file
svgOut = ""

#=========================================================
# function to draw the arrays to the screen
#=========================================================
def display():
    global theScreen
    global cursorChar
    global xpos
    global ypos

    prevCursorChar = theScreen[ypos][xpos] 
    theScreen[ypos][xpos] = cursorChar
    os.system("clear")
    print fullRow
    for y in theScreen:
        printLine = "#"
        for x in y:
            printLine += str(x)
        print printLine+"#"
    print fullRow
    theScreen[ypos][xpos] = prevCursorChar

#=========================================================
# function to draw a line given a start and end xy position
# useInSvg determines of the line should be saved in the svg
# the line that is drawn while still holding down the button
# will not be saved (to prevent filling the whole screen when
# moving around )
#=========================================================
def drawLine(x1, y1, x2, y2, useInSvg):
    global theScreen
    global svgOut
    global svgScale
    global svgHeightCorrection

    bx1, by1, bx2, by2 = int(x1), int(y1), int(x2), int(y2) #remember the original values

    x1 = max(0, min(x1, width))
    x2 = max(0, min(x2, width))
    y1 = max(0, min(y1, height))
    y2 = max(0, min(y2, height))
        
    #d = direction
    d = math.atan2(y2-y1, x2 - x1)*180/math.pi 
    d2 = d
    
    # calculate a different character for each direction the line can go
    if d2 < 0:
        d2 = (360-abs(d2))        
    dist = math.floor(math.sqrt((x1-x2)*(x1-x2) + (y1-y2)*(y1-y2)))
    lineChar = "*"
    if (d2 >= 360-22.53 or d2 < 22.5) or (d2>=157.5 and d2<202.5):
       lineChar = "~"
    elif (d2 >= 22.5 and d2 < 67.5) or (d2>=202.5 and d2<247.5):
       lineChar = "\\"
    elif (d2>=67.5 and d2<112.5) or (d2>=247.5 and d2<292.5):
       lineChar = "|"
    elif (d2>=112.5 and d2<157.5) or (d2>=292.5 and d2<337.5):
       lineChar = "/"
    
    step = 0
    while(step < dist):
        x1 += math.cos(d*math.pi/180)
        y1 += math.sin(d*math.pi/180)
        if x1>width:
            x1 = width-1
            x2 = width-1
        elif x1<0:
            x1 = 0
            x2 = 0
        if y1>height:
            y1 = height-1
            y2 = height-1
        elif y2<0:
            y1 = 0
            y2 = 0
        drawY = int(math.floor(y1))
        drawX = int(math.floor(x1))
        if drawY < len(theScreen):
            if drawX < len(theScreen[drawY]):
                theScreen[drawY][drawX] = lineChar
        step = step + 1
    theScreen[by1][bx1], theScreen[by2][bx2] = "*", "*"

    if useInSvg==1:
        svgOut = svgOut + """<line x1="{0}" y1="{1}" x2="{2}" y2="{3}"
    style="stroke:rgb(0,0,0);stroke-width:{4}"/>""".format(bx1*svgScale, by1*(svgScale*svgHeightCorrection), bx2*svgScale, by2*(svgScale*svgHeightCorrection), svgScale)


#=========================================================
# Save svg image and show it in google chrome 
#=========================================================
def saveOpenAndQuit():
    global svgOut
    global width
    global height
    global svgScale
    svgOut = """<svg xmlns="http://www.w3.org/2000/svg" version="1.1" width="{0}" height="{1}">""".format(svgScale*width, svgScale*(height*svgHeightCorrection)) + svgOut
    svgOut = svgOut + "</svg>"

    fo = open("image.svg", "wb")
    fo.write( svgOut);
    fo.close()
    os.system("google-chrome image.svg &")

#=========================================================
# Main loop
# Read data from Arduino, update cursor position and such, display screen
#=========================================================
# open a serial connection to the Arduino. 
# Make sure the last number, the baudrate, is the same as in your Arduino code
ser = serial.Serial('/dev/ttyACM0', 19200)

# keep looping forever
while 1:
    # read the serial data coming in and split it on the comma so you'll have a list you can work with.
    serialData = ser.readline()
    serialData = serialData.replace("\r\n", "")
    dataList = serialData.split(',')

    # only go on if there are 6 elements
    # this prevents errors when somehow the script didn't receive the full message
    if len(dataList)==6: 
        prevZbut = zbut
        accx = float(dataList[0])
        accy = float(dataList[1])
        zbut = float(dataList[2])
        cbut = float(dataList[3])
        joyx = float(dataList[4])
        joyy = float(dataList[5])
        
        #move cursor around
        # using 150/100 leaves some room for weird offsets that the controller might have
        if joyx>150: 
            xpos = xpos + 1
        elif joyx<100:
            xpos = xpos - 1

        if joyy>150:
            ypos = ypos - 1
        elif joyy<100:
            ypos = ypos + 1

        # limit the cursor position to positions in the drawing area
        xpos = max(0, min(xpos, width-1)) 
        ypos = max(0, min(ypos, height-1))

        #remember previous char on cursor location
        prevCursorChar = theScreen[ypos][xpos]
        
        if cbut==1:
            saveOpenAndQuit()
        cursorChar = "o"

        if zbut==1 and prevZbut==0:#click, starting to draw
            sx, sy = xpos, ypos
            theScreen[ypos][xpos] = "*"
            cursorChar = "x"
            display()
        elif zbut==1 and prevZbut==1:#dragging around
            cursorChar = "*"
            #somehow I need to iterate over the tempScreen and theScreen arrays
            #instead of just copying them with .extend
            tempScreen = []
            for y in range(0, height):
                xlist = []
                for x in range(0, width):
                    xlist.append(theScreen[y][x])
                tempScreen.append(xlist)

            drawLine(sx, sy, xpos, ypos, 0)
            display()
            theScreen = []
            for y in range(0, height):
                xlist = []
                for x in range(0, width):
                    xlist.append(tempScreen[y][x])
                theScreen.append(xlist)

        elif zbut==0 and prevZbut==1:#released
            theScreen[ypos][xpos] = "*"
            cursorChar = "x"
            if sx!=xpos or sy!=ypos:
                drawLine(sx, sy, xpos, ypos, 1)
            sx, sy = -1, -1
            display()

        elif zbut==0 and prevZbut==0:#doing nuthin'
            cursorChar = "o"
            display()