melodium/main.py at master · kelvin-jiang/melodium · GitHub

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
import argparse
import numpy as np

from contour import create_contours, plot_contours, plot_melody, select_melody
from salience import compute_saliences, plot_saliences
from spectral import equal_loudness_filter, hop_size, inverse_spectral_transform, plot_spectral_transform, spectral_transform
from utils import load_wav_file, write_melody, write_wav_file

def extract_melody(audio, fs, args):
    # sinusoid extraction
    if not args.omit_elf:
        print("Applying equal loudness filter...")
        audio = equal_loudness_filter(audio)
    print("Applying spectral transform...")
    f, t, Zxx = spectral_transform(audio, fs)
    print("Plotting spectral transform...")
    plot_spectral_transform(f, t, Zxx, fs, args.spectral_plot)
    # TODO: frequency correction

    # compute saliences
    t_size = t.shape[0] if args.duration < 0 else -(-args.duration * fs) // hop_size
    if args.use_cached_saliences:
        print("Loading salience...")
        saliences = np.load(args.cached_saliences)
    else:
        print("Computing salience...")
        saliences = compute_saliences(f, t, Zxx, args.workers, fs, args.cached_saliences, t_size)
    print("Plotting salience...")
    plot_saliences(t, saliences, fs, args.salience_plot, t_size)

    # pitch contour creation
    print("Creating pitch contours...")
    contours, space = create_contours(saliences, fs)
    print("Plotting pitch contours...")
    plot_contours(space, fs, len(contours), args.contour_plot)

    # melody selection
    print("Selecting melody...")
    melody = select_melody(contours, space.shape[1], fs)
    print("Plotting melody...")
    plot_melody(melody, fs, './output/melody')

    # TODO: reconstruct melody as audio signal
    # _, melody_audio = inverse_spectral_transform(melody, fs)

    return melody

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True)
    parser.add_argument('--output_data', default='./output/melody.txt')
    # parser.add_argument('--output_audio')
    parser.add_argument('--omit_elf', action='store_true')
    parser.add_argument('--spectral_plot', default='./output/spectral')
    parser.add_argument('--use_cached_saliences', action='store_true')
    parser.add_argument('--cached_saliences', default='./output/salience.npy')
    parser.add_argument('--salience_plot', default='./output/salience')
    parser.add_argument('--contour_plot', default='./output/contours')
    parser.add_argument('--melody_plot', default='./output/melody')
    parser.add_argument('--workers', type=int, default=1)
    parser.add_argument('--duration', type=int, default=-1)
    args = parser.parse_args()

    fs, audio = load_wav_file(args.input)
    assert fs == 44100

    melody = extract_melody(audio, fs, args)
    write_melody(melody, args.output_data, fs, hop_size)
    # if args.output_audio:
    #     write_wav_file(melody_audio, args.output_audio, fs)

if __name__ == '__main__':
    main()