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()