File size: 4,242 Bytes
1eabf9f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
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')