-
Notifications
You must be signed in to change notification settings - Fork 7
/
main.py
124 lines (97 loc) · 3.75 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
# -*- coding: utf-8 -*-
"""
This is the main file to launch pyqtgraph-spectrographer.
"""
from collections import deque
import numpy as np
import matplotlib.pyplot as plt
import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtGui
from src.microphone import MicrophoneRecorder
def generatePgColormap(cm_name):
"""Converts a matplotlib colormap to a pyqtgraph colormap.
See https://github.com/pyqtgraph/pyqtgraph/issues/561 for source."""
colormap = plt.get_cmap(cm_name)
colormap._init()
lut = (colormap._lut * 255).view(np.ndarray) # Convert matplotlib colormap from 0-1 to 0 -255 for Qt
return lut
CHUNKSIZE = 2048
SAMPLE_RATE = 44100
TIME_VECTOR = np.arange(CHUNKSIZE) / SAMPLE_RATE
N_FFT = 4096
FREQ_VECTOR = np.fft.rfftfreq(N_FFT, d=TIME_VECTOR[1] - TIME_VECTOR[0])
WATERFALL_FRAMES = int(1000 * 2048 // N_FFT)
TIMEOUT = TIME_VECTOR.max()
EPS = 1e-8
recorder = MicrophoneRecorder(sample_rate=SAMPLE_RATE, chunksize=CHUNKSIZE)
recorder.start()
win = pg.GraphicsWindow()
win.resize(1000, 600)
win.setWindowTitle('pyqtgraph spectrographer')
waveform_plot = win.addPlot(title="Waveform")
waveform_plot.showGrid(x=True, y=True)
waveform_plot.enableAutoRange('xy', False)
waveform_plot.setXRange(TIME_VECTOR.min(), TIME_VECTOR.max())
waveform_plot.setYRange(-2 ** 15, 2 ** 15 - 1)
waveform_plot.setLabel('left', "Microphone signal", units='A.U.')
waveform_plot.setLabel('bottom', "Time", units='s')
curve = waveform_plot.plot(pen='y')
def update_waveform():
global curve, data, ptr, waveform_plot, recorder
frames = recorder.get_frames()
if len(frames) == 0:
data = np.zeros((recorder.chunksize,), dtype=np.int)
else:
data = frames[-1]
curve.setData(x=TIME_VECTOR, y=data)
timer = QtCore.QTimer()
timer.timeout.connect(update_waveform)
timer.start(TIMEOUT)
fft_plot = win.addPlot(title='FFT plot')
fft_curve = fft_plot.plot(pen='y')
fft_plot.enableAutoRange('xy', False)
fft_plot.showGrid(x=True, y=True)
fft_plot.setXRange(FREQ_VECTOR.min(), FREQ_VECTOR.max())
fft_plot.setYRange(20 * np.log10(2 ** 11 * CHUNKSIZE) - 100, 20 * np.log10(2 ** 11 * CHUNKSIZE))
fft_plot.setLabel('left', "Amplitude", units='A.U.')
fft_plot.setLabel('bottom', "Frequency", units='Hz')
waterfall_data = deque(maxlen=WATERFALL_FRAMES)
def update_fft():
global data, fft_curve, fft_plot
if data.max() > 1:
X = np.abs(np.fft.rfft(np.hanning(data.size) * data, n=N_FFT))
magn = 20 * np.log10(X + EPS)
fft_curve.setData(x=FREQ_VECTOR, y=magn)
waterfall_data.append(np.log10(X + 1e-12))
timer_fft = QtCore.QTimer()
timer_fft.timeout.connect(update_fft)
timer_fft.start(TIMEOUT)
win.nextRow()
image_data = np.random.rand(20, 20)
waterfall_plot = win.addPlot(title='Waterfall plot', colspan=2)
waterfall_plot.setLabel('left', "Frequency", units='Hz')
waterfall_plot.setLabel('bottom', "Time", units='s')
waterfall_plot.setXRange(0, WATERFALL_FRAMES * TIME_VECTOR.max())
waterfall_image = pg.ImageItem()
waterfall_plot.addItem(waterfall_image)
waterfall_image.setImage(image_data)
lut = generatePgColormap('viridis')
waterfall_image.setLookupTable(lut)
# set scale: x in seconds, y in Hz
waterfall_image.scale(CHUNKSIZE / SAMPLE_RATE, FREQ_VECTOR.max() * 2. / N_FFT)
def update_waterfall():
global waterfall_data, waterfall_image
arr = np.c_[waterfall_data]
if arr.size > 0:
if arr.ndim == 1:
arr = arr[:, np.newaxis]
max = arr.max()
min = max / 10
waterfall_image.setImage(arr, levels=(min, max))
timer_waterfall = QtCore.QTimer()
timer_waterfall.timeout.connect(update_waterfall)
timer_waterfall.start(2 * TIMEOUT)
if __name__ == '__main__':
import sys
if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
QtGui.QApplication.instance().exec_()