70 lines
2.5 KiB
Python
Executable File
70 lines
2.5 KiB
Python
Executable File
#! /usr/bin/env python3
|
|
|
|
import csv
|
|
import sys
|
|
import numpy as np
|
|
from scipy import signal
|
|
from scipy.io import wavfile
|
|
from tones import TONES
|
|
from filters import bandpass_filter, note, smoothing_filter, anti_alias
|
|
|
|
pure_sample_length = 0.1
|
|
|
|
if __name__ == '__main__':
|
|
file_name = sys.argv[1]
|
|
sample_rate, data = wavfile.read(file_name)
|
|
print(f"{file_name}: {len(data)} samples @ {sample_rate} Hz")
|
|
|
|
if len(data.shape) == 2:
|
|
data = data.mean(axis=1)
|
|
|
|
# Normalise
|
|
print(np.max(data))
|
|
data = data / np.max(data)
|
|
|
|
# TODO JMT: Find out why the correlation step fails when max frequency <= 2 * nyquist rate
|
|
data, sample_rate, decimation = anti_alias(data, sample_rate, 4800)
|
|
print(f'Length after decimation: {len(data)} samples (/{decimation}, {sample_rate})')
|
|
|
|
pure_signals = {tone:note(freq, pure_sample_length, 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()}
|
|
massaged = {tone:smoothing_filter(correlation) 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([])
|
|
layout = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")
|
|
layout.setGeometry(0, 0, 1600, 960)
|
|
|
|
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(massaged.items(), colors):
|
|
line = plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0.1, name=tone)
|
|
legend.addItem(line, tone)
|
|
|
|
y_max = max(line.getData()[1]) # Maximum y-value
|
|
x_max = line.getData()[0][np.argmax(line.getData()[1])] # Corresponding x-coordinate
|
|
|
|
label = pg.TextItem(html=f'<div style="text-align: center"><span style="color: #FFFFFF; font-size: 12pt;">{line.opts["name"]}</span></div>', anchor=(0.5, 0.5))
|
|
plot.addItem(label)
|
|
label.setPos(x_max, y_max)
|
|
label.setZValue(100) # Ensure label is above other items
|
|
|
|
plot.setLabel('left', 'Signal Correlation')
|
|
plot.setLabel('bottom', 'Time (samples)')
|
|
plot.showGrid(x=True, y=True)
|
|
|
|
legend_view.setFixedWidth(80)
|
|
layout.ci.layout.setColumnFixedWidth(1, 80)
|
|
|
|
app.exec_() |