121 lines
3.9 KiB
Python
121 lines
3.9 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import scipy.signal
|
|
from scipy.io import wavfile
|
|
import sounddevice as sd
|
|
import random
|
|
from pesq import pesq
|
|
import json
|
|
|
|
|
|
|
|
#SIGNAL_PATH = "speechfiles/sp01.wav"
|
|
#NOISE_PATH = "noisefiles/white.dat"
|
|
|
|
|
|
# Scale the signal to the range [-1,1]
|
|
def normalize_signal(signal):
|
|
min_amp = np.min(signal)
|
|
normalized_signal = signal - min_amp
|
|
max_amp = np.max(normalized_signal)
|
|
normalized_signal *= 2.0/max_amp
|
|
normalized_signal -= 1
|
|
return normalized_signal
|
|
|
|
|
|
# Load an audio file from disk
|
|
def load_audiofile(path):
|
|
sound_data = []
|
|
sample_rate = 8000
|
|
# Load .dat files as sound files sampled at 8[kHz]
|
|
if path[-3:] == "dat":
|
|
with open(path, "r") as sound_file:
|
|
sound_data_strings = sound_file.readlines()
|
|
for data_string in sound_data_strings:
|
|
sound_data.append(eval(data_string.strip()))
|
|
sound_data = np.array(sound_data, dtype=np.float64)
|
|
elif path[-3:] == "wav":
|
|
sample_rate, sound_data = wavfile.read(path)
|
|
# Make sure it is nparray of floats (trust me bro, normalizing yells at you if its ints)
|
|
sound_data = np.array(sound_data, dtype=np.float64)
|
|
return sample_rate, sound_data
|
|
|
|
|
|
# Add noise to a signal with a desired SNR
|
|
def add_noise(signal, noise, snr):
|
|
len_signal = len(signal)
|
|
len_noise = len(noise)
|
|
|
|
# Get a random crop of the noise to match the length of the signal
|
|
noise_crop_start = random.randrange(len_noise-len_signal)
|
|
noise_crop = noise[noise_crop_start:noise_crop_start+len_signal]
|
|
|
|
# Calculate the power of the signal and noise
|
|
noise_power = np.linalg.norm(noise_crop, 2)
|
|
signal_power = np.linalg.norm(signal, 2)
|
|
|
|
# Adjust the noise level to match desired SNR
|
|
u = 10**(snr/20)
|
|
desired_noise_power = signal_power/u
|
|
ratio = desired_noise_power / noise_power
|
|
noise_crop *= ratio
|
|
|
|
noisy_signal = signal + noise_crop
|
|
return noisy_signal
|
|
|
|
|
|
def main():
|
|
noise_paths = ("noisefiles/white.dat", "noisefiles/train.dat", "noisefiles/street.dat", "noisefiles/exhibition.dat")
|
|
|
|
# Compose signal paths for the 30 sentences
|
|
signal_paths = []
|
|
for i in range(1,30+1):
|
|
signal_paths.append(f"speechfiles/sp{i:02}.wav")
|
|
|
|
# SNR in dB
|
|
snrs = (0, 10, 20, 30)
|
|
|
|
pesqs = {"unfiltered": {}, "filtered": {}}
|
|
|
|
for snr in snrs:
|
|
pesqs["unfiltered"][snr] = {}
|
|
pesqs["filtered"][snr] = {}
|
|
|
|
for noise_path in noise_paths:
|
|
pesqs["filtered"][snr][noise_path[:-4]] = []
|
|
pesqs["unfiltered"][snr][noise_path[:-4]] = []
|
|
|
|
for signal_path in signal_paths:
|
|
noise_sample_rate, noise_data = load_audiofile(noise_path)
|
|
signal_sample_rate, signal_data = load_audiofile(signal_path)
|
|
|
|
assert signal_sample_rate == noise_sample_rate, "Signal and noise sampling rates didn't match."
|
|
sample_rate = signal_sample_rate
|
|
|
|
noisy_signal = add_noise(signal_data, noise_data, snr)
|
|
filtered_signal = scipy.signal.wiener(noisy_signal)
|
|
|
|
noisy_pesq = pesq(sample_rate, signal_data, noisy_signal, mode='nb')
|
|
filtered_pesq = pesq(sample_rate, signal_data, filtered_signal, mode='nb')
|
|
|
|
pesqs["filtered"][snr][noise_path[:-4]].append(noisy_pesq)
|
|
pesqs["unfiltered"][snr][noise_path[:-4]].append(filtered_pesq)
|
|
|
|
pesqs_json = json.dumps(pesqs, indent=4)
|
|
with open("pesqs.json", "w") as outfile:
|
|
outfile.write(pesqs_json)
|
|
|
|
plt.plot(snrs, noisy_pesqs, label="PESQ Unfiltered")
|
|
plt.plot(snrs, filtered_pesqs, label="PESQ Filtered")
|
|
plt.xlabel("SNR [dB]")
|
|
plt.ylabel("PESQ")
|
|
plt.legend()
|
|
plt.show()
|
|
|
|
#sd.play(normalize_signal(noisy_signal), samplerate=sample_rate, blocking=True)
|
|
#sd.play(normalize_signal(filtered_signal), samplerate=sample_rate, blocking=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|