| """ |
| Copyright (c) : Calixte Mayoraz 2024 |
| https://gitlab.com/calixtemayoraz |
| """ |
| import os |
| import tensorflow as tf |
| import logging |
| import numpy as np |
| from .preprocess import preprocess |
|
|
|
|
| class MusicSquidModel: |
| MODEL_PATH = os.path.abspath(os.path.join(__file__, '..', 'model')) |
|
|
| def __init__(self): |
| self._model: tf.keras.Model = tf.keras.models.load_model(MusicSquidModel.MODEL_PATH) |
| logging.info("Loaded model %s", self.model_version) |
|
|
| @property |
| def model_version(self) -> str: |
| return self._model.name |
|
|
| @property |
| def output_shape(self) -> tuple: |
| return self._model.output_shape |
|
|
| def embed(self, filename) -> 'np.ndarray|tuple[np.ndarray,np.ndarray]': |
| """ |
| Returns an embedding for the input file. |
| |
| Since the model was trained on 30s segments of tracks, |
| we compute an embedding for all 30s segments of the input data and |
| return a median vector for the output embeddings. |
| |
| |
| Parameters |
| ---------- |
| filename: str |
| |
| Returns |
| ------- |
| np.ndarray|tuple[np.ndarray,np.ndarray]: |
| the embedded track or the embedded track and class probabilities |
| """ |
| return np.median(self._model(preprocess(filename)), axis=0) |
|
|
