diff --git a/src/main.py b/src/main.py index 7a809f9..b3e606a 100644 --- a/src/main.py +++ b/src/main.py @@ -31,39 +31,41 @@ LABELS_CSV = DATA_ROOT / "labels.csv" EMBEDDINGS_PATH = Path("./src/emoset_test_embeddings.npy") LABELS_PATH = Path("./src/emoset_test_labels.npy") -# Параметры эксперимента NUM_CHOICES = 6 TOTAL_ROUNDS = 10 -st.set_page_config(page_title="EmoSet & Music Recommendation Demo", layout="wide") +# Словарь для расшифровки меток (алфавитный порядок EmoSet) +EMO_NAMES = { + 0: "amusement (веселье)", + 1: "anger (гнев)", + 2: "awe (трепет)", + 3: "contentment (удовлетворение)", + 4: "disgust (отвращение)", + 5: "excitement (возбуждение)", + 6: "fear (страх)", + 7: "sadness (грусть)" +} + +st.set_page_config(page_title="Debug Mode: EmoSet & Music", layout="wide") -# Инициализация музыкального движка с кэшированием @st.cache_resource def load_music_engine(): - # Путь относительно src: ../dataset/DEAM/music_db.csv - db_path = Path(__file__).parent.parent / "dataset" / "DEAM" / "music_db.csv" - if not db_path.exists(): - return None - return MusicMatcher(db_path) + base_dir = Path(__file__).resolve().parent + db_path = base_dir.parent / "dataset" / "DEAM" / "music_db.csv" + model_path = base_dir / "music_engine" / "va_regressor.pkl" + if not db_path.exists(): return None + return MusicMatcher(db_path=db_path, model_path=model_path) matcher = load_music_engine() -# ---------------------------- -# 2️⃣ Загрузка данных EmoSet -# ---------------------------- @st.cache_data def load_emoset_data(): - if not IMAGES_DIR.exists() or not EMBEDDINGS_PATH.exists(): + if not IMAGES_DIR.exists() or not EMBEDDINGS_PATH.exists() or not LABELS_CSV.exists(): return None, None, None - - image_files = sorted([f.name for f in IMAGES_DIR.glob("*.jpg")]) + df = pd.read_csv(LABELS_CSV) + image_files = df['filename'].tolist() embeddings = np.load(EMBEDDINGS_PATH) labels_array = np.load(LABELS_PATH) - - if len(image_files) != len(embeddings) or len(image_files) != len(labels_array): - st.error("Размеры массивов данных не совпадают!") - st.stop() - return image_files, embeddings, labels_array image_files, embeddings, labels_array = load_emoset_data() @@ -72,69 +74,81 @@ image_files, embeddings, labels_array = load_emoset_data() # 3️⃣ Логика приложения # ---------------------------- if image_files is None: - st.error("Данные EmoSet не найдены. Проверьте папку dataset.") + st.error("Данные не найдены.") else: if 'round' not in st.session_state: st.session_state.round = 1 st.session_state.chosen_indices = [] st.session_state.current_options = random.sample(range(len(image_files)), NUM_CHOICES) - st.title("Выбор эмоциональных образов") - st.write(f"Раунд {st.session_state.round} из {TOTAL_ROUNDS}. Выберите изображение, которое больше всего соответствует вашему настроению.") - + st.title("🧪 Отладка эмоционального маппинга") + if st.session_state.round <= TOTAL_ROUNDS: - # Отображение сетки изображений + st.write(f"**Раунд {st.session_state.round} из {TOTAL_ROUNDS}**") + + # Сетка выбора с отладочной информацией cols = st.columns(3) for i, idx in enumerate(st.session_state.current_options): with cols[i % 3]: img_path = IMAGES_DIR / image_files[idx] img = Image.open(img_path) st.image(img, use_container_width=True) - if st.button(f"Выбрать {i+1}", key=f"btn_{idx}"): + + # --- ИНФОРМАЦИЯ ДЛЯ ОТЛАДКИ --- + if matcher: + v, a = matcher.predict_va(embeddings[idx]) + label_id = labels_array[idx] + label_name = EMO_NAMES.get(label_id, "Unknown") + + st.caption(f"**Класс:** {label_name}") + st.caption(f"**Прогноз:** V: {v:.2f} | A: {a:.2f}") + # ------------------------------ + + if st.button(f"Выбрать образ {i+1}", key=f"btn_{idx}", use_container_width=True): st.session_state.chosen_indices.append(idx) st.session_state.round += 1 if st.session_state.round <= TOTAL_ROUNDS: st.session_state.current_options = random.sample(range(len(image_files)), NUM_CHOICES) st.rerun() else: - # ФИНАЛЬНЫЙ ЭТАП: Анализ и Музыка - st.success("Анализ завершен! Ваш эмоциональный профиль сформирован.") + # ФИНАЛЬНЫЙ ЭТАП + st.success("✅ Профиль сформирован!") - # Расчет среднего вектора пользователя - chosen_embeddings = embeddings[st.session_state.chosen_indices] - user_vector = np.mean(chosen_embeddings, axis=0) - - # РАЗДЕЛ МУЗЫКАЛЬНЫХ РЕКОМЕНДАЦИЙ - st.divider() - st.header("🎵 Рекомендованный плейлист") + # Считаем итог + chosen_embs = embeddings[st.session_state.chosen_indices] + all_v, all_a = [], [] + for emb in chosen_embs: + v, a = matcher.predict_va(emb) + all_v.append(v) + all_a.append(a) - if matcher is None: - st.warning("База данных DEAM (music_db.csv) не найдена. Подбор музыки недоступен.") - else: - with st.spinner("Сопоставляем визуальный профиль с музыкальной базой..."): - target_v, target_a, playlist = matcher.get_playlist(user_vector, top_k=5) - - # Визуализация VA-метрик пользователя - m1, m2, m3 = st.columns(3) - m1.metric("Позитивность (Valence)", f"{target_v:.2f}", help="Шкала 1-9") - m2.metric("Энергия (Arousal)", f"{target_a:.2f}", help="Шкала 1-9") - m3.metric("Найдено треков", len(playlist)) + target_v, target_a = np.mean(all_v), np.mean(all_a) + + # Вывод результатов + col1, col2 = st.columns([2, 1]) + + with col1: + st.header("🎵 Ваш плейлист") + distances = np.sqrt((matcher.music_db['valence'] - target_v)**2 + (matcher.music_db['arousal'] - target_a)**2) + playlist = matcher.music_db.copy() + playlist['distance'] = distances + st.table(playlist.sort_values(by='distance').head(5)[['song_id', 'valence', 'arousal', 'distance']]) - # Таблица с результатами - st.subheader("Топ-5 подходящих композиций") - st.table(playlist[['song_id', 'valence', 'arousal', 'distance']]) - - st.info("💡 Вы можете найти эти треки по ID в папке audio датасета DEAM.") - - # Визуализация вектора (графики) + with col2: + st.header("📊 Профиль") + st.metric("Valence (Итог)", f"{target_v:.2f}") + st.metric("Arousal (Итог)", f"{target_a:.2f}") + + # Показываем, что именно выбрал пользователь st.divider() - st.subheader("Визуализация эмоционального вектора") - fig, ax = plt.subplots(figsize=(10, 3)) - ax.plot(user_vector[:100]) # Показываем первые 100 измерений для наглядности - ax.set_title("Эмбеддинг вашего настроения (фрагмент)") - st.pyplot(fig) + st.subheader("Выбранные вами образы и их веса:") + sum_cols = st.columns(5) + for i, idx in enumerate(st.session_state.chosen_indices): + with sum_cols[i % 5]: + v_i, a_i = matcher.predict_va(embeddings[idx]) + st.image(Image.open(IMAGES_DIR / image_files[idx]), use_container_width=True) + st.write(f"V:{v_i:.1f} A:{a_i:.1f}") if st.button("Начать заново"): - for key in list(st.session_state.keys()): - del st.session_state[key] + for key in list(st.session_state.keys()): del st.session_state[key] st.rerun() \ No newline at end of file diff --git a/src/music_engine/matcher.py b/src/music_engine/matcher.py index c789427..cefa65c 100644 --- a/src/music_engine/matcher.py +++ b/src/music_engine/matcher.py @@ -1,53 +1,52 @@ import numpy as np import pandas as pd from pathlib import Path +import joblib class MusicMatcher: - def __init__(self, db_path: Path): + 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. Загрузка обученного регрессора + if Path(model_path).exists(): + self.regressor = joblib.load(model_path) + print("✅ Регрессионная модель успешно загружена.") + else: + self.regressor = None + print("⚠️ ВНИМАНИЕ: Модель va_regressor.pkl не найдена!") def predict_va(self, embedding: np.ndarray): """ - Умный хак для демо: используем статистику вектора. - Мрачные картинки имеют меньшую среднюю активацию. + Использование обученной ML-модели (Ridge) для маппинга + эмбеддингов в пространство Valence-Arousal. """ - # 1. Средняя сила активации (прокси для Valence) - mean_act = float(np.mean(embedding)) - - # 2. Разреженность вектора (прокси для Arousal) - # Чем больше нулей или близких к нулю значений, тем "спокойнее" картинка - sparsity = float(np.mean(embedding < 0.1)) - - # Масштабируем: если средняя активация низкая (мрачное), уводим Valence вниз - # (Типичный mean_act для ResNet обычно от 0.2 до 0.8) - v = np.interp(mean_act, [0.2, 0.6], [2.0, 8.0]) - - # Если много нулей (пустота/мрак), Arousal падает. - # Если нулей мало (пестрый мусор) - Arousal растет. - a = np.interp(sparsity, [0.3, 0.8], [8.0, 2.0]) # Обратная зависимость - - # Добавляем "соль" из суммы вектора, чтобы избежать одинаковых результатов - # для похожих, но разных наборов картинок - salt = (float(np.sum(embedding)) % 2.0) - 1.0 - v += salt - a += salt - + 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) 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-норма) до треков в базе distances = np.sqrt( (self.music_db['valence'] - target_v)**2 + (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) diff --git a/src/music_engine/va_regressor.pkl b/src/music_engine/va_regressor.pkl new file mode 100644 index 0000000..3b8697e Binary files /dev/null and b/src/music_engine/va_regressor.pkl differ diff --git a/src/scripts/train_regressor.py b/src/scripts/train_regressor.py index 9179c52..9953a17 100644 --- a/src/scripts/train_regressor.py +++ b/src/scripts/train_regressor.py @@ -1,61 +1,65 @@ import numpy as np import pandas as pd from pathlib import Path -from sklearn.linear_model import Ridge +from sklearn.linear_model import RidgeCV from sklearn.multioutput import MultiOutputRegressor +from sklearn.preprocessing import StandardScaler +from sklearn.pipeline import Pipeline from sklearn.model_selection import train_test_split from sklearn.metrics import mean_squared_error, r2_score import joblib -# 1. Эталонный маппинг EmoSet -> Valence/Arousal (шкала 1-9) -# Убедись, что индексы ключей (0-7) совпадают с тем, как они размечены в твоем labels.npy -# Стандартный порядок EmoSet: +# 1. Алфавитный маппинг EmoSet EMO_VA_MAP = { - 0: (7.5, 6.5), # amusement (радость/веселье) - позитивно, средне-активно - 1: (6.5, 5.0), # awe (трепет/восхищение) - позитивно, спокойно - 2: (7.0, 3.0), # contentment (удовлетворение) - позитивно, очень спокойно - 3: (8.0, 8.0), # excitement (возбуждение) - очень позитивно, очень активно - 4: (2.0, 8.0), # anger (гнев) - негативно, очень активно - 5: (3.0, 6.0), # disgust (отвращение) - негативно, средне-активно - 6: (2.5, 7.5), # fear (страх) - негативно, очень активно - 7: (2.0, 2.0), # sadness (грусть) - негативно, пассивно + 0: (7.5, 6.5), # amusement + 1: (2.0, 8.0), # anger + 2: (6.5, 5.0), # awe + 3: (7.0, 3.0), # contentment + 4: (3.0, 6.0), # disgust + 5: (8.0, 8.0), # excitement + 6: (2.5, 7.5), # fear + 7: (2.0, 2.0), # sadness } -# 2. Загрузка данных -# Укажи пути к твоим эмбеддингам и меткам (можно взять train или test, для демо не так критично) -EMBEDDINGS_PATH = Path("../emoset_test_embeddings.npy") -LABELS_PATH = Path("../emoset_test_labels.npy") +BASE_DIR = Path(__file__).resolve().parent.parent +EMBEDDINGS_PATH = BASE_DIR / "emoset_test_embeddings.npy" +LABELS_PATH = BASE_DIR / "emoset_test_labels.npy" print("Загрузка данных...") X = np.load(EMBEDDINGS_PATH) y_labels = np.load(LABELS_PATH) -# Преобразуем дискретные метки в целевые координаты V-A -print("Формирование целевых координат (Valence, Arousal)...") y_va = np.array([EMO_VA_MAP[label] for label in y_labels]) - -# Разделение на train/val X_train, X_test, y_train, y_test = train_test_split(X, y_va, test_size=0.2, random_state=42) -# 3. Обучение модели -print("Обучение Ridge регрессора...") -# Ridge отлично справляется с многомерными эмбеддингами, избегая переобучения -base_estimator = Ridge(alpha=1.0) -model = MultiOutputRegressor(base_estimator) +# 2. НОВАЯ, ПРАВИЛЬНАЯ АРХИТЕКТУРА (Pipeline) +print("Обучение масштабатора и RidgeCV регрессора...") +# Pipeline гарантирует, что при предсказании в main.py новые векторы тоже будут масштабированы +model = Pipeline([ + ('scaler', StandardScaler()), + ('regressor', MultiOutputRegressor(RidgeCV(alphas=[0.1, 1.0, 10.0, 100.0, 1000.0]))) +]) model.fit(X_train, y_train) -# 4. Оценка +# 3. Диагностика и Оценка y_pred = model.predict(X_test) + mse = mean_squared_error(y_test, y_pred) r2 = r2_score(y_test, y_pred) -print(f"Обучение завершено!") -print(f"MSE (Среднеквадратичная ошибка): {mse:.4f}") -print(f"R^2 Score (Коэффициент детерминации): {r2:.4f}") +print(f"\n[УСПЕХ] Обучение завершено!") +print(f"MSE: {mse:.4f}") +print(f"R^2 Score: {r2:.4f}") -# 5. Сохранение -output_model_path = Path("../src/music_engine/va_regressor.pkl") +# === ТОТ САМЫЙ ТЕСТ НА КОЛЛАПС === +print("\n--- ДИАГНОСТИКА РАЗБРОСА ПРЕДСКАЗАНИЙ ---") +print(f"Valence: от {y_pred[:, 0].min():.2f} до {y_pred[:, 0].max():.2f} (Эталон: 2.0 - 8.0)") +print(f"Arousal: от {y_pred[:, 1].min():.2f} до {y_pred[:, 1].max():.2f} (Эталон: 2.0 - 8.0)") +# =============================================== + +# 4. Сохранение (Pipeline сохраняется целиком со StandardScaler) +output_model_path = BASE_DIR / "music_engine" / "va_regressor.pkl" output_model_path.parent.mkdir(parents=True, exist_ok=True) joblib.dump(model, output_model_path) -print(f"Модель сохранена в: {output_model_path}") \ No newline at end of file +print(f"\nМодель сохранена в: {output_model_path}") \ No newline at end of file