diff --git a/scripts/selcal-fft.py b/scripts/selcal-fft.py index a9a6bc9..c730cda 100755 --- a/scripts/selcal-fft.py +++ b/scripts/selcal-fft.py @@ -5,29 +5,11 @@ from scipy.io import wavfile from scipy.fft import fft from filters import anti_alias from tones import TONES +from utilities import * import matplotlib.pyplot as plt import sys file_name = sys.argv[1] - - -def decibels(f): - return 10 * np.log10(f) - -def find_top_two_keys(d, threshold): - sorted_items = sorted(d.items(), key=lambda item: item[1], reverse=True) - - if len(sorted_items) < 3: - return None - - top_two = sorted_items[:2] - third_value = sorted_items[2][1] - if top_two[0][1] - third_value < threshold or top_two[1][1] - third_value < threshold: - return None - - return top_two[0][0], top_two[1][0] - -# Step 1: Read the WAV file sample_rate, data = wavfile.read(file_name) # Handle stereo audio by converting to mono if needed @@ -92,7 +74,7 @@ for i in range(num_segments): # Only import if we're actually plotting, these imports are pretty heavy. import pyqtgraph as pg -from pyqtgraph.Qt import QtGui, QtWidgets, QtCore, mkQApp +from pyqtgraph.Qt import QtWidgets, QtCore app = pg.mkQApp("ImageView Example") window = QtWidgets.QMainWindow() @@ -108,8 +90,7 @@ fft_view.setImage(fft_results) fft_view.setRect(QtCore.QRectF(0, 0, len(data) // sample_rate, max_freq)) plot.addItem(fft_view) -# Note JMT: Requires matplotlib installed to use this colormap -colormap = pg.colormap.get("CMRmap", source='matplotlib') +colormap = pg.colormap.get("inferno") colorbar = pg.ColorBarItem( values=(0,1), colorMap=colormap) colorbar.setImageItem(fft_view, insert_in=plot) diff --git a/scripts/utilities.py b/scripts/utilities.py new file mode 100644 index 0000000..9b6fbb8 --- /dev/null +++ b/scripts/utilities.py @@ -0,0 +1,17 @@ +import numpy as np + +def decibels(f): + return 10 * np.log10(f) + +def find_top_two_keys(d, threshold): + sorted_items = sorted(d.items(), key=lambda item: item[1], reverse=True) + + if len(sorted_items) < 3: + return None + + top_two = sorted_items[:2] + third_value = sorted_items[2][1] + if top_two[0][1] - third_value < threshold or top_two[1][1] - third_value < threshold: + return None + + return top_two[0][0], top_two[1][0]