-
Notifications
You must be signed in to change notification settings - Fork 0
/
onair_example.py
72 lines (58 loc) · 1.86 KB
/
onair_example.py
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
"""
script that demonstrate the use of musdb for evaluating separation results of musdb
"""
import onair
import museval
import stempeg
import sys
import os
import numpy as np
import argparse
from onair.config import stem_ids
def mix_as_estimate(track):
# get the mixture path for external processing
track.path
# get the sample rate
track.rate
# return any number of targets
estimates = {
'vocals': track.audio,
'accompaniment': track.audio,
}
return estimates
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="onair-py test program")
parser.add_argument("--root", type=str, help="root path of dataset")
parser.add_argument(
"--out-stem-dir",
type=str,
default=None,
help="provide output stem dir",
)
# Model Parameters
parser.add_argument(
"--eval-dir",
type=str,
default=None,
help="perform BSS evaluation, store eval this dir",
)
args = parser.parse_args()
# initiate musdb
onair_db = onair.DB(root=args.root)
for track in onair_db:
print(f'track: {track}')
track_stems = track.stems
if args.out_stem_dir is not None:
for stem, stem_idx in stem_ids.items():
curr_track = track_stems[stem_idx, :, :]
target_path = os.path.join(args.out_stem_dir, f'{track.name}-{stem}.wav')
stempeg.write_audio(
path=target_path,
data=curr_track,
sample_rate=track.rate
)
print(f'wrote audio for stem {stem}, stem index {stem_idx}, to {target_path}')
if args.eval_dir is not None:
estimates = mix_as_estimate(track)
scores = museval.eval_mus_track(track, estimates, output_dir=args.eval_dir)
print(scores)