File size: 5,013 Bytes
2e1dc7f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
#%% Imports
from pathlib import Path
import subprocess as sp
import essentia.standard as es
import config as cfg
import numpy as np
import librosa
from multipledispatch import dispatch
from torch import Tensor
import torchaudio

from data.labels import GENRE_LABELS, MOOD_THEME_CLASSES, INSTRUMENT_CLASSES
from utils import audio as audio_utils
import essentia

#%% Download models
if False:
    sp.call([
        "curl",
        "https://essentia.upf.edu/models/classification-heads/genre_discogs400/genre_discogs400-discogs-effnet-1.pb",
        "--output", "genre_discogs400-discogs-effnet-1.pb"
    ])
    sp.call([
        "curl",
        "https://essentia.upf.edu/models/feature-extractors/discogs-effnet/discogs-effnet-bs64-1.pb",
        "--output", "discogs-effnet-bs64-1.pb"
    ])
    sp.call([
        "curl",
        "https://essentia.upf.edu/models/classification-heads/mtg_jamendo_moodtheme/mtg_jamendo_moodtheme-discogs-effnet-1.pb",
        "--output", "mtg_jamendo_moodtheme-discogs-effnet-1.pb"
    ])
    sp.call([
        "curl",
        "https://essentia.upf.edu/models/classification-heads/mtg_jamendo_instrument/mtg_jamendo_instrument-discogs-effnet-1.pb",
        "--output", "mtg_jamendo_instrument-discogs-effnet-1.pb"
    ])


def filter_predictions(predictions, class_list, threshold=0.1):
    predictions_mean = np.mean(predictions, axis=0)
    sorted_indices = np.argsort(predictions_mean)[::-1]
    filtered_indices = [
        i for i in sorted_indices if predictions_mean[i] > threshold
    ]
    filtered_labels = [class_list[i] for i in filtered_indices]
    filtered_values = [predictions_mean[i] for i in filtered_indices]
    return filtered_labels, filtered_values


def make_comma_separated_unique(tags):
    seen_tags = set()
    result = []
    for tag in ', '.join(tags).split(', '):
        if tag not in seen_tags:
            result.append(tag)
            seen_tags.add(tag)
    return ', '.join(result)


@dispatch(Path)
def get_audio_features(audio_filename: Path):  # type: ignore
    audio = audio_utils.load_audio(audio_filename, 16_000, False).squeeze()
    # audio = es.MonoLoader(filename=str(audio_filename),
    #                       sampleRate=16000,
    #                       resampleQuality=4)()

    return get_audio_features(audio, 16_000)


@dispatch(Tensor, int, Path)
def get_audio_features(audio: Tensor, sr: int,
                       models_dir: Path):  # type: ignore

    essentia.log.infoActive = False
    audio = audio_utils.to_mono(audio)
    audio = torchaudio.functional.resample(audio, sr, 16_000).squeeze()
    audio = audio.numpy()

    embedding_model = es.TensorflowPredictEffnetDiscogs(
        graphFilename=str(models_dir / "discogs-effnet-bs64-1.pb"),
        output="PartitionedCall:1")
    embeddings = embedding_model(audio)

    result_dict = {}

    # Predicting genres
    genre_model = es.TensorflowPredict2D(
        graphFilename=str(models_dir / "genre_discogs400-discogs-effnet-1.pb"),
        input="serving_default_model_Placeholder",
        output="PartitionedCall:0")
    predictions = genre_model(embeddings)

    filtered_labels, _ = filter_predictions(predictions, GENRE_LABELS)
    filtered_labels = ', '.join(filtered_labels).replace("---",
                                                         ", ").split(', ')
    result_dict['genres'] = make_comma_separated_unique(filtered_labels)

    # Predicting mood/theme
    mood_model = es.TensorflowPredict2D(
        graphFilename=str(models_dir /
                          "mtg_jamendo_moodtheme-discogs-effnet-1.pb"))
    predictions = mood_model(embeddings)
    filtered_labels, _ = filter_predictions(predictions,
                                            MOOD_THEME_CLASSES,
                                            threshold=0.05)
    result_dict['moods'] = make_comma_separated_unique(filtered_labels)

    bpm, key = get_bpm_key(audio, sr)

    result_dict["bpm"] = bpm
    result_dict["key"] = key

    # Predicting instruments
    # instrument_model = es.TensorflowPredict2D(
    #     graphFilename="mtg_jamendo_instrument-discogs-effnet-1.pb")
    # predictions = instrument_model(embeddings)
    # filtered_labels, _ = filter_predictions(predictions, INSTRUMENT_CLASSES)
    # result_dict['instruments'] = filtered_labels

    return result_dict


@dispatch(Path)
def get_bpm_key(audio_filename):
    y, sr = librosa.load(str(audio_filename))
    get_bpm_key(y, sr)


@dispatch(np.ndarray, int)
def get_bpm_key(audio: np.ndarray, sr: int):
    tempo, _ = librosa.beat.beat_track(y=audio, sr=sr)
    tempo = round(tempo[0])
    chroma = librosa.feature.chroma_stft(y=audio, sr=sr)
    key = np.argmax(np.sum(chroma, axis=1))
    key = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B'][key]
    length = librosa.get_duration(y=audio, sr=sr)

    return tempo, key


#%% Test on demo audio
if __name__ == "__main__":
    audio_filename = cfg.AUDIO_DIR / "cake.wav"
    features = get_audio_features(audio_filename)

# %