diff --git a/lib/pickling.py b/lib/pickling.py new file mode 100644 index 0000000..379cbf7 --- /dev/null +++ b/lib/pickling.py @@ -0,0 +1,19 @@ +import bz2 +import pickle +import _pickle as cPickle +import sys + + +# Pickle a file and then compress it into a file with extension +def compress_pickle(filename, data): + sys.stderr.write("loading pickle %s..." % filename) + with bz2.BZ2File(filename, 'w') as f: + cPickle.dump(data, f) + +# Load any compressed pickle file +def decompress_pickle(filename): + sys.stderr.write("saving pickle %s..." % filename) + data = bz2.BZ2File(filename, 'rb') + data = cPickle.load(data) + return data + diff --git a/main.py b/main.py index cd84a43..fea049f 100644 --- a/main.py +++ b/main.py @@ -10,6 +10,9 @@ from scipy import signal from tqdm import tqdm +from lib.pickling import * +from os.path import exists + def load_wav(fpath): """Loads a .wav file and returns the data and sample rate. @@ -214,21 +217,36 @@ def plot_series_ax(ax, series, label=None): ax.plot(t, series, label=label) +def wav_to_enf(filename, nominal_freq, freq_band_size, harmonic_n=1): + """ + This code attempts to load a pickle file that contains the enf samples. + If the pickle file does not exist, the code loads the corresponding wav file, + processes it, and saves the enf samples in a new pickle file. + This approach prevents unnecessary loading and computing. + """ + pklfilename = "." + filename + ".pkl" + if exists(pklfilename): + return decompress_pickle(pklfilename) + else: + ref_data, ref_fs = load_wav(filename) + enf = enf_series(ref_data, ref_fs, nominal_freq, freq_band_size, harmonic_n) + compress_pickle(pklfilename, enf) + return enf + if __name__ == "__main__": nominal_freq = 50 freq_band_size = 0.2 - # !!!: make sure to run ./bin/download-example-files first - ref_data, ref_fs = load_wav("./001_ref.wav") + refwav = '001_ref.wav' + wav = "001.wav" - ref_enf_output = enf_series(ref_data, ref_fs, nominal_freq, freq_band_size) + # !!!: make sure to run ./bin/download-example-files first + ref_enf_output = wav_to_enf(refwav, nominal_freq, freq_band_size, harmonic_n=1) ref_enf = ref_enf_output['enf'] # !!!: make sure to run ./bin/download-example-files first - data, fs = load_wav("./001.wav") - harmonic_n = 1 - enf_output = enf_series(data, fs, nominal_freq, freq_band_size, harmonic_n=2) + enf_output = wav_to_enf(refwav, nominal_freq, freq_band_size, harmonic_n=2) target_enf = enf_output['enf'] stft = enf_output['stft']