Cleaned up the selcal-detect.py file, so I understand it better

This commit is contained in:
Jono Targett 2024-05-29 22:03:59 +09:30
parent b65fa14f5e
commit cd83d1683b
2 changed files with 32 additions and 71 deletions

View File

@ -1,20 +1,35 @@
import numpy as np
from scipy.signal import butter, lfilter, decimate from scipy.signal import butter, lfilter, decimate
def anti_alias(data, sample_rate, max_frequency): def anti_alias(data, sample_rate, max_frequency):
FILTER_HEADROOM = 1.2
nyquist_rate = 2 * max_frequency nyquist_rate = 2 * max_frequency
downsample_factor = 1
if sample_rate > nyquist_rate: if sample_rate > nyquist_rate:
filtered_data = lowpass_filter(data, max_frequency, sample_rate) filtered_data = lowpass_filter(data, max_frequency * FILTER_HEADROOM, sample_rate)
downsample_factor = int(sample_rate / nyquist_rate) downsample_factor = int(sample_rate / nyquist_rate)
filtered_data = decimate(filtered_data, downsample_factor) filtered_data = decimate(filtered_data, downsample_factor)
sample_rate = sample_rate // downsample_factor sample_rate = sample_rate // downsample_factor
else: else:
filtered_data = data filtered_data = data
return filtered_data, sample_rate return filtered_data, sample_rate, downsample_factor
def smoothing_filter(data, window_size=256):
window = np.ones(window_size) / window_size
return np.convolve(data, window, mode='same')
# Stolen from selcald
def note(freq, length, amp=1.0, rate=44100):
if freq == 0:
data = np.zeros(int(length * rate))
else:
t = np.linspace(0, length, int(length * rate))
data = np.sin(2 * np.pi * freq * t) * amp
return data
# These originally came from https://scipy.github.io/old-wiki/pages/Cookbook/ButterworthBandpass, # These originally came from https://scipy.github.io/old-wiki/pages/Cookbook/ButterworthBandpass,
# but they've been copied around the internet so many times that ChatGPT now produces them verbatim. # but they've been copied around the internet so many times that ChatGPT now produces them verbatim.

View File

@ -5,85 +5,31 @@ import sys
import numpy as np import numpy as np
from scipy import signal from scipy import signal
from scipy.io import wavfile from scipy.io import wavfile
from scipy.signal import butter, lfilter from tones import TONES
from filters import bandpass_filter, note, smoothing_filter, anti_alias
tones = {}
with open('tones.csv', newline='') as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
tones[row['designator']] = float(row['frequency'])
'''
def freq_of_key(midi_key):
return 440.0 * (2 ** ((midi_key - 69)/12))
tones = {}
for c in range(65, 90):
tones[c] = freq_of_key(c)
'''
# Shamelessly lifted from
# https://scipy.github.io/old-wiki/pages/Cookbook/ButterworthBandpass
def butter_bandpass(lowcut, highcut, fs, order=5):
nyq = 0.5 * fs
low = lowcut / nyq
high = highcut / nyq
b, a = butter(order, [low, high], btype='band')
return b, a
def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
b, a = butter_bandpass(lowcut, highcut, fs, order=order)
y = lfilter(b, a, data)
return y
# tone synthesis
def note(freq, cycles, amp=32767.0, rate=44100):
len = cycles * (1.0/rate)
t = np.linspace(0, len, int(len * rate))
if freq == 0:
data = np.zeros(int(len * rate))
else:
data = np.sin(2 * np.pi * freq * t) * amp
return data.astype(int)
def decimate_from_sample_rate(sample_rate):
if sample_rate == 44100:
return 4 # rate = 11025, Fmax = 5512.5 Hz
elif sample_rate == 48000:
return 5 # rate = 9600, Fmax = 4800 Hz
elif sample_rate == 22050:
return 2 # rate = 11025, Fmax = 5512.5 Hz
elif sample_rate == 11025:
return 1 # rate = 11025, Fmax = 5512.5 Hz
else:
raise ValueError("Sample rate not supported")
pure_sample_length = 0.1
if __name__ == '__main__': if __name__ == '__main__':
# TODO JMT: What is this?
FLT_LEN = 2000
file_name = sys.argv[1] file_name = sys.argv[1]
sample_rate, data = wavfile.read(file_name) sample_rate, data = wavfile.read(file_name)
print(f"{file_name}: {len(data)} samples @ {sample_rate} Hz") print(f"{file_name}: {len(data)} samples @ {sample_rate} Hz")
decimate = decimate_from_sample_rate(sample_rate) if len(data.shape) == 2:
if decimate > 1: data = data.mean(axis=1)
data = signal.decimate(data, decimate)
sample_rate = sample_rate / decimate
print(f'Length after decimation: {len(data)} samples') # Normalise
print(np.max(data))
data = data / np.max(data)
data = butter_bandpass_filter(data, 270, 1700, sample_rate, order=8) # 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, FLT_LEN, rate=sample_rate) for tone,freq in tones.items()} 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()} 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()}
N = FLT_LEN // 8 # 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. # Only import if we're actually plotting, these imports are pretty heavy.
import pyqtgraph as pg import pyqtgraph as pg
@ -100,7 +46,7 @@ if __name__ == '__main__':
legend.setParentItem(legend_view) legend.setParentItem(legend_view)
color_map = pg.colormap.get('CET-C6s') color_map = pg.colormap.get('CET-C6s')
colors = color_map.getLookupTable(nPts=len(tones)) colors = color_map.getLookupTable(nPts=len(TONES))
for (tone, correlation), color in zip(massaged.items(), colors): for (tone, correlation), color in zip(massaged.items(), colors):
line = plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0.1, name=tone) line = plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0.1, name=tone)