Python script for processing songs with aubio, then playing them back with pygame and controlling LEDs (with smithje/RGB_LED_Driver) in time with the music -- all on a Raspberry Pi

song2rgb.py

#!/usr/bin/python

import pygame.mixer as mixer
from pygame.mixer import music
#from led_driver import set_channel_value
from ledlibs import RGB_Driver

from colorsys import hsv_to_rgb
import cPickle
import argparse
from collections import deque
import subprocess as sp
import time

TOL = 0.1 # 100 ms
aubiopitch_cmd = lambda x: ['aubiopitch', '-u', 'midi', '-i', '%s' % x]
aubioonset_cmd = lambda x: ['aubioonset', '-i', '%s' % x]
channels = [(15,1,2),(3,4,5),(6,7,8),(9,10,11),(12,13,14)]

led_drivers = []

def aubiopitch(song):
    out = sp.check_output(aubiopitch_cmd(song))
    pitches = deque([(float(y), float(z)) for y,z in [x.split('\t') for x in out.splitlines()]])
    return pitches

def aubioonset(song):
    out = sp.check_output(aubioonset_cmd(song))
    onsets = deque([float(y) for y in out.splitlines()])
    return onsets

def playsong(pitches, onsets):
    music.play()
    start_time = time.time() 

    while len(pitches) > 0 or len(onsets) > 0:
        #pos = music.get_pos() / 1000.
        pos = time.time() - start_time
        pitch = None
        onset = False
        # if we've fallen behind the song, throw out pitches/onsets we missed
        while len(pitches) > 0 and pitches[0][0] < pos:
            pitches.popleft()
        while len(onsets) > 0 and onsets[0] < pos:
            onsets.popleft()

        # collect pitches/onsets happening within TOL (100ms) of now in the song
        while len(pitches) > 0 and pitches[0][1] != -1.0 and abs(pitches[0][0] - pos) < TOL:
            pitch = pitches.popleft()
        while len(onsets) > 0 and abs(onsets[0] - pos) < TOL:
            onsets.popleft()
            onset = True

        # TODO: doesn't use an onset detected if a pitch wasn't also detected
        if pitch:
            # TODO: only using the first pitch detected at a certain time
            rgb = pitch2rgb(pitch[1], onset)
            #print pos, len(pitches), len(onsets), '-', pitches_popped, onsets_popped, '-', rgb
            write_rgb(rgb)
        else:
            write_rgb((0., 0., 0.))

    write_rgb((0.,0.,0.))

def write_rgb(rgb):
    rgb = [int(x * 4095) for x in rgb]
    for led in led_drivers:
        led.set_rgb(rgb[0], rgb[1], rgb[2])
    return
    #used with pi-blaster
    set_channel_value(0, rgb[0])
    set_channel_value(1, rgb[1])
    set_channel_value(2, rgb[2])
    set_channel_value(4, rgb[1])
    set_channel_value(6, rgb[0])
    set_channel_value(7, rgb[2])
    #time.sleep(0.3)

def pitch2rgb(pitch, onset=False):
    rgb = hsv_to_rgb(*pitch2hsv(pitch, onset))
    return rgb

def pitch2hsv(pitch, onset=False):
    # map pitch [20,100] to hue [0.0,1.0]
    hue = (pitch - 20.0) / 80
    sat = 1.0
    val = 1.0 if onset else 0.5
    return hue, sat, val


def main():
    parser = argparse.ArgumentParser(description='Outputs RGB values during song playback.')
    parser.add_argument('song', type=str, help='input song file')
    parser.add_argument('-p', '--pitches', dest='fout_pitches',
                        type=argparse.FileType('w'), help='file to output aubiopitch results to')
    parser.add_argument('-o', '--onsets', dest='fout_onsets',
                        type=argparse.FileType('w'), help='file to output aubioonset results to')
    parser.add_argument('-P', '--in-pitches', dest='fin_pitches',
                        type=argparse.FileType('r'), help='file to read aubiopitch values from')
    parser.add_argument('-O', '--in-onsets', dest='fin_onsets',
                        type=argparse.FileType('r'), help='file to read aubioonset values from')
    parser.add_argument('-C', '--led-count', dest='led_count', type=int, default=1,
                        help='number of leds to output to (from 1 to 5)', choices=range(1,6))
    args = parser.parse_args()

    if args.fout_pitches:
        pitches = aubiopitch(args.song)
        cPickle.dump(pitches, args.fout_pitches)
    if args.fout_onsets:
        onsets = aubioonset(args.song)
        cPickle.dump(onsets, args.fout_onsets)
    if args.fout_pitches or args.fout_onsets:
        return

    pitches = cPickle.load(args.fin_pitches) if args.fin_pitches else []
    onsets = cPickle.load(args.fin_onsets) if args.fin_onsets else []

    global led_drivers
    led_drivers = [RGB_Driver.RGB_Driver(None, channels[i][0], channels[i][1], channels[i][2]) for i in range(args.led_count)]
    
    #mixer.pre_init(22050, -16, 2, 8192)
    mixer.pre_init(44100, -16, 2, 16777216)
    mixer.init()
    music.load(args.song)
    playsong(pitches, onsets)

if __name__ == '__main__':
    main()