diff --git a/src/acc_images_model.ipynb b/src/scripts/acc_images_model.ipynb similarity index 100% rename from src/acc_images_model.ipynb rename to src/scripts/acc_images_model.ipynb diff --git a/src/download_DEAM.ipynb b/src/scripts/download_DEAM.ipynb similarity index 100% rename from src/download_DEAM.ipynb rename to src/scripts/download_DEAM.ipynb diff --git a/src/extract_embeddings.ipynb b/src/scripts/extract_embeddings.ipynb similarity index 100% rename from src/extract_embeddings.ipynb rename to src/scripts/extract_embeddings.ipynb diff --git a/src/prep_deam.ipynb b/src/scripts/prep_deam.ipynb similarity index 100% rename from src/prep_deam.ipynb rename to src/scripts/prep_deam.ipynb diff --git a/src/prepare_dataset_to_train.ipynb b/src/scripts/prepare_dataset_to_train.ipynb similarity index 100% rename from src/prepare_dataset_to_train.ipynb rename to src/scripts/prepare_dataset_to_train.ipynb diff --git a/src/test.ipynb b/src/scripts/test.ipynb similarity index 100% rename from src/test.ipynb rename to src/scripts/test.ipynb diff --git a/src/train_images.ipynb b/src/scripts/train_images.ipynb similarity index 100% rename from src/train_images.ipynb rename to src/scripts/train_images.ipynb diff --git a/src/scripts/train_regressor.py b/src/scripts/train_regressor.py new file mode 100644 index 0000000..9179c52 --- /dev/null +++ b/src/scripts/train_regressor.py @@ -0,0 +1,61 @@ +import numpy as np +import pandas as pd +from pathlib import Path +from sklearn.linear_model import Ridge +from sklearn.multioutput import MultiOutputRegressor +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: +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 (грусть) - негативно, пассивно +} + +# 2. Загрузка данных +# Укажи пути к твоим эмбеддингам и меткам (можно взять train или test, для демо не так критично) +EMBEDDINGS_PATH = Path("../emoset_test_embeddings.npy") +LABELS_PATH = Path("../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) + +model.fit(X_train, y_train) + +# 4. Оценка +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}") + +# 5. Сохранение +output_model_path = Path("../src/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