Adding noise, Wiener Filtering, and PESQ working

This commit is contained in:
Adog64 2024-04-25 21:07:00 -04:00
parent 50a5e57e18
commit b2dd75319f
2 changed files with 82 additions and 41 deletions

View File

@ -0,0 +1,82 @@
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
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():
signal_sample_rate, signal_data = load_audiofile(SIGNAL_PATH)
noise_sample_rate, noise_data = load_audiofile(NOISE_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, 0)
filtered_signal = scipy.signal.wiener(noisy_signal)
print(pesq(sample_rate, signal_data, noisy_signal, mode='nb'))
print(pesq(sample_rate, signal_data, filtered_signal, mode='nb'))
#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()

View File

@ -1,41 +0,0 @@
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal
import sounddevice as sd
SOUND_PATH = "noisefiles/train.dat"
def normalize_signal(signal):
min_amp = np.min(signal)
normalized_signal = signal - min_amp
max_amp = np.max(normalized_signal)
normalized_signal *= 2/max_amp
normalized_signal -= 1
return normalized_signal
def load_audiofile(path):
sound_data = []
sample_rate = 8000
if path[-3:] == "dat":
with open(SOUND_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)
elif path[-3:] == "wav":
sample_rate, sound_data = wavfile.read(path)
return sample_rate, sound_data
def main():
sample_rate, sound_data = load_audiofile(SOUND_PATH)
print(sample_rate)
sd.play(normalize_signal(sound_data), samplerate=sample_rate, blocking=True)
if __name__ == "__main__":
main()