Beta v.1.0

src/music_engine/matcher.py

@@ -5,50 +5,64 @@ import joblib
 
 class MusicMatcher:
     def __init__(self, db_path: Path | str, model_path: Path | str):
-        # 1. Загрузка базы музыки
+        """
+        Инициализация движка сопоставления музыки.
+        """
         self.music_db = pd.read_csv(db_path)
         self.music_db['valence'] = pd.to_numeric(self.music_db['valence'], errors='coerce')
         self.music_db['arousal'] = pd.to_numeric(self.music_db['arousal'], errors='coerce')
         self.music_db = self.music_db.dropna()
         
-        # 2. Загрузка обученного регрессора
+        self.audio_dir = Path(db_path).parent / "DEAM_audio" / "MEMD_audio"
+        
+        if self.audio_dir.exists():
+            print(f"✅ Музыкальный архив найден: {self.audio_dir}")
+        else:
+            print(f"⚠️ ПРЕДУПРЕЖДЕНИЕ: Папка {self.audio_dir} не найдена!")
+        
         if Path(model_path).exists():
             self.regressor = joblib.load(model_path)
-            print("✅ Регрессионная модель успешно загружена.")
+            print("✅ ML-регрессор загружен.")
         else:
             self.regressor = None
-            print("⚠️ ВНИМАНИЕ: Модель va_regressor.pkl не найдена!")
+            print("⚠️ Файл модели .pkl не найден.")
 
     def predict_va(self, embedding: np.ndarray):
-        """
-        Использование обученной ML-модели (Ridge) для маппинга
-        эмбеддингов в пространство Valence-Arousal.
-        """
+        """Честный прогноз координат Valence-Arousal."""
         if self.regressor is not None:
-            # Модель ожидает двумерный массив (batch_size, features)
             emb_2d = embedding.reshape(1, -1)
-            prediction = self.regressor.predict(emb_2d)[0] # Получаем [Valence, Arousal]
-            
-            v, a = prediction[0], prediction[1]
-        else:
-            # Fallback на случай, если файл модели потеряется
-            v, a = 5.0, 5.0 
-            
-        return np.clip(v, 1.0, 9.0), np.clip(a, 1.0, 9.0)
+            prediction = self.regressor.predict(emb_2d)[0]
+            return np.clip(prediction[0], 1.0, 9.0), np.clip(prediction[1], 1.0, 9.0)
+        return 5.0, 5.0
 
-    def get_playlist(self, user_vector: np.ndarray, top_k: int = 5):
-        # 1. Предсказываем координаты через ML-модель
-        target_v, target_a = self.predict_va(user_vector)
-        
-        # 2. Считаем Евклидово расстояние (L2-норма) до треков в базе
+    def get_audio_path(self, song_id):
+        """Поиск mp3 файла по его номеру."""
+        if not self.audio_dir.exists():
+            return None
+            
+        clean_id = str(int(float(song_id)))
+        for ext in ['.mp3', '.wav']:
+            file_path = self.audio_dir / f"{clean_id}{ext}"
+            if file_path.exists():
+                return file_path
+        return None
+
+    def find_nearest_tracks(self, target_v: float, target_a: float, top_k: int = 5):
+        """
+        Поиск с использованием Взвешенного Евклидова расстояния (Weighted KNN).
+        Энергия (Arousal) получает больший вес, так как она сильнее 
+        определяет жанр и ритм композиции.
+        """
+        # Вес для Arousal = 2.0, для Valence = 1.0
+        # Это не позволит спокойным трекам (A < 4) попадать в выдачу 
+        # для энергичных запросов (A > 6).
         distances = np.sqrt(
-            (self.music_db['valence'] - target_v)**2 + 
-            (self.music_db['arousal'] - target_a)**2
+            1.0 * (self.music_db['valence'] - target_v)**2 + 
+            2.5 * (self.music_db['arousal'] - target_a)**2  # Жесткий штраф за разницу в энергии
         )
         
-        # 3. Формируем финальную выдачу
         df_result = self.music_db.copy()
         df_result['distance'] = distances
-        playlist = df_result.sort_values(by='distance').head(top_k)
         
-        return target_v, target_a, playlist
+        # Сортируем по расстоянию и берем топ-K
+        return df_result.sort_values(by='distance').head(top_k)