83 lines
2.7 KiB
Python
83 lines
2.7 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
|
|
|
|
|
|
|
|
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()
|