107 lines
3.2 KiB
Python
Executable File
107 lines
3.2 KiB
Python
Executable File
#! /usr/bin/env python3
|
|
|
|
import numpy as np
|
|
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]
|
|
sample_rate, data = wavfile.read(file_name)
|
|
|
|
# Handle stereo audio by converting to mono if needed
|
|
if len(data.shape) == 2:
|
|
data = data.mean(axis=1)
|
|
|
|
max_freq = 1600.0
|
|
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
|
|
frequency_resolution = sample_rate / fft_size
|
|
max_bin = int(max_freq / frequency_resolution)
|
|
|
|
segment_interval = 0.2 # seconds
|
|
samples_per_interval = int(sample_rate * segment_interval)
|
|
num_segments = len(data) // samples_per_interval
|
|
|
|
|
|
fft_results = np.zeros((num_segments, max_bin))
|
|
|
|
for i in range(num_segments):
|
|
# Segment window is current position to fft_size samples in the past. As such some segments
|
|
# will have overlap in which samples are used when fft_size > samples_per_interval
|
|
end = (i + 1) * samples_per_interval
|
|
start = end - fft_size
|
|
|
|
try:
|
|
segment = data[start:end]
|
|
|
|
fft_result = fft(segment)
|
|
|
|
magnitudes = np.abs(fft_result)
|
|
total_energy = np.sum(magnitudes ** 2)
|
|
normalized_magnitudes = magnitudes / np.sqrt(total_energy)
|
|
|
|
# Store the normalised magnitude spectrum only for the desired frequency range
|
|
fft_results[i, :] = normalized_magnitudes[:max_bin]
|
|
|
|
tone_width = 6.25 # Hz
|
|
def band(centre_frequency, width=tone_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):
|
|
low_bin, high_bin = bins(band)
|
|
return np.sum(normalized_magnitudes[low_bin:high_bin])
|
|
|
|
scores = {
|
|
tone:decibels(magnitude_in_band(band(frequency)))
|
|
for tone,frequency in TONES.items()
|
|
}
|
|
|
|
active_tones = find_top_two_keys(scores, 3.0)
|
|
if active_tones:
|
|
print(active_tones)
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
|
|
# Only import if we're actually plotting, these imports are pretty heavy.
|
|
import pyqtgraph as pg
|
|
from pyqtgraph.Qt import QtWidgets, QtCore
|
|
|
|
app = pg.mkQApp("ImageView Example")
|
|
window = QtWidgets.QMainWindow()
|
|
window.resize(1280, 720)
|
|
window.setWindowTitle(f"SELCAL FFT Analysis: {file_name}")
|
|
|
|
layout = pg.GraphicsLayoutWidget(show=True)
|
|
window.setCentralWidget(layout)
|
|
plot = layout.addPlot()
|
|
|
|
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)
|
|
|
|
colormap = pg.colormap.get("inferno")
|
|
colorbar = pg.ColorBarItem( values=(0,1), colorMap=colormap)
|
|
colorbar.setImageItem(fft_view, insert_in=plot)
|
|
|
|
tone_pen = pg.mkPen(color=(20, 20, 20), width=1, style=QtCore.Qt.DashLine)
|
|
for frequency in TONES.values():
|
|
tone_line = pg.InfiniteLine(pos=frequency, angle=0, pen=tone_pen)
|
|
plot.addItem(tone_line)
|
|
|
|
yticks = [(frequency, tone) for tone,frequency in TONES.items()]
|
|
plot.getAxis('left').setTicks([yticks])
|
|
|
|
window.show()
|
|
pg.exec()
|