Converted fft detector to pyqtgraph as well

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()
+    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()