Spaces:
Sleeping
Sleeping
import os | |
from BeatNet.BeatNet import BeatNet | |
import time | |
import datetime | |
from tqdm import tqdm | |
import soundfile as sf | |
import librosa | |
import numpy as np | |
device = 'cuda' # 'cpu' or 'cuda', I found there is no difference | |
def traverse_dir( | |
root_dir, | |
extension, | |
amount=None, | |
str_include=None, | |
str_exclude=None, | |
is_pure=False, | |
is_sort=False, | |
is_ext=True): | |
file_list = [] | |
cnt = 0 | |
for root, _, files in os.walk(root_dir): | |
for file in files: | |
if file.endswith(extension): | |
# path | |
mix_path = os.path.join(root, file) | |
pure_path = mix_path[len(root_dir)+1:] if is_pure else mix_path | |
# amount | |
if (amount is not None) and (cnt == amount): | |
if is_sort: | |
file_list.sort() | |
return file_list | |
# check string | |
if (str_include is not None) and (str_include not in pure_path): | |
continue | |
if (str_exclude is not None) and (str_exclude in pure_path): | |
continue | |
if not is_ext: | |
ext = pure_path.split('.')[-1] | |
pure_path = pure_path[:-(len(ext)+1)] | |
file_list.append(pure_path) | |
cnt += 1 | |
if is_sort: | |
file_list.sort() | |
return file_list | |
def estimate_beat_beatnet(path_audio): | |
estimator = BeatNet( | |
1, | |
mode='offline', | |
inference_model='DBN', | |
plot=[], | |
thread=False, | |
device=device) | |
beats = estimator.process(path_audio) | |
return beats | |
def estimate_beat_madmom(path_audio): | |
from madmom.features.downbeats import DBNDownBeatTrackingProcessor | |
from madmom.features.downbeats import RNNDownBeatProcessor | |
# print('[*] estimating beats...') | |
proc = DBNDownBeatTrackingProcessor(beats_per_bar=[3, 4], fps=100) | |
act = RNNDownBeatProcessor()(path_audio) | |
proc_res = proc(act) | |
return proc_res | |
def export_audio_with_click(proc_res, path_audio, path_output, sr=44100): | |
# extract time | |
times_beat = proc_res[np.where(proc_res[:, 1]!=1)][:, 0] | |
times_downbeat = proc_res[np.where(proc_res[:, 1]==1)][:, 0] | |
# load | |
y, _ = librosa.core.load(path_audio, sr=sr) | |
# click audio | |
y_beat = librosa.clicks(times=times_beat, sr=sr, click_freq=1200, click_duration=0.5) * 0.6 | |
y_downbeat = librosa.clicks(times=times_downbeat, sr=sr, click_freq=600, click_duration=0.5) | |
# merge | |
max_len = max(len(y), len(y_beat), len(y_downbeat)) | |
y_integrate = np.zeros(max_len) | |
y_integrate[:len(y_beat)] += y_beat | |
y_integrate[:len(y_downbeat)] += y_downbeat | |
y_integrate[:len(y)] += y | |
# librosa.output.write_wav(path_output, y_integrate, sr) | |
sf.write(path_output, y_integrate, sr) | |
if __name__ == '__main__': | |
path_rootdir = '../audiocraft/dataset/example/full' | |
audio_base = 'no_vocals' | |
ext = 'wav' | |
st, ed = 0, None | |
filelist = traverse_dir( | |
path_rootdir, | |
extension=ext, | |
str_include=audio_base, | |
is_sort=True) | |
num_files = len(filelist) | |
print(' > num files:', num_files) | |
if ed is None: | |
ed = num_files | |
# run | |
start_time_all = time.time() | |
for i in range(num_files-1,-1,-1): | |
start_time_one = time.time() | |
print("==={}/{}======[{} - {}]========".format( | |
i, num_files, st, ed)) | |
path_audio = filelist[i] | |
path_outfile = path_audio.replace('no_vocals.wav', 'beats.npy') | |
print(' inp >', path_audio) | |
print(' out >', path_outfile) | |
if os.path.exists(path_outfile): | |
print('[o] existed') | |
continue | |
beats = estimate_beat_beatnet(path_audio) | |
# save | |
np.save(path_outfile, beats) | |
end_time_one = time.time() | |
runtime = end_time_one - start_time_one | |
print(f' > runtime:', str(datetime.timedelta(seconds=runtime))+'\n') | |
end_time_all = time.time() | |
runtime = end_time_all - start_time_all | |
print(f'testing time:', str(datetime.timedelta(seconds=runtime))+'\n') | |