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

« back to all changes in this revision

Viewing changes to examples/advanced/spec_waterfall.py

  • Committer: Bazaar Package Importer
  • Author(s): Varun Hiremath
  • Date: 2011-04-06 19:03:54 UTC
  • mfrom: (7.2.2 sid)
  • Revision ID: james.westby@ubuntu.com-20110406190354-rwd55l2ezjecfo41
Tags: 3.4.0-2
d/rules: fix pyshared directory path (Closes: #621116)

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
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
# Major library imports
 
8
import pyaudio
 
9
from numpy import zeros, linspace, short, fromstring, hstack, transpose, array
 
10
from scipy import fft
 
11
 
 
12
# Enthought library imports
 
13
from enthought.chaco.default_colormaps import jet
 
14
from enthought.enable.api import Window, Component, ComponentEditor
 
15
from enthought.traits.api import HasTraits, Instance, List, Range
 
16
from enthought.traits.ui.api import Item, Group, View, Handler
 
17
from enthought.enable.example_support import DemoFrame, demo_main
 
18
from enthought.pyface.timer.api import Timer
 
19
 
 
20
# Chaco imports
 
21
from enthought.chaco.api import (Plot, ArrayPlotData, HPlotContainer, VPlotContainer,
 
22
    AbstractMapper, LinePlot, LinearMapper, DataRange1D, OverlayPlotContainer)
 
23
 
 
24
NUM_SAMPLES = 1024
 
25
SAMPLING_RATE = 11025
 
26
SPECTROGRAM_LENGTH = 50
 
27
 
 
28
class WaterfallRenderer(LinePlot):
 
29
 
 
30
    # numpy arrays of the same length
 
31
    values = List(args=[])
 
32
 
 
33
    # Maps each array in values into a contrained, short screen space
 
34
    y2_mapper = Instance(AbstractMapper)
 
35
    
 
36
    _cached_data_pts = List()
 
37
    _cached_screen_pts = List()
 
38
 
 
39
    def _gather_points(self):
 
40
        if not self._cache_valid:
 
41
            if not self.index or len(self.values) == 0:
 
42
                return
 
43
 
 
44
            index = self.index.get_data()
 
45
            values = self.values
 
46
 
 
47
            numindex = len(index)
 
48
            if numindex == 0 or all(len(v)==0 for v in values) or all(numindex != len(v) for v in values):
 
49
                self._cached_data_pts = []
 
50
                self._cache_valid = True
 
51
            
 
52
            self._cached_data_pts = [transpose(array((index, v))) for v in values]
 
53
            self._cache_value = True
 
54
        return
 
55
 
 
56
    def get_screen_points(self):
 
57
        self._gather_points()
 
58
        return [self.map_screen(pts, i) for i, pts in enumerate(self._cached_data_pts)]
 
59
 
 
60
    def map_screen(self, data_array, data_offset=None):
 
61
        """ data_offset, if provided, is a float that will be mapped 
 
62
        into screen space using self.value_mapper and then added to
 
63
        mapping data_array with y2_mapper.  If data_offset is not
 
64
        provided, then y2_mapper is used.
 
65
        """
 
66
        if len(data_array) == 0:
 
67
            return []
 
68
        x_ary, y_ary = transpose(data_array)
 
69
        sx = self.index_mapper.map_screen(x_ary)
 
70
        if data_offset is not None:
 
71
            dy = self.value_mapper.map_screen(data_offset)
 
72
            sy = self.y2_mapper.map_screen(y_ary) + dy
 
73
        else:
 
74
            sy = self.value_mapper.map_screen(y_ary)
 
75
 
 
76
        if self.orientation == "h":
 
77
            return transpose(array((sx, sy)))
 
78
        else:
 
79
            return transpose(array((sy, sx)))
 
80
 
 
81
#============================================================================
 
82
# Create the Chaco plot.
 
83
#============================================================================
 
84
 
 
85
def _create_plot_component(obj):
 
86
    # Setup the spectrum plot
 
87
    frequencies = linspace(0.0, float(SAMPLING_RATE)/2, num=NUM_SAMPLES/2)
 
88
    obj.spectrum_data = ArrayPlotData(frequency=frequencies)
 
89
    empty_amplitude = zeros(NUM_SAMPLES/2)
 
90
    obj.spectrum_data.set_data('amplitude', empty_amplitude)
 
91
 
 
92
    obj.spectrum_plot = Plot(obj.spectrum_data)
 
93
    spec_renderer = obj.spectrum_plot.plot(("frequency", "amplitude"), name="Spectrum",
 
94
                           color="red")[0]
 
95
    obj.spectrum_plot.padding = 50
 
96
    obj.spectrum_plot.title = "Spectrum"
 
97
    spec_range = obj.spectrum_plot.plots.values()[0][0].value_mapper.range
 
98
    spec_range.low = 0.0
 
99
    spec_range.high = 5.0
 
100
    obj.spectrum_plot.index_axis.title = 'Frequency (hz)'
 
101
    obj.spectrum_plot.value_axis.title = 'Amplitude'
 
102
 
 
103
    # Time Series plot
 
104
    times = linspace(0.0, float(NUM_SAMPLES)/SAMPLING_RATE, num=NUM_SAMPLES)
 
105
    obj.time_data = ArrayPlotData(time=times)
 
106
    empty_amplitude = zeros(NUM_SAMPLES)
 
107
    obj.time_data.set_data('amplitude', empty_amplitude)
 
108
 
 
109
    obj.time_plot = Plot(obj.time_data)
 
110
    obj.time_plot.plot(("time", "amplitude"), name="Time", color="blue")
 
111
    obj.time_plot.padding = 50
 
112
    obj.time_plot.title = "Time"
 
113
    obj.time_plot.index_axis.title = 'Time (seconds)'
 
114
    obj.time_plot.value_axis.title = 'Amplitude'
 
115
    time_range = obj.time_plot.plots.values()[0][0].value_mapper.range
 
116
    time_range.low = -0.2
 
117
    time_range.high = 0.2
 
118
 
 
119
    # Spectrogram plot
 
120
    values = [zeros(NUM_SAMPLES/2) for i in xrange(SPECTROGRAM_LENGTH)]
 
121
    p = WaterfallRenderer(index = spec_renderer.index, values = values,
 
122
            index_mapper = LinearMapper(range = obj.spectrum_plot.index_mapper.range),
 
123
            value_mapper = LinearMapper(range = DataRange1D(low=0, high=SPECTROGRAM_LENGTH)),
 
124
            y2_mapper = LinearMapper(low_pos=0, high_pos=8, 
 
125
                            range=DataRange1D(low=0, high=15)),
 
126
            )
 
127
    spectrogram_plot = p
 
128
    obj.spectrogram_plot = p
 
129
    dummy = Plot()
 
130
    dummy.padding = 50
 
131
    dummy.index_axis.mapper.range = p.index_mapper.range
 
132
    dummy.index_axis.title = "Frequency (hz)"
 
133
    dummy.add(p)
 
134
 
 
135
    container = HPlotContainer()
 
136
    container.add(obj.spectrum_plot)
 
137
    container.add(obj.time_plot)
 
138
 
 
139
    c2 = VPlotContainer()
 
140
    c2.add(dummy)
 
141
    c2.add(container)
 
142
 
 
143
    return c2
 
144
 
 
145
 
 
146
def get_audio_data():
 
147
    pa = pyaudio.PyAudio()
 
148
    stream = pa.open(format=pyaudio.paInt16, channels=1, rate=SAMPLING_RATE,
 
149
                     input=True, frames_per_buffer=NUM_SAMPLES)
 
150
    audio_data  = fromstring(stream.read(NUM_SAMPLES), dtype=short)
 
151
    stream.close()
 
152
    normalized_data = audio_data / 32768.0
 
153
    return (abs(fft(normalized_data))[:NUM_SAMPLES/2], normalized_data)
 
154
 
 
155
 
 
156
# HasTraits class that supplies the callable for the timer event.
 
157
class TimerController(HasTraits):
 
158
 
 
159
    def onTimer(self, *args):
 
160
        spectrum, time = get_audio_data()
 
161
        self.spectrum_data.set_data('amplitude', spectrum)
 
162
        self.time_data.set_data('amplitude', time)
 
163
        spec_data = self.spectrogram_plot.values[1:] + [spectrum]
 
164
        self.spectrogram_plot.values = spec_data
 
165
        self.spectrum_plot.request_redraw()
 
166
        return
 
167
 
 
168
#============================================================================
 
169
# Attributes to use for the plot view.
 
170
size = (900,850)
 
171
title = "Audio Spectrum Waterfall"
 
172
 
 
173
#============================================================================
 
174
# Demo class that is used by the demo.py application.
 
175
#============================================================================
 
176
 
 
177
class DemoHandler(Handler):
 
178
 
 
179
    def closed(self, info, is_ok):
 
180
        """ Handles a dialog-based user interface being closed by the user.
 
181
        Overridden here to stop the timer once the window is destroyed.
 
182
        """
 
183
 
 
184
        info.object.timer.Stop()
 
185
        return
 
186
 
 
187
class Demo(HasTraits):
 
188
 
 
189
    plot = Instance(Component)
 
190
 
 
191
    controller = Instance(TimerController, ())
 
192
 
 
193
    timer = Instance(Timer)
 
194
 
 
195
    traits_view = View(
 
196
                    Group(
 
197
                        Item('plot', editor=ComponentEditor(size=size),
 
198
                             show_label=False),
 
199
                        orientation = "vertical"),
 
200
                    resizable=True, title=title,
 
201
                    width=size[0], height=size[1]+25,
 
202
                    handler=DemoHandler
 
203
                    )
 
204
 
 
205
    def __init__(self, **traits):
 
206
        super(Demo, self).__init__(**traits)
 
207
        self.plot = _create_plot_component(self.controller)
 
208
 
 
209
    def edit_traits(self, *args, **kws):
 
210
        # Start up the timer! We should do this only when the demo actually
 
211
        # starts and not when the demo object is created.
 
212
        self.timer = Timer(20, self.controller.onTimer)
 
213
        return super(Demo, self).edit_traits(*args, **kws)
 
214
 
 
215
    def configure_traits(self, *args, **kws):
 
216
        # Start up the timer! We should do this only when the demo actually
 
217
        # starts and not when the demo object is created.
 
218
        self.timer = Timer(20, self.controller.onTimer)
 
219
        return super(Demo, self).configure_traits(*args, **kws)
 
220
 
 
221
 
 
222
popup = Demo()
 
223
 
 
224
#============================================================================
 
225
# Stand-alone frame to display the plot.
 
226
#============================================================================
 
227
 
 
228
from enthought.etsconfig.api import ETSConfig
 
229
 
 
230
if ETSConfig.toolkit == "wx":
 
231
 
 
232
    import wx
 
233
    class PlotFrame(DemoFrame):
 
234
 
 
235
        def _create_window(self):
 
236
 
 
237
            self.controller = TimerController()
 
238
            container = _create_plot_component(self.controller)
 
239
            # Bind the exit event to the onClose function which will force the
 
240
            # example to close. The PyAudio package causes problems that normally
 
241
            # prevent the user from closing the example using the 'X' button.
 
242
            # NOTE: I believe it is sufficient to just stop the timer-Vibha.
 
243
            self.Bind(wx.EVT_CLOSE, self.onClose)
 
244
 
 
245
            # Set the timer to generate events to us
 
246
            timerId = wx.NewId()
 
247
            self.timer = wx.Timer(self, timerId)
 
248
            self.Bind(wx.EVT_TIMER, self.controller.onTimer, id=timerId)
 
249
            self.timer.Start(20.0, wx.TIMER_CONTINUOUS)
 
250
 
 
251
            # Return a window containing our plots
 
252
            return Window(self, -1, component=container)
 
253
 
 
254
        def onClose(self, event):
 
255
            #sys.exit()
 
256
            self.timer.Stop()
 
257
            event.Skip()
 
258
 
 
259
elif ETSConfig.toolkit == "qt4":
 
260
 
 
261
    from enthought.qt import QtGui, QtCore
 
262
 
 
263
    class PlotFrame(DemoFrame):
 
264
        def _create_window(self):
 
265
            self.controller = TimerController()
 
266
            container = _create_plot_component(self.controller)
 
267
 
 
268
            # start a continuous timer
 
269
            self.timer = QtCore.QTimer()
 
270
            self.timer.timeout.connect(self.controller.onTimer)
 
271
            self.timer.start(20)
 
272
 
 
273
            return Window(self, -1, component=container)
 
274
 
 
275
        def closeEvent(self, event):
 
276
            # stop the timer
 
277
            if getattr(self, "timer", None):
 
278
                self.timer.stop()
 
279
            return super(PlotFrame, self).closeEvent(event)
 
280
 
 
281
if __name__ == "__main__":
 
282
    demo_main(PlotFrame, size=size, title=title)