Updated selcal-detect.py to have a nice pretty graph

This commit is contained in:
Jono Targett 2024-05-25 16:14:15 +09:30
parent d65b9c1fd4
commit 5d18b8636c

View File

@ -4,12 +4,9 @@ import csv
import sys
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
from scipy.io import wavfile
from scipy.signal import butter, lfilter
from math import log10
import pyqtgraph as pg
from PyQt5 import QtWidgets
tones = {}
with open('tones.csv', newline='') as csvfile:
@ -55,9 +52,9 @@ def decimate_from_sample_rate(sample_rate):
raise ValueError("Sample rate not supported")
# analyze wav file
if __name__ == '__main__':
FLT_LEN = 2000 # Samples
# TODO JMT: What is this?
FLT_LEN = 2000
file_name = sys.argv[1]
sample_rate, data = wavfile.read(file_name)
@ -76,41 +73,36 @@ if __name__ == '__main__':
pure_signals = {tone:note(freq, FLT_LEN, rate=sample_rate) for tone,freq in tones.items()}
correlations = {tone:np.abs(signal.correlate(data, pure, mode='same')) for tone,pure in pure_signals.items()}
N = FLT_LEN # Rolling average length
cumsum_convolution = np.ones(N)/N
massaged = {tone:np.convolve(correlation, cumsum_convolution, mode='valid') for tone,correlation in correlations.items()}
# Only import if we're actually plotting, these imports are pretty heavy.
import pyqtgraph as pg
from PyQt5 import QtWidgets
app = QtWidgets.QApplication([])
win = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")
layout = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")
layout.setGeometry(0, 0, 1600, 960)
plot = win.addPlot(title=file_name)
plot.addLegend()
plot = layout.addPlot(title=file_name)
legend_view = layout.addViewBox()
legend = pg.LegendItem(offset=(0, 0))
legend.setParentItem(legend_view)
color_map = pg.colormap.get('CET-C6s')
colors = color_map.getLookupTable(nPts=len(tones))
for (tone, correlation), color in zip(correlations.items(), colors):
plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0, name=tone)
for (tone, correlation), color in zip(massaged.items(), colors):
line = plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0.1, name=tone)
legend.addItem(line, tone)
plot.setLabel('left', 'Signal Correlation')
plot.setLabel('bottom', 'Time (samples)')
plot.showGrid(x=True, y=True)
app.exec_()
legend_view.setFixedWidth(80)
layout.ci.layout.setColumnFixedWidth(1, 80)
# matplotlib is too slow
'''
fig, ax = plt.subplots()
# Step 4: Plot the data
for tone,correlation in correlations.items():
ax.plot(correlation, label=tone)
# Step 5: Customize the plot
ax.set_title('Multiple Lines Sharing the Same X Data')
ax.set_xlabel('Time (samples)')
ax.set_ylabel('Signal Correlation')
ax.legend() # Add a legend
ax.grid(True) # Add a grid
# Step 6: Display the plot
plt.show()
'''
app.exec_()