-
Notifications
You must be signed in to change notification settings - Fork 2
/
predict.py
65 lines (51 loc) · 2.17 KB
/
predict.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
import numpy as np
from essentia.standard import (
MonoLoader,
TensorflowPredictEffnetDiscogs,
TensorflowPredict2D,
)
from labels import labels
def process_labels(label):
genre, style = label.split("---")
return f"{style}\n({genre})"
processed_labels = list(map(process_labels, labels))
class Predictor:
def __init__(self):
"""Load the model into memory and create the Essentia network for predictions"""
self.embedding_model_file = "./models/discogs-effnet-bs64-1.pb"
self.classification_model_file = "./models/genre_discogs400-discogs-effnet-1.pb"
self.output = "activations"
self.sample_rate = 16000
self.loader = MonoLoader()
self.tensorflowPredictEffnetDiscogs = TensorflowPredictEffnetDiscogs(
graphFilename=self.embedding_model_file,
output="PartitionedCall:1",
patchHopSize=128, # remove overlap between patches for efficiency
)
self.classification_model = TensorflowPredict2D(
graphFilename=self.classification_model_file,
input="serving_default_model_Placeholder",
output="PartitionedCall:0",
)
def predict(self, audio=None):
"""Run a single prediction on the model"""
print("loading audio...")
self.loader.configure(
sampleRate=self.sample_rate,
resampleQuality=4,
filename=str(audio),
)
waveform = self.loader()
# Model Inferencing
print("running the model...")
embeddings = self.tensorflowPredictEffnetDiscogs(waveform)
activations = self.classification_model(embeddings)
activations_mean = np.mean(activations, axis=0)
# Parsing Genres
result_dict = dict(zip(labels, activations_mean.tolist()))
sorted_genres = sorted(result_dict.items(), key=lambda x: x[1], reverse=True)
top_genre = sorted_genres[0][0]
genre_primary, genre_full = map(str.strip, top_genre.split("---"))
genre_secondary_full = sorted_genres[1][0]
genre_secondary = genre_secondary_full.split("---")[1].strip()
return genre_primary, genre_full, genre_secondary