jono/selcal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Proper cleanup of selcal-fft.py

Browse Source
This commit is contained in:
Jono Targett 2024-05-27 00:37:47 +09:30
parent 700f34e729
commit e4b0645ac6
1 changed files with 31 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
scripts/selcal-fft.py
View File

@ -1,6 +1,5 @@
#! /usr/bin/env python3 #! /usr/bin/env python3
import numpy as np import numpy as np
from scipy.io import wavfile from scipy.io import wavfile
from scipy.fft import fft from scipy.fft import fft
@ -40,14 +39,14 @@ data, sample_rate = anti_alias(data, sample_rate, max_freq)
fft_size = 1024 # Must be larger than max_freq TODO JMT: fix this, zero-pad fft_size = 1024 # Must be larger than max_freq TODO JMT: fix this, zero-pad
frequency_resolution = sample_rate / fft_size frequency_resolution = sample_rate / fft_size
high_bin = int(max_freq / frequency_resolution) max_bin = int(max_freq / frequency_resolution)
segment_interval = 0.2 # seconds segment_interval = 0.2 # seconds
samples_per_interval = int(sample_rate * segment_interval) samples_per_interval = int(sample_rate * segment_interval)
num_segments = len(data) // samples_per_interval num_segments = len(data) // samples_per_interval
fft_results = np.zeros((num_segments, high_bin)) fft_results = np.zeros((num_segments, max_bin))
for i in range(num_segments): for i in range(num_segments):
# Segment window is current position to fft_size samples in the past. As such some segments # Segment window is current position to fft_size samples in the past. As such some segments
@ -65,16 +64,18 @@ for i in range(num_segments):
normalized_magnitudes = magnitudes / np.sqrt(total_energy) normalized_magnitudes = magnitudes / np.sqrt(total_energy)
# Store the normalised magnitude spectrum only for the desired frequency range # Store the normalised magnitude spectrum only for the desired frequency range
fft_results[i, :] = normalized_magnitudes[:high_bin] fft_results[i, :] = normalized_magnitudes[:max_bin]
tone_width = 6.25 # Hz tone_width = 6.25 # Hz
def band(centre_frequency, width=tone_width): def band(centre_frequency, width=tone_width):
return np.array([centre_frequency - width, centre_frequency + width]) return (centre_frequency - width, centre_frequency + width)
def bins(band):
return (int(band[0]/frequency_resolution), int(band[1]/frequency_resolution))
def magnitude_in_band(band): def magnitude_in_band(band):
return np.sum( low_bin, high_bin = bins(band)
normalized_magnitudes[int(band[0]/frequency_resolution):int(band[1]/frequency_resolution)] return np.sum(normalized_magnitudes[low_bin:high_bin])
)
scores = { scores = {
tone:decibels(magnitude_in_band(band(frequency))) tone:decibels(magnitude_in_band(band(frequency)))
@ -85,8 +86,8 @@ for i in range(num_segments):
if active_tones: if active_tones:
print(active_tones) print(active_tones)
except: except Exception as e:
pass print(e)
# Only import if we're actually plotting, these imports are pretty heavy. # Only import if we're actually plotting, these imports are pretty heavy.
@ -94,39 +95,31 @@ import pyqtgraph as pg
from pyqtgraph.Qt import QtGui, QtWidgets, QtCore, mkQApp from pyqtgraph.Qt import QtGui, QtWidgets, QtCore, mkQApp
app = pg.mkQApp("ImageView Example") app = pg.mkQApp("ImageView Example")
window = QtWidgets.QMainWindow()
window.resize(1280, 720)
window.setWindowTitle(f"SELCAL FFT Analysis: {file_name}")
## Create window with ImageView widget layout = pg.GraphicsLayoutWidget(show=True)
win = QtWidgets.QMainWindow() window.setCentralWidget(layout)
win.resize(1280, 720) plot = layout.addPlot()
imv = pg.GraphicsLayoutWidget(show=True)
win.setCentralWidget(imv)
fft_view = pg.ImageItem()
fft_view.setImage(fft_results)
fft_view.setRect(QtCore.QRectF(0, 0, len(data) // sample_rate, max_freq))
plot.addItem(fft_view)
correlogram = pg.ImageItem() # Note JMT: Requires matplotlib installed to use this colormap
correlogram.setImage(fft_results) colormap = pg.colormap.get("CMRmap", source='matplotlib')
img_rect = QtCore.QRectF(0, 0, len(data) // sample_rate, max_freq) colorbar = pg.ColorBarItem( values=(0,1), colorMap=colormap)
correlogram.setRect(img_rect) colorbar.setImageItem(fft_view, insert_in=plot)
plotItem = imv.addPlot() # add PlotItem to the main GraphicsLayoutWidget tone_pen = pg.mkPen(color=(20, 20, 20), width=1, style=QtCore.Qt.DashLine)
#plotItem.setDefaultPadding(0.0) # plot without padding data range for frequency in TONES.values():
plotItem.addItem(correlogram) # display correlogram tone_line = pg.InfiniteLine(pos=frequency, angle=0, pen=tone_pen)
#plotItem.showAxes(True, showValues=(True, True, False, False), size=20) plot.addItem(tone_line)
freq_pen = pg.mkPen(color=(20, 20, 20), width=1, style=QtCore.Qt.DashLine)
for freq in TONES.values():
horizontal_line = pg.InfiniteLine(pos=freq, angle=0, pen=freq_pen)
plotItem.addItem(horizontal_line)
yticks = [(frequency, tone) for tone,frequency in TONES.items()] yticks = [(frequency, tone) for tone,frequency in TONES.items()]
plotItem.getAxis('left').setTicks([yticks]) plot.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')
window.show()
pg.exec() pg.exec()