83 lines
4.2 KiB
Python
83 lines
4.2 KiB
Python
import joblib
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pathlib import Path
|
|
|
|
class MusicMatcher:
|
|
def __init__(self, db_path: Path | str, model_path: Path | str):
|
|
# Загрузка базы данных музыкальных произведений
|
|
self.music_db = pd.read_csv(db_path)
|
|
self.acoustic_features = ['energy', 'flux', 'centroid', 'pitch', 'hnr', 'zcr']
|
|
|
|
# Удаление записей с пропущенными целевыми или акустическими признаками
|
|
target_columns = ['valence', 'arousal'] + self.acoustic_features
|
|
self.music_db = self.music_db.dropna(subset=target_columns)
|
|
|
|
# Масштабирование акустических параметров к диапазону [0, 1]
|
|
self.norm_db = self.music_db.copy()
|
|
for feat in self.acoustic_features:
|
|
f_min = self.norm_db[feat].min()
|
|
f_max = self.norm_db[feat].max()
|
|
if f_max > f_min:
|
|
self.norm_db[f"norm_{feat}"] = (self.norm_db[feat] - f_min) / (f_max - f_min)
|
|
else:
|
|
self.norm_db[f"norm_{feat}"] = 0.0
|
|
|
|
# Определение путей к аудиофайлам и загрузка модели регрессии
|
|
self.audio_dir = Path(db_path).parent / "DEAM_audio" / "MEMD_audio"
|
|
|
|
if Path(model_path).exists():
|
|
self.regressor = joblib.load(model_path)
|
|
else:
|
|
self.regressor = None
|
|
|
|
def predict_va(self, embedding: np.ndarray) -> tuple[float, float]:
|
|
# Прогнозирование координат Valence/Arousal по визуальному эмбеддингу
|
|
if not self.regressor:
|
|
return 5.0, 5.0
|
|
|
|
raw_prediction = self.regressor.predict(embedding.reshape(1, -1))[0]
|
|
valence_pred = np.clip(raw_prediction[0], 1.0, 9.0)
|
|
arousal_pred = np.clip(raw_prediction[1], 1.0, 9.0)
|
|
|
|
return float(valence_pred), float(arousal_pred)
|
|
|
|
def get_audio_path(self, song_id: int | float | str) -> Path | None:
|
|
# Поиск физического пути к аудиофайлу в зависимости от расширения
|
|
if not self.audio_dir.exists():
|
|
return None
|
|
|
|
clean_id = str(int(float(song_id)))
|
|
for ext in ['.mp3', '.wav']:
|
|
path = self.audio_dir / f"{clean_id}{ext}"
|
|
if path.exists():
|
|
return path
|
|
return None
|
|
|
|
def find_nearest_tracks(self, target_v: float, target_a: float, llm_profile: dict = None, top_k: int = 5) -> pd.DataFrame:
|
|
# Расчет евклидова расстояния в эмоциональном пространстве Рассела
|
|
v_dist = (self.norm_db['valence'] - target_v) ** 2
|
|
a_dist = (self.norm_db['arousal'] - target_a) ** 2
|
|
|
|
# Взвешенное расстояние с приоритетом оси активации (Arousal)
|
|
self.norm_db['emo_distance'] = np.sqrt(1.0 * v_dist + 2.5 * a_dist)
|
|
|
|
# Ранжирование только по эмоциональному критерию при отсутствии профиля LLM
|
|
if not llm_profile:
|
|
self.norm_db['final_score'] = self.norm_db['emo_distance']
|
|
return self.norm_db.sort_values(by='final_score').head(top_k)
|
|
|
|
# Расчет отклонений по вектору акустических параметров LLM
|
|
acoustic_penalty = np.zeros(len(self.norm_db))
|
|
for feat in self.acoustic_features:
|
|
if feat in llm_profile:
|
|
target_val = llm_profile[feat]
|
|
acoustic_penalty += np.abs(self.norm_db[f"norm_{feat}"] - target_val)
|
|
|
|
# Нормирование акустической дистанции
|
|
self.norm_db['acoustic_distance'] = acoustic_penalty / len(self.acoustic_features)
|
|
|
|
# Вычисление интегральной метрики соответствия (мультимодальный скоринг)
|
|
self.norm_db['final_score'] = self.norm_db['emo_distance'] + (self.norm_db['acoustic_distance'] * 2.0)
|
|
|
|
return self.norm_db.sort_values(by='final_score').head(top_k) |