← Back to branch summary

~ubuntu-branches/ubuntu/utopic/python-chaco/utopic

« back to all changes in this revision

Viewing changes to examples/demo/vtk/spectrum.py

  • Committer: Package Import Robot
  • Author(s): Andrew Starr-Bochicchio
  • Date: 2014-06-01 17:04:08 UTC
  • mfrom: (7.2.5 sid)
  • Revision ID: package-import@ubuntu.com-20140601170408-m86xvdjd83a4qon0
Tags: 4.4.1-1ubuntu1
* Merge from Debian unstable. Remaining Ubuntu changes:
 - Let the binary-predeb target work on the usr/lib/python* directory
   as we don't have usr/share/pyshared anymore.

Show diffs side-by-side

added added

removed removed

Lines of Context:
examples/demo/vtk/spectrum.py
1
 
#!/usr/bin/env python
2
 
"""
3
 
This plot displays the audio spectrum from the microphone.
4
 
 
5
 
Based on updating_plot.py
6
 
"""
7
 
 
8
 
# Major library imports
9
 
import pyaudio
10
 
from numpy import zeros, linspace, short, fromstring, hstack, transpose
11
 
from scipy import fft
12
 
 
13
 
# Enthought library imports
14
 
from chaco.default_colormaps import jet
15
 
from traits.api import HasTraits, Instance, Any
16
 
 
17
 
# Chaco imports
18
 
from chaco.api import Plot, ArrayPlotData, HPlotContainer, \
19
 
        OverlayPlotContainer
20
 
from chaco.tools.api import MoveTool, PanTool, ZoomTool
21
 
 
22
 
NUM_SAMPLES = 1024
23
 
SAMPLING_RATE = 11025
24
 
SPECTROGRAM_LENGTH = 100
25
 
 
26
 
def create_plot_component(obj):
27
 
    # Setup the spectrum plot
28
 
    frequencies = linspace(0., float(SAMPLING_RATE)/2, num=NUM_SAMPLES/2)
29
 
    obj.spectrum_data = ArrayPlotData(frequency=frequencies)
30
 
    empty_amplitude = zeros(NUM_SAMPLES/2)
31
 
    obj.spectrum_data.set_data('amplitude', empty_amplitude)
32
 
 
33
 
    obj.spectrum_plot = Plot(obj.spectrum_data)
34
 
    obj.spectrum_plot.plot(("frequency", "amplitude"), name="Spectrum",
35
 
                           color="red")
36
 
    obj.spectrum_plot.padding = 50
37
 
    obj.spectrum_plot.title = "Spectrum"
38
 
    spec_range = obj.spectrum_plot.plots.values()[0][0].value_mapper.range
39
 
    spec_range.low = 0.0
40
 
    spec_range.high = 5.0
41
 
    obj.spectrum_plot.index_axis.title = 'Frequency (hz)'
42
 
    obj.spectrum_plot.value_axis.title = 'Amplitude'
43
 
 
44
 
    # Time Series plot
45
 
    times = linspace(0., float(NUM_SAMPLES)/SAMPLING_RATE, num=NUM_SAMPLES)
46
 
    obj.time_data = ArrayPlotData(time=times)
47
 
    empty_amplitude = zeros(NUM_SAMPLES)
48
 
    obj.time_data.set_data('amplitude', empty_amplitude)
49
 
 
50
 
    obj.time_plot = Plot(obj.time_data)
51
 
    obj.time_plot.plot(("time", "amplitude"), name="Time", color="blue")
52
 
    obj.time_plot.padding = 50
53
 
    obj.time_plot.title = "Time"
54
 
    obj.time_plot.index_axis.title = 'Time (seconds)'
55
 
    obj.time_plot.value_axis.title = 'Amplitude'
56
 
    time_range = obj.time_plot.plots.values()[0][0].value_mapper.range
57
 
    time_range.low = -0.2
58
 
    time_range.high = 0.2
59
 
 
60
 
    # Spectrogram plot
61
 
    spectrogram_data = zeros(( NUM_SAMPLES/2, SPECTROGRAM_LENGTH))
62
 
    obj.spectrogram_plotdata = ArrayPlotData()
63
 
    obj.spectrogram_plotdata.set_data('imagedata', spectrogram_data)
64
 
    spectrogram_plot = Plot(obj.spectrogram_plotdata)
65
 
    spectrogram_time = linspace(
66
 
        0.0, float(SPECTROGRAM_LENGTH*NUM_SAMPLES)/float(SAMPLING_RATE),
67
 
        num=SPECTROGRAM_LENGTH)
68
 
    spectrogram_freq = linspace(0.0, float(SAMPLING_RATE/2), num=NUM_SAMPLES/2)
69
 
    spectrogram_plot.img_plot('imagedata',
70
 
                              name='Spectrogram',
71
 
                              xbounds=spectrogram_time,
72
 
                              ybounds=spectrogram_freq,
73
 
                              colormap=jet,
74
 
                              )
75
 
    range_obj = spectrogram_plot.plots['Spectrogram'][0].value_mapper.range
76
 
    range_obj.high = 5
77
 
    range_obj.low = 0.0
78
 
    spectrogram_plot.title = 'Spectrogram'
79
 
    obj.spectrogram_plot = spectrogram_plot
80
 
 
81
 
    return obj.spectrum_plot, obj.time_plot, obj.spectrogram_plot
82
 
 
83
 
_stream = None
84
 
def get_audio_data():
85
 
    global _stream
86
 
    if _stream is None:
87
 
        # The audio stream is opened the first time this function gets called.
88
 
        # The stream is always closed (if it was opened) in a try finally
89
 
        # block at the end of this file,
90
 
        pa = pyaudio.PyAudio()
91
 
        _stream = pa.open(format=pyaudio.paInt16, channels=1,
92
 
                          rate=SAMPLING_RATE,
93
 
                          input=True, frames_per_buffer=NUM_SAMPLES)
94
 
 
95
 
    audio_data  = fromstring(_stream.read(NUM_SAMPLES), dtype=short)
96
 
    normalized_data = audio_data / 32768.0
97
 
    return (abs(fft(normalized_data))[:NUM_SAMPLES/2], normalized_data)
98
 
 
99
 
class TimerController(HasTraits):
100
 
    interactor = Any()
101
 
    timer_id = Any()
102
 
 
103
 
    def on_timer(self, vtk_obj=None, eventname=""):
104
 
        try:
105
 
            spectrum, time = get_audio_data()
106
 
        except IOError:
107
 
            return
108
 
        self.spectrum_data.set_data('amplitude', spectrum)
109
 
        self.time_data.set_data('amplitude', time)
110
 
        spectrogram_data = self.spectrogram_plotdata.get_data('imagedata')
111
 
        spectrogram_data = hstack((spectrogram_data[:,1:],
112
 
                                   transpose([spectrum])))
113
 
 
114
 
        self.spectrogram_plotdata.set_data('imagedata', spectrogram_data)
115
 
        self.spectrum_plot.request_redraw()
116
 
        return
117
 
 
118
 
def main():
119
 
    from tvtk.api import tvtk
120
 
    from mayavi import mlab
121
 
    from enable.vtk_backend.vtk_window import EnableVTKWindow
122
 
    f = mlab.figure(size=(900,850))
123
 
    m = mlab.test_mesh()
124
 
    scene = mlab.gcf().scene
125
 
    render_window = scene.render_window
126
 
    renderer = scene.renderer
127
 
    rwi = scene.interactor
128
 
 
129
 
    # Create the plot
130
 
    timer_controller = TimerController()
131
 
    plots = create_plot_component(timer_controller)
132
 
    specplot, timeplot, spectrogram = plots
133
 
 
134
 
    for i, p in enumerate(plots):
135
 
        p.set(resizable = "", bounds = [200,200], outer_x = 0,
136
 
                bgcolor = "transparent",
137
 
                )
138
 
        p.outer_y = i*250
139
 
        p.tools.append(MoveTool(p, drag_button="right"))
140
 
        p.tools.append(PanTool(p))
141
 
        p.tools.append(ZoomTool(p))
142
 
 
143
 
    spectrogram.tools[-1].set(tool_mode="range", axis="value")
144
 
    spectrogram.tools[-2].set(constrain=True, constrain_direction="y")
145
 
 
146
 
    container = OverlayPlotContainer(bgcolor = "transparent",
147
 
                    fit_window = True)
148
 
    container.add(*plots)
149
 
    container.timer_callback = timer_controller.on_timer
150
 
 
151
 
    window = EnableVTKWindow(rwi, renderer,
152
 
            component = container,
153
 
            istyle_class = tvtk.InteractorStyleTrackballCamera,
154
 
            bgcolor = "transparent",
155
 
            event_passthrough = True,
156
 
            )
157
 
 
158
 
    mlab.show()
159
 
 
160
 
if __name__ == "__main__":
161
 
    main()
162