wavelet decomposition viewer with PyWavelets & PyQtGraph

dwtui.py

"""
Simple wavelet decomposition viewer components

"""

import time
import numpy as np
import pywt
import pyqtgraph as pg

def wtrect(n,h):
    rect = pg.QtCore.QRect()
    rect.setTop(h)
    rect.setBottom(h)
    rect.setLeft(0)
    rect.setRight(n)
    return rect

class Signal(pg.PlotItem):
    def __init__(self, fs, data, parent=None):
        pg.PlotItem.__init__(self, parent)
        self.time = np.r_[:data.size]*1.0/fs
        self.data = data
        self.pdi = pg.PlotDataItem(x=self.time, y=self.data)
        self.addItem(self.pdi)
        self.setLabel('bottom', 'time')

class WaveDec(pg.PlotItem):
    def __init__(self, fs, data, parent=None):
        pg.PlotItem.__init__(self, parent)
        self.endtime = len(data)*1.0/fs
        self.data = data
        self.compute_wavedec()
        self.ticks = []
        for i, cdi in enumerate(self.cD):
            im = pg.ImageItem(cdi[:, np.newaxis])
            im.setRect(wtrect(self.endtime, i))
            self.addItem(im)
            self.ticks.append((i+0.5, str(fs/(2**(len(self.cD)-i)))))
        self.getAxis('left').setTicks([self.ticks])

    def compute_wavedec(self):
        self.wavelet = pywt.Wavelet('sym20')
        tic = time.time()
        cs = pywt.wavedec(self.data, self.wavelet)
        print 'wavedec required %.2f s' % (time.time() - tic, )
        self.cA, self.cD = cs[0], cs[1:]


class WaveView(pg.GraphicsWindow):

    def __init__(self, fs, signal, parent=None):
        pg.GraphicsWindow.__init__(self, parent)
        self.fs = fs
        self.signal = signal
        self.p_ts = Signal(fs, signal)
        self.addItem(self.p_ts)
        self.nextRow()
        self.p_wd = WaveDec(fs, signal)
        self.p_wd.setXLink(self.p_ts)
        self.addItem(self.p_wd)
        self.show()

if __name__ == '__main__':

    app = pg.QtGui.QApplication([])

    fs = 1024.0
    n = 10000
    y = np.cumsum(np.random.randn(n)) + 10*np.sin(np.r_[:n]/40.)

    wv = WaveView(fs, y)
    wv.show()

    app.exec_()

Hi! Is it possible to put color to the spectogram?