148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
import argparse
|
|
import csv
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import joblib
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
import seaborn as sns
|
|
|
|
from xgboost import XGBRegressor
|
|
from sklearn.preprocessing import StandardScaler
|
|
from sklearn.metrics import mean_absolute_error
|
|
|
|
# =========================================================
|
|
# Hyper Parameters
|
|
# =========================================================
|
|
|
|
MAX_NOTES = 2000 # 분석할 최대 노트 수
|
|
FACTOR_COUNT = 4 # [type, bpm, scroll, delta]
|
|
INPUT_DIM = MAX_NOTES * FACTOR_COUNT
|
|
|
|
TRAIN_SIZE = 0
|
|
VALID_SIZE = 0
|
|
RANDOM_STATE = 42
|
|
|
|
N_ESTIMATORS = 500
|
|
MAX_DEPTH = 6
|
|
LEARNING_RATE = 0.05
|
|
CONTINUE_TRAINING = True
|
|
ERROR_TOLERANCE = 0.1
|
|
|
|
# =========================================================
|
|
# 파일명
|
|
# =========================================================
|
|
|
|
FACTORS_FILENAME = "factors.json"
|
|
MEASURE_FILENAME = "measure.csv"
|
|
MODEL_FILENAME = "model.pkl"
|
|
SCALER_FILENAME = "scaler.pkl"
|
|
|
|
def safe_float(value):
|
|
if value is None: return 0.0
|
|
x = float(value)
|
|
return x if math.isfinite(x) else 0.0
|
|
|
|
def train_model(working_dir: str, data_dir: str):
|
|
random.seed(RANDOM_STATE)
|
|
|
|
factors_path = os.path.join(working_dir, FACTORS_FILENAME)
|
|
measure_path = os.path.join(data_dir, MEASURE_FILENAME)
|
|
model_path = os.path.join(working_dir, MODEL_FILENAME)
|
|
scaler_path = os.path.join(working_dir, SCALER_FILENAME)
|
|
|
|
with open(factors_path, "r", encoding="utf-8") as f:
|
|
factor_data = json.load(f)
|
|
|
|
# feature_map build: key -> list of factors
|
|
feature_map = {(str(item["songno"]), str(item["difficulty"])): item["factors"] for item in factor_data}
|
|
|
|
dataset = []
|
|
with open(measure_path, "r", encoding="utf-8") as f:
|
|
reader = csv.reader(f)
|
|
next(reader, None)
|
|
for row in reader:
|
|
if len(row) < 3: continue
|
|
measure, songno, diff = safe_float(row[0]), str(row[1]), str(row[2])
|
|
key = (songno, diff)
|
|
|
|
if key in feature_map:
|
|
raw_factors = feature_map[key]
|
|
# 고정 길이 벡터로 변환 (Padding or Truncating)
|
|
vector = np.zeros(INPUT_DIM, dtype=np.float32)
|
|
for i in range(min(len(raw_factors), MAX_NOTES)):
|
|
for j in range(FACTOR_COUNT):
|
|
vector[i * FACTOR_COUNT + j] = safe_float(raw_factors[i][j])
|
|
|
|
dataset.append((vector, measure, songno, diff))
|
|
|
|
random.shuffle(dataset)
|
|
if len(dataset) < (TRAIN_SIZE + VALID_SIZE):
|
|
raise ValueError(f"Dataset size {len(dataset)} < required {TRAIN_SIZE + VALID_SIZE}")
|
|
|
|
train_ds = dataset[:TRAIN_SIZE]
|
|
valid_ds = dataset[TRAIN_SIZE:TRAIN_SIZE + VALID_SIZE]
|
|
|
|
X_train = np.array([x for x, _, _, _ in train_ds])
|
|
y_train = np.array([y for _, y, _, _ in train_ds])
|
|
X_valid = np.array([x for x, _, _, _ in valid_ds])
|
|
y_valid = np.array([y for _, y, _, _ in valid_ds])
|
|
valid_info = [(s, d) for _, _, s, d in valid_ds]
|
|
|
|
if CONTINUE_TRAINING and os.path.exists(scaler_path):
|
|
scaler = joblib.load(scaler_path)
|
|
else:
|
|
scaler = StandardScaler()
|
|
scaler.fit(X_train)
|
|
|
|
X_train = scaler.transform(X_train)
|
|
X_valid = scaler.transform(X_valid)
|
|
|
|
if CONTINUE_TRAINING and os.path.exists(model_path):
|
|
model = joblib.load(model_path)
|
|
model.fit(X_train, y_train, xgb_model=model.get_booster())
|
|
else:
|
|
model = XGBRegressor(n_estimators=N_ESTIMATORS, max_depth=MAX_DEPTH, learning_rate=LEARNING_RATE, objective="reg:squarederror", random_state=RANDOM_STATE)
|
|
model.fit(X_train, y_train)
|
|
|
|
pred = model.predict(X_valid)
|
|
mae = mean_absolute_error(y_valid, pred)
|
|
accuracy = np.sum(np.abs(pred - y_valid) <= ERROR_TOLERANCE) / len(y_valid)
|
|
print(f"MAE: {mae:.4f} | Accuracy: {accuracy:.4f}")
|
|
|
|
# Results save
|
|
validate_details = []
|
|
for i in range(len(y_valid)):
|
|
validate_details.append({"songno": valid_info[i][0], "diff": valid_info[i][1], "actual": float(y_valid[i]), "predicted": float(pred[i]), "error": float(y_valid[i] - pred[i])})
|
|
|
|
validate_details.sort(key=lambda x: abs(x["error"]), reverse=True)
|
|
with open(os.path.join(working_dir, "validate.json"), "w", encoding="utf-8") as f:
|
|
json.dump({"summary": {"mae": float(mae), "accuracy": float(accuracy)}, "details": validate_details}, f, indent=2)
|
|
|
|
# Plot
|
|
plt.switch_backend('Agg')
|
|
df_plot = pd.DataFrame(validate_details)
|
|
df_plot['abs_error'] = df_plot['error'].abs()
|
|
df_plot = df_plot.sort_values('abs_error', ascending=False).reset_index(drop=True)
|
|
plt.figure(figsize=(12, 6))
|
|
sns.scatterplot(data=df_plot, x=df_plot.index, y='abs_error', color='crimson')
|
|
plt.axhline(0.2, color='green', linestyle='--')
|
|
plt.ylim(0, 4)
|
|
plt.savefig(os.path.join(working_dir, "validate.png"))
|
|
|
|
joblib.dump(model, model_path)
|
|
joblib.dump(scaler, scaler_path)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--workingDir", required=True)
|
|
parser.add_argument("--dataDir", required=True)
|
|
parser.add_argument("--trainSize", required=True, type=int)
|
|
parser.add_argument("--validSize", required=True, type=int)
|
|
args = parser.parse_args()
|
|
TRAIN_SIZE, VALID_SIZE = args.trainSize, args.validSize
|
|
train_model(args.workingDir, args.dataDir)
|