diff --git a/scripts/selcal-fft.py b/scripts/selcal-fft.py index b37d4d4..bbe411d 100755 --- a/scripts/selcal-fft.py +++ b/scripts/selcal-fft.py @@ -8,6 +8,8 @@ from scipy.signal import butter, lfilter, decimate import matplotlib.pyplot as plt import sys +file_name = sys.argv[1] + frequencies = np.array([10 ** ((n + 22) * 0.0225 + 2) for n in range(32)]) widths = np.array([6.25 * 10 ** (n * 0.0225) for n in range(32)]) widths = np.array([6.25 for n in range(32)]) @@ -54,7 +56,7 @@ def get_largest_two_indices(numbers, threshold): return None # Step 1: Read the WAV file -sample_rate, data = wavfile.read(sys.argv[1]) +sample_rate, data = wavfile.read(file_name) # Handle stereo audio by converting to mono if needed if len(data.shape) == 2: @@ -132,16 +134,44 @@ for i in range(act_segs): pass -# Step 5: Plot the spectrogram -plt.figure(figsize=(12, 8)) -extent = [0, num_segments * increment, 0, max_freq] -plt.imshow(dft_results.T, aspect='auto', origin='lower', extent=extent, cmap='viridis') -plt.colorbar(label='Magnitude') -plt.title('Spectrogram (100 Hz to 2000 Hz)') -plt.xlabel('Time (s)') -plt.ylabel('Frequency (Hz)') +# Only import if we're actually plotting, these imports are pretty heavy. +import pyqtgraph as pg +from pyqtgraph.Qt import QtGui, QtWidgets, QtCore, mkQApp + +app = pg.mkQApp("ImageView Example") + +## Create window with ImageView widget +win = QtWidgets.QMainWindow() +win.resize(1280, 720) +imv = pg.GraphicsLayoutWidget(show=True) +win.setCentralWidget(imv) +correlogram = pg.ImageItem() +correlogram.setImage(dft_results) +img_rect = QtCore.QRectF(0, 0, len(filtered_data) // sample_rate, max_freq) +correlogram.setRect(img_rect) + +plotItem = imv.addPlot() # add PlotItem to the main GraphicsLayoutWidget +#plotItem.setDefaultPadding(0.0) # plot without padding data range +plotItem.addItem(correlogram) # display correlogram +#plotItem.showAxes(True, showValues=(True, True, False, False), size=20) + +freq_pen = pg.mkPen(color=(20, 20, 20), width=1, style=QtCore.Qt.DashLine) for freq in frequencies: - plt.axhline(y=freq, color='r', linestyle='--', linewidth=0.5) -plt.show() \ No newline at end of file + horizontal_line = pg.InfiniteLine(pos=freq, angle=0, pen=freq_pen) + plotItem.addItem(horizontal_line) + +yticks = [(i, str(round(i))) for i in frequencies] +plotItem.getAxis('left').setTicks([yticks]) + +colorMap = pg.colormap.get("CMRmap", source='matplotlib') # choose perceptually uniform, diverging color map +# generate an adjustabled color bar, initially spanning -1 to 1: +bar = pg.ColorBarItem( values=(0,1), colorMap=colorMap) +# link color bar and color map to correlogram, and show it in plotItem: +bar.setImageItem(correlogram, insert_in=plotItem) + +win.show() +win.setWindowTitle('pyqtgraph example: ImageView') + +pg.exec()