Comparing OTs

import os
import librosa
import librosa.display
import numpy as np
from scipy.signal import butter, lfilter
from scipy.io import wavfile
import matplotlib.pyplot as plt

def butter_bandpass(lowcut, highcut, fs, order=5):
    nyquist = 0.5 * fs
    low = lowcut / nyquist
    high = highcut / nyquist
    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


AUDIO_DIR = 'data/audio'
MP3s = [os.path.join(AUDIO_DIR, o) for o in os.listdir(AUDIO_DIR) if o.endswith('.mp3')]

fig, ax = plt.subplots(figsize=(12, 6))

# Create frequency bins (logarithmically spaced in the audible range)
bins = np.logspace(np.log10(20), np.log10(20000), 20)  # 20 bins from 20Hz to 20kHz

# First pass: calculate the maximum amplitude across all files
max_amplitude = 0
all_binned_spectra = []

for mp3 in MP3s:
    # Load the audio
    y, sr = librosa.load(mp3)
    
    # Compute the frequency distribution
    spectrum = np.abs(np.fft.rfft(y))
    freq = np.fft.rfftfreq(len(y), 1/sr)
    
    # Bin the spectrum data
    binned_spectrum = np.zeros_like(bins[:-1])
    for j in range(len(bins) - 1):
        mask = (freq >= bins[j]) & (freq < bins[j+1])
        if np.any(mask):
            binned_spectrum[j] = np.mean(spectrum[mask])
    
    # Track the maximum value across all files
    file_max = np.max(binned_spectrum)
    if file_max > max_amplitude:
        max_amplitude = file_max
    
    all_binned_spectra.append(binned_spectrum)

# Second pass: plot with normalization against the global maximum
for i, mp3 in enumerate(MP3s):
    # Get basename for display
    basename = os.path.basename(mp3).replace('.mp3', '')
    
    # Normalize using the global maximum value
    binned_spectrum = all_binned_spectra[i] / max_amplitude
    
    # Plot the binned frequency distribution
    ax.plot(bins[:-1], binned_spectrum, label=basename, alpha=0.7)

# Set labels and title
ax.set_xlabel('Frequency')
ax.set_ylabel('Normed Amplitude')
ax.set_title('Frequency Distributions - OTs in JS Comparison (first clips)')
ax.set_xscale('log')
ax.grid(True, which="both", ls="-", alpha=0.5)
ax.set_xlim(20, 15000)  # Focus on audible range

# Format x-axis with standard frequency labels
frequencies = [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000]
ax.set_xticks(frequencies)
ax.set_xticklabels([f'{f}Hz' if f < 1000 else f'{f//1000}kHz' for f in frequencies])

ax.legend()

plt.tight_layout()

# Save the plot
plt.savefig('data/audio_frequency_distributions.png')

# Apply bandpass filters and save filtered audio
for i, mp3 in enumerate(MP3s):
    # Get basename for display
    basename = os.path.basename(mp3).replace('.mp3', '')
    
    # Load the audio
    y, sr = librosa.load(mp3)
    
    # Apply the bandpass filter
    y_filtered = butter_bandpass_filter(y, lowcut, highcut, fs)

    wavfile.write(f'data/audio/{basename}_filtered.wav', sr, y_filtered.astype(np.float32))