diff --git a/compare.ts b/.old/compare.ts similarity index 100% rename from compare.ts rename to .old/compare.ts diff --git a/script/extract.ts b/.old/extract.ts similarity index 100% rename from script/extract.ts rename to .old/extract.ts diff --git a/predict-all.ts b/.old/predict-all.ts similarity index 100% rename from predict-all.ts rename to .old/predict-all.ts diff --git a/predict.sh b/.old/predict.sh similarity index 100% rename from predict.sh rename to .old/predict.sh diff --git a/script/predict.ts b/.old/predict.ts similarity index 100% rename from script/predict.ts rename to .old/predict.ts diff --git a/spec.md b/.old/spec.md similarity index 100% rename from spec.md rename to .old/spec.md diff --git a/t.sh b/.old/t.sh similarity index 100% rename from t.sh rename to .old/t.sh diff --git a/test/compare.json b/.old/test/compare.json similarity index 100% rename from test/compare.json rename to .old/test/compare.json diff --git a/test/factors.json b/.old/test/factors.json similarity index 100% rename from test/factors.json rename to .old/test/factors.json diff --git a/test/model.pth b/.old/test/model.pth similarity index 100% rename from test/model.pth rename to .old/test/model.pth diff --git a/test/result.json b/.old/test/result.json similarity index 100% rename from test/result.json rename to .old/test/result.json diff --git a/test/scaler.json b/.old/test/scaler.json similarity index 100% rename from test/scaler.json rename to .old/test/scaler.json diff --git a/test/validate_result_1.csv b/.old/test/validate_result_1.csv similarity index 100% rename from test/validate_result_1.csv rename to .old/test/validate_result_1.csv diff --git a/test/validate_result_10.csv b/.old/test/validate_result_10.csv similarity index 100% rename from test/validate_result_10.csv rename to .old/test/validate_result_10.csv diff --git a/test/validate_result_2.csv b/.old/test/validate_result_2.csv similarity index 100% rename from test/validate_result_2.csv rename to .old/test/validate_result_2.csv diff --git a/test/validate_result_3.csv b/.old/test/validate_result_3.csv similarity index 100% rename from test/validate_result_3.csv rename to .old/test/validate_result_3.csv diff --git a/test/validate_result_4.csv b/.old/test/validate_result_4.csv similarity index 100% rename from test/validate_result_4.csv rename to .old/test/validate_result_4.csv diff --git a/test/validate_result_5.csv b/.old/test/validate_result_5.csv similarity index 100% rename from test/validate_result_5.csv rename to .old/test/validate_result_5.csv diff --git a/test/validate_result_6.csv b/.old/test/validate_result_6.csv similarity index 100% rename from test/validate_result_6.csv rename to .old/test/validate_result_6.csv diff --git a/test/validate_result_7.csv b/.old/test/validate_result_7.csv similarity index 100% rename from test/validate_result_7.csv rename to .old/test/validate_result_7.csv diff --git a/test/validate_result_8.csv b/.old/test/validate_result_8.csv similarity index 100% rename from test/validate_result_8.csv rename to .old/test/validate_result_8.csv diff --git a/test/validate_result_9.csv b/.old/test/validate_result_9.csv similarity index 100% rename from test/validate_result_9.csv rename to .old/test/validate_result_9.csv diff --git a/script/train.py b/.old/train.py similarity index 100% rename from script/train.py rename to .old/train.py diff --git a/train.sh b/.old/train.sh similarity index 100% rename from train.sh rename to .old/train.sh diff --git a/script/train.ts b/.old/train.ts similarity index 100% rename from script/train.ts rename to .old/train.ts diff --git a/bun.lock b/bun.lock index 285f7af..4ef8732 100644 --- a/bun.lock +++ b/bun.lock @@ -8,6 +8,7 @@ "tja-parser": "^0.2.9", }, "devDependencies": { + "@types/bun": "^1.3.13", "@types/iconv-lite": "^0.0.1", "@types/node": "^25.6.0", "typescript": "^6.0.3", @@ -17,10 +18,14 @@ "packages": { "@babel/runtime": ["@babel/runtime@7.29.2", "", {}, "sha512-JiDShH45zKHWyGe4ZNVRrCjBz8Nh9TMmZG1kh4QTK8hCBTWBi8Da+i7s1fJw7/lYpM4ccepSNfqzZ/QvABBi5g=="], + "@types/bun": ["@types/bun@1.3.13", "", { "dependencies": { "bun-types": "1.3.13" } }, "sha512-9fqXWk5YIHGGnUau9TEi+qdlTYDAnOj+xLCmSTwXfAIqXr2x4tytJb43E9uCvt09zJURKXwAtkoH4nLQfzeTXw=="], + "@types/iconv-lite": ["@types/iconv-lite@0.0.1", "", { "dependencies": { "@types/node": "*" } }, "sha512-SsRBQxGw7/2/NxYJfBdiUx5a7Ms/voaUhOO9u2y9FTeTNBO1PXohzE4i3JfD8q2Te42HLTn5pyZtDf8j1bPKgQ=="], "@types/node": ["@types/node@25.6.0", "", { "dependencies": { "undici-types": "~7.19.0" } }, "sha512-+qIYRKdNYJwY3vRCZMdJbPLJAtGjQBudzZzdzwQYkEPQd+PJGixUL5QfvCLDaULoLv+RhT3LDkwEfKaAkgSmNQ=="], + "bun-types": ["bun-types@1.3.13", "", { "dependencies": { "@types/node": "*" } }, "sha512-QXKeHLlOLqQX9LgYaHJfzdBaV21T63HhFJnvuRCcjZiaUDpbs5ED1MgxbMra71CsryN/1dAoXuJJJwIv/2drVA=="], + "complex.js": ["complex.js@2.4.3", "", {}, "sha512-UrQVSUur14tNX6tiP4y8T4w4FeJAX3bi2cIv0pu/DTLFNxoq7z2Yh83Vfzztj6Px3X/lubqQ9IrPp7Bpn6p4MQ=="], "csv-parse": ["csv-parse@6.2.1", "", {}, "sha512-LRLMV+UCyfMokp8Wb411duBf1gaBKJfOfBWU9eHMJ+b+cJYZsNu3AFmjJf3+yPGd59Exz1TsMjaSFyxnYB9+IQ=="], diff --git a/docs/1. tja factorize.md b/docs/1. tja factorize.md new file mode 100644 index 0000000..42946a6 --- /dev/null +++ b/docs/1. tja factorize.md @@ -0,0 +1,12 @@ +# TJA factorize +TJA를 DNN에 사용하기 위해 factorize한다. + +각 노트를 다음과 같이 변환한다. +- [Note type, BPM, Scroll, Timing] +- Note type은 0과 1만 사용한다. + - Don -> 0 + - Ka -> 1 +- BPM은 100을 나누어 사용한다. +- Scroll은 5를 나누어 사용한다. +- Timing은 이전 노트와의 시간 차이를 사용한다. (단위 s) + diff --git a/docs/2. features.md b/docs/2. features.md new file mode 100644 index 0000000..be3e058 --- /dev/null +++ b/docs/2. features.md @@ -0,0 +1,40 @@ +# Feature +DNN 분석을 위해 채보의 특성을 추출한다. + +## Average density +평균 밀도를 나타낸다. + +$\frac{노트 수}{마지막 노트 타이밍 - 첫 노트 타이밍}$ + +## Peak density +최고 밀도를 나타낸다. + +어떤 노트를 기준으로 앞으로 1초내에 있는 노트들의 개수 중 최대값으로 구한다. + +## Average BPM +평균 BPM을 나타낸다. + +$\frac{\Sigma BPM}{노트 수}$ + +## Average BPM 2 +평균 BPM을 나타내나, 다른 방식으로 구한다. +극소수의 노트만 BPM이 다를 경우 평균 BPM에 미치는 영향이 클 수 있기 때문에, BPM의 제곱을 총합하여 평균을 구한다. + +$\sqrt{\frac{\Sigma BPM^2}{노트수}}$ + +## BPM Change +BPM 변화 횟수를 나타낸다. BPM 흔들림을 제외하기 위해, 이전 노트와의 BPM차이가 1.5 이상일 떄 1 증가한다. + +## Scroll Change +스크롤 변화 횟수를 나타낸다. 노트의 $BPM \times Scroll$의 값이 이전 노트와 1.5 이상 차이날 때 1 증가한다. + +## Rhythm Complexity +i번째 노트와 i-1번쨰 노트의 간격이 i-1번쨰 노트와 i-2번째 노트의 간격의 비율이 2의 거듭제곱이 아닐 때 1 증가한다. + +## Color Complexity +i번째 노트와 i-2번쨰 노트가 다른 종류일 때 증가하며, 간격의 제곱의 역수에 비례한다. + +$\frac{\mathrm{color\ changed\ ?\ 1\ :\ 0}}{\Delta t^2}$ + +## Note Count +노트의 개수 \ No newline at end of file diff --git a/package.json b/package.json index b216a5e..1c8216f 100644 --- a/package.json +++ b/package.json @@ -5,6 +5,7 @@ "tja-parser": "^0.2.9" }, "devDependencies": { + "@types/bun": "^1.3.13", "@types/iconv-lite": "^0.0.1", "@types/node": "^25.6.0", "typescript": "^6.0.3" diff --git a/predict/predict_xgboost.py b/predict/predict_xgboost.py new file mode 100644 index 0000000..9d1a21a --- /dev/null +++ b/predict/predict_xgboost.py @@ -0,0 +1,155 @@ +import argparse +import json +import math +import os +import joblib +import numpy as np + + +# ========================================================= +# 파일명 +# ========================================================= + +FEATURES_FILENAME = "features.json" + +MODEL_FILENAME = "model.pkl" +SCALER_FILENAME = "scaler.pkl" +FEATURE_NAMES_FILENAME = "features.txt" + + +# ========================================================= +# safe float +# ========================================================= + +def safe_float(value): + if value is None: + return 0.0 + + x = float(value) + + if not math.isfinite(x): + return 0.0 + + return x + + +# ========================================================= +# 예측 함수 +# ========================================================= + +def predict( + working_dir: str, + songno: str +): + # ===================================================== + # 경로 + # ===================================================== + + features_path = os.path.join( + working_dir, + FEATURES_FILENAME + ) + + model_path = os.path.join( + working_dir, + MODEL_FILENAME + ) + + scaler_path = os.path.join( + working_dir, + SCALER_FILENAME + ) + + feature_names_path = os.path.join( + working_dir, + FEATURE_NAMES_FILENAME + ) + + # ===================================================== + # 모델 로드 + # ===================================================== + + model = joblib.load(model_path) + scaler = joblib.load(scaler_path) + + # ===================================================== + # feature 이름 로드 + # ===================================================== + + with open(feature_names_path, "r", encoding="utf-8") as f: + feature_names = [ + line.strip() + for line in f.readlines() + if line.strip() + ] + + # ===================================================== + # features.json 로드 + # ===================================================== + + with open(features_path, "r", encoding="utf-8") as f: + data = json.load(f) + + # ===================================================== + # target 찾기 + # ===================================================== + + targets = [] + + for item in data: + if str(item["songno"]) == str(songno): + targets.append(item) + + if len(targets) == 0: + raise ValueError(f"Chart not found: songno={songno}") + + # ===================================================== + # feature vector 생성 + # ===================================================== + + row = [] + + for target in targets: + row = [] + + for k in feature_names: + value = target.get(k, 0) + row.append(safe_float(value)) + + X = np.array([row], dtype=np.float32) + + X = scaler.transform(X) + + pred = model.predict(X)[0] + + diff = target.get("difficulty", "unknown") + + print( + f"{diff:10} " + f"{pred:.1f}" + ) + + +# ========================================================= +# main +# ========================================================= + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--workingDir", + required=True + ) + + parser.add_argument( + "--songno", + required=True + ) + + args = parser.parse_args() + + predict( + args.workingDir, + args.songno + ) \ No newline at end of file diff --git a/preprocess/factorize.ts b/preprocess/factorize.ts new file mode 100644 index 0000000..b919ac6 --- /dev/null +++ b/preprocess/factorize.ts @@ -0,0 +1,18 @@ +import { Course } from "tja-parser"; +import { filterHitNotes } from "./util"; + +export type NoteFactor = [type: 0 | 1, bpm: number, scroll: number, delta: number] +export function factorize(course: Course) { + const hitNotes = filterHitNotes(course); + const factors: NoteFactor[] = []; + for (let i = 0; i < hitNotes.length; i++) { + const note = hitNotes[i]; + factors.push([ + (note.type === 1 || note.type === 3) ? 0 : 1, + note.getBPM() / 100, + note.getScroll() / 5, + i === 0 ? 0 : (note.getTimingMS() - hitNotes[i - 1].getTimingMS()) / 1000 + ]) + } + return factors; +} \ No newline at end of file diff --git a/script/factorize.ts b/preprocess/featurize.ts similarity index 66% rename from script/factorize.ts rename to preprocess/featurize.ts index b3b83ab..5f8938b 100644 --- a/script/factorize.ts +++ b/preprocess/featurize.ts @@ -1,6 +1,6 @@ import { Course, Bar, Note, HitNote } from 'tja-parser'; -export namespace Factor { +export namespace Feature { export function filterHitNotes(course: Course) { const notes: HitNote[] = []; course.noteGroups.forEach((g) => { @@ -31,6 +31,9 @@ export namespace Factor { // 밀도 관련 export function getAverageDensity(notes: HitNote[]) { + if (notes.length === 0) { + return 0; + } return notes.length / (notes[notes.length - 1].getTimingMS() - notes[0].getTimingMS()) * 1000 } @@ -50,7 +53,7 @@ export namespace Factor { // BPM 관련 export function getAverageBPM(notes: HitNote[]) { - + if(notes.length === 0) return 0; const averageBPM = notes.reduce((p, note) => p + note.getBPM().valueOf(), 0) / notes.length; return averageBPM; } @@ -65,29 +68,6 @@ export namespace Factor { return bpmChange; } - // 복잡성 - export function getComplexity(notes: HitNote[]) { - let complexity = 0; - for (let i = 2; i < notes.length; i++) { - let localComplexity = 0; - - // ddk 또는 dkk류면 1, 아니면 0.5 - if ( - (notes[i].type % 2 === notes[i - 1].type % 2 && notes[i - 1].type % 2 !== notes[i - 2].type % 2) || - (notes[i].type % 2 !== notes[i - 1].type % 2 && notes[i - 1].type % 2 === notes[i - 2].type % 2) - ) { - localComplexity = 1; - } else { - localComplexity = 0.5 - } - - // 시간 차가 짧을 수록 complexity 증가 - localComplexity *= (1 / (notes[i].getTimingMS() - notes[i - 1].getTimingMS()) + 1 / (notes[i - 1].getTimingMS() - notes[i - 2].getTimingMS())) - complexity += localComplexity; - } - return complexity / notes.length; - } - // 스크롤 변화 export function getScrollChange(notes: HitNote[]) { let bpmChange = 0; @@ -98,17 +78,44 @@ export namespace Factor { }; return bpmChange; } + + export function getRhythmComplexity(notes: HitNote[]) { + let complexity = 0; + for (let i = 2; i < notes.length; i++) { + const d1 = notes[i].getTimingMS() - notes[i - 1].getTimingMS() + const d2 = notes[i - 1].getTimingMS() - notes[i - 2].getTimingMS() + + const ratio = d1 / d2; + const log = Math.log2(ratio); + + if (Math.abs(log - Math.round(log)) < 1e-3) { + complexity++; + } + } + return complexity; + } + + export function getColorComplexity(notes: HitNote[]) { + let complexity = 0; + for (let i = 2; i < notes.length; i++) { + if (notes[i].type % 2 != notes[i - 2].type % 2) { + complexity += 1 / ((notes[i].getTimingMS() - notes[i - 2].getTimingMS()) ** 2) + } + } + return complexity; + } } -export function factorize(course: Course) { - const notes = Factor.filterHitNotes(course) +export function featurize(course: Course) { + const notes = Feature.filterHitNotes(course) return { note_count: notes.length, - density_avg: Factor.getAverageDensity(notes), - density_peak: Factor.getPeakDensity(notes), - bpm_avg: Factor.getAverageBPM(notes), - bpm_change: Factor.getBpmChange(notes), - complexity: Factor.getComplexity(notes), - scroll_change: Factor.getScrollChange(notes) + density_avg: Feature.getAverageDensity(notes), + density_peak: Feature.getPeakDensity(notes), + bpm_avg: Feature.getAverageBPM(notes), + bpm_change: Feature.getBpmChange(notes), + scroll_change: Feature.getScrollChange(notes), + rhythm_complexity: Feature.getRhythmComplexity(notes), + color_complexity: Feature.getColorComplexity(notes) }; } \ No newline at end of file diff --git a/script/parse.ts b/preprocess/parse.ts similarity index 100% rename from script/parse.ts rename to preprocess/parse.ts diff --git a/preprocess/util.ts b/preprocess/util.ts new file mode 100644 index 0000000..f8e51d1 --- /dev/null +++ b/preprocess/util.ts @@ -0,0 +1,15 @@ +import { Bar, Course, HitNote } from "tja-parser"; + +export function filterHitNotes(course: Course): HitNote[] { + const notes: HitNote[] = []; + course.noteGroups.forEach((g) => { + if (g instanceof Bar) { + for (const note of g.getNotes()) { + if (note instanceof HitNote) { + notes.push(note); + } + } + } + }); + return notes; +} \ No newline at end of file diff --git a/script/preprocess.ts b/script/preprocess.ts new file mode 100644 index 0000000..a682191 --- /dev/null +++ b/script/preprocess.ts @@ -0,0 +1,65 @@ +import Bun from 'bun'; +import path from 'node:path'; +import { parseArgs } from 'node:util'; +import fs, { mkdirSync } from 'node:fs'; +import { Song } from 'tja-parser'; +import { featurize } from '../preprocess/featurize'; +import { parseTja } from '../preprocess/parse' + +const { values } = parseArgs({ + args: Bun.argv, + options: { + outputDir: { + type: "string" + }, + dataDir: { + type: "string" + } + }, + strict: true, + allowPositionals: true, +}) + +if (!values.dataDir || !values.outputDir) { + console.error("--outputDir --dataDir"); + process.exit(1); +} + +const outputDir = values.outputDir ?? ''; +if (!fs.existsSync(outputDir)) mkdirSync(outputDir) +const dataDir = values.dataDir ?? ''; + +const tjaDir = path.join(dataDir, 'tja'); +const files = fs.readdirSync(tjaDir); + +const features: ({ songno: string, difficulty: 'oni' | 'ura' } & {})[] = []; +for (const file of files) { + const tja = fs.readFileSync(path.join(tjaDir, file), 'utf-8'); + const songno = path.basename(file, '.tja'); + try { + const parsed = parseTja(tja); + const oni = parsed?.oni; + const edit = parsed?.edit; + if (oni) { + features.push({ + songno, + difficulty: 'oni', + ...featurize(oni) + }) + } + if (edit) { + features.push({ + songno, + difficulty: 'ura', + ...featurize(edit) + }) + } + } + catch (err) { + console.error(err); + console.error(file); + } +} + +const featurePath = path.join(outputDir, 'features.json'); +fs.writeFileSync(featurePath, JSON.stringify(features, null, 2), 'utf-8'); \ No newline at end of file diff --git a/train/train_xgboost.py b/train/train_xgboost.py new file mode 100644 index 0000000..c30edb3 --- /dev/null +++ b/train/train_xgboost.py @@ -0,0 +1,313 @@ +import argparse +import csv +import json +import math +import os +import joblib +import numpy as np + +from xgboost import XGBRegressor +from sklearn.model_selection import train_test_split +from sklearn.preprocessing import StandardScaler +from sklearn.metrics import mean_absolute_error + +# ========================================================= +# Hyper Parameters +# ========================================================= + +TEST_SIZE = 0.2 +RANDOM_STATE = 42 + +N_ESTIMATORS = 500 +MAX_DEPTH = 6 +LEARNING_RATE = 0.05 +SUBSAMPLE = 0.8 +COLSAMPLE_BYTREE = 0.8 + +CONTINUE_TRAINING = True + +# 예측 성공으로 간주할 허용 오차 +ERROR_TOLERANCE = 0.2 + +# ========================================================= +# 파일명 +# ========================================================= + +FEATURES_FILENAME = "features.json" +MEASURE_FILENAME = "measure.csv" + +MODEL_FILENAME = "model.pkl" +SCALER_FILENAME = "scaler.pkl" +FEATURE_NAMES_FILENAME = "features.txt" + +# ========================================================= +# 무시할 key +# ========================================================= + +IGNORE_KEYS = { + "songno", + "difficulty" +} + + +def safe_float(value): + if value is None: + return 0.0 + + x = float(value) + + if not math.isfinite(x): + return 0.0 + + return x + + +def train_model( + working_dir: str, + data_dir: str +): + # ===================================================== + # path + # ===================================================== + + features_path = os.path.join( + working_dir, + FEATURES_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 + ) + + feature_names_path = os.path.join( + working_dir, + FEATURE_NAMES_FILENAME + ) + + # ===================================================== + # features.json + # ===================================================== + + with open(features_path, "r", encoding="utf-8") as f: + feature_data = json.load(f) + + if len(feature_data) == 0: + raise ValueError("features.json is empty") + + # ===================================================== + # feature map + # ===================================================== + + feature_map = {} + + for item in feature_data: + key = ( + str(item["songno"]), + str(item["difficulty"]) + ) + + feature_map[key] = item + + # ===================================================== + # feature names + # ===================================================== + + feature_names = sorted([ + k for k in feature_data[0].keys() + if k not in IGNORE_KEYS + ]) + + # ===================================================== + # measure.csv + # ===================================================== + + X = [] + y = [] + + 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 = safe_float(row[0]) + songno = str(row[1]) + diff = str(row[2]) + + key = (songno, diff) + + if key not in feature_map: + print( + f"[WARN] feature not found: " + f"{songno} {diff}" + ) + continue + + feature_item = feature_map[key] + + features = [ + safe_float(feature_item.get(k, 0)) + for k in feature_names + ] + + X.append(features) + y.append(measure) + + if len(X) == 0: + raise ValueError("No training data") + + X = np.array(X, dtype=np.float32) + y = np.array(y, dtype=np.float32) + + print(f"Dataset Size: {len(X)}") + print(f"Feature Count: {len(feature_names)}") + + # ===================================================== + # split + # ===================================================== + + X_train, X_valid, y_train, y_valid = train_test_split( + X, + y, + test_size=TEST_SIZE, + random_state=RANDOM_STATE + ) + + # ===================================================== + # scaler + # ===================================================== + + if CONTINUE_TRAINING and os.path.exists(scaler_path): + print("Loading existing scaler...") + + scaler = joblib.load(scaler_path) + + else: + print("Creating new scaler...") + + scaler = StandardScaler() + scaler.fit(X_train) + + X_train = scaler.transform(X_train) + X_valid = scaler.transform(X_valid) + + # ===================================================== + # model + # ===================================================== + + if CONTINUE_TRAINING and os.path.exists(model_path): + print("Loading existing model...") + + model = joblib.load(model_path) + + previous_booster = model.get_booster() + + model.fit( + X_train, + y_train, + xgb_model=previous_booster + ) + + else: + print("Creating new model...") + + model = XGBRegressor( + n_estimators=N_ESTIMATORS, + max_depth=MAX_DEPTH, + learning_rate=LEARNING_RATE, + subsample=SUBSAMPLE, + colsample_bytree=COLSAMPLE_BYTREE, + objective="reg:squarederror", + random_state=RANDOM_STATE + ) + + model.fit(X_train, y_train) + + # ===================================================== + # evaluate + # ===================================================== + + pred = model.predict(X_valid) + + mae = mean_absolute_error(y_valid, pred) + + correct = np.sum( + np.abs(pred - y_valid) <= ERROR_TOLERANCE + ) + + accuracy = correct / len(y_valid) + + print(f"\nMAE: {mae:.4f}") + print( + f"Accuracy " + f"(±{ERROR_TOLERANCE}): " + f"{accuracy:.4f}" + ) + + # ===================================================== + # feature importance + # ===================================================== + + print("\nFeature Importance:") + + importance = model.feature_importances_ + + pairs = list(zip(feature_names, importance)) + pairs.sort(key=lambda x: x[1], reverse=True) + + for name, score in pairs: + print(f"{name:25} {score:.6f}") + + # ===================================================== + # save + # ===================================================== + + joblib.dump(model, model_path) + joblib.dump(scaler, scaler_path) + + with open(feature_names_path, "w", encoding="utf-8") as f: + for name in feature_names: + f.write(name + "\n") + + print("\nSaved:") + print(model_path) + print(scaler_path) + print(feature_names_path) + + +# ========================================================= +# main +# ========================================================= + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--workingDir", + required=True + ) + + parser.add_argument( + "--dataDir", + required=True + ) + + args = parser.parse_args() + + train_model( + args.workingDir, + args.dataDir + ) \ No newline at end of file diff --git a/train/train_xgboost_pick.py b/train/train_xgboost_pick.py new file mode 100644 index 0000000..028cfe3 --- /dev/null +++ b/train/train_xgboost_pick.py @@ -0,0 +1,369 @@ +import argparse +import csv +import json +import math +import os +import random +import joblib +import numpy as np + +from xgboost import XGBRegressor +from sklearn.preprocessing import StandardScaler +from sklearn.metrics import mean_absolute_error + +# ========================================================= +# Hyper Parameters +# ========================================================= + +TRAIN_SIZE = 0 +VALID_SIZE = 0 + +RANDOM_STATE = 42 + +N_ESTIMATORS = 500 +MAX_DEPTH = 6 +LEARNING_RATE = 0.05 +SUBSAMPLE = 0.8 +COLSAMPLE_BYTREE = 0.8 + +CONTINUE_TRAINING = True + +# 예측 성공으로 간주할 허용 오차 +ERROR_TOLERANCE = 0.5 + +# ========================================================= +# 파일명 +# ========================================================= + +FEATURES_FILENAME = "features.json" +MEASURE_FILENAME = "measure.csv" + +MODEL_FILENAME = "model.pkl" +SCALER_FILENAME = "scaler.pkl" +FEATURE_NAMES_FILENAME = "features.txt" + +# ========================================================= +# 무시할 key +# ========================================================= + +IGNORE_KEYS = { + "songno", + "difficulty" +} + + +# ========================================================= +# safe float +# ========================================================= + +def safe_float(value): + if value is None: + return 0.0 + + x = float(value) + + if not math.isfinite(x): + return 0.0 + + return x + + +# ========================================================= +# train +# ========================================================= + +def train_model( + working_dir: str, + data_dir: str +): + random.seed(RANDOM_STATE) + + # ===================================================== + # path + # ===================================================== + + features_path = os.path.join( + working_dir, + FEATURES_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 + ) + + feature_names_path = os.path.join( + working_dir, + FEATURE_NAMES_FILENAME + ) + + # ===================================================== + # features.json + # ===================================================== + + with open(features_path, "r", encoding="utf-8") as f: + feature_data = json.load(f) + + if len(feature_data) == 0: + raise ValueError("features.json is empty") + + # ===================================================== + # feature map + # ===================================================== + + feature_map = {} + + for item in feature_data: + key = ( + str(item["songno"]), + str(item["difficulty"]) + ) + + feature_map[key] = item + + # ===================================================== + # feature names + # ===================================================== + + feature_names = sorted([ + k for k in feature_data[0].keys() + if k not in IGNORE_KEYS + ]) + + # ===================================================== + # dataset build + # ===================================================== + + 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 = safe_float(row[0]) + songno = str(row[1]) + diff = str(row[2]) + + key = (songno, diff) + + if key not in feature_map: + print( + f"[WARN] feature not found: " + f"{songno} {diff}" + ) + continue + + feature_item = feature_map[key] + + features = [ + safe_float(feature_item.get(k, 0)) + for k in feature_names + ] + + dataset.append(( + features, + measure + )) + + # ===================================================== + # shuffle + # ===================================================== + + random.shuffle(dataset) + + required_size = TRAIN_SIZE + VALID_SIZE + + if len(dataset) < required_size: + raise ValueError( + f"Not enough dataset " + f"({len(dataset)} < {required_size})" + ) + + # ===================================================== + # split + # ===================================================== + + train_dataset = dataset[:TRAIN_SIZE] + valid_dataset = dataset[ + TRAIN_SIZE: + TRAIN_SIZE + VALID_SIZE + ] + + X_train = np.array( + [x for x, _ in train_dataset], + dtype=np.float32 + ) + + y_train = np.array( + [y for _, y in train_dataset], + dtype=np.float32 + ) + + X_valid = np.array( + [x for x, _ in valid_dataset], + dtype=np.float32 + ) + + y_valid = np.array( + [y for _, y in valid_dataset], + dtype=np.float32 + ) + + print(f"Train Size: {len(X_train)}") + print(f"Valid Size: {len(X_valid)}") + print(f"Feature Count: {len(feature_names)}") + + # ===================================================== + # scaler + # ===================================================== + + if CONTINUE_TRAINING and os.path.exists(scaler_path): + print("Loading existing scaler...") + + scaler = joblib.load(scaler_path) + + else: + print("Creating new scaler...") + + scaler = StandardScaler() + scaler.fit(X_train) + + X_train = scaler.transform(X_train) + X_valid = scaler.transform(X_valid) + + # ===================================================== + # model + # ===================================================== + + if CONTINUE_TRAINING and os.path.exists(model_path): + print("Loading existing model...") + + model = joblib.load(model_path) + + previous_booster = model.get_booster() + + model.fit( + X_train, + y_train, + xgb_model=previous_booster + ) + + else: + print("Creating new model...") + + model = XGBRegressor( + n_estimators=N_ESTIMATORS, + max_depth=MAX_DEPTH, + learning_rate=LEARNING_RATE, + subsample=SUBSAMPLE, + colsample_bytree=COLSAMPLE_BYTREE, + objective="reg:squarederror", + random_state=RANDOM_STATE + ) + + model.fit(X_train, y_train) + + # ===================================================== + # evaluate + # ===================================================== + + pred = model.predict(X_valid) + + mae = mean_absolute_error(y_valid, pred) + + correct = np.sum( + np.abs(pred - y_valid) <= ERROR_TOLERANCE + ) + + accuracy = correct / len(y_valid) + + print(f"\nMAE: {mae:.4f}") + + print( + f"Accuracy " + f"(±{ERROR_TOLERANCE}): " + f"{accuracy:.4f}" + ) + + # ===================================================== + # feature importance + # ===================================================== + + print("\nFeature Importance:") + + importance = model.feature_importances_ + + pairs = list(zip(feature_names, importance)) + pairs.sort(key=lambda x: x[1], reverse=True) + + for name, score in pairs: + print(f"{name:25} {score:.6f}") + + # ===================================================== + # save + # ===================================================== + + joblib.dump(model, model_path) + joblib.dump(scaler, scaler_path) + + with open(feature_names_path, "w", encoding="utf-8") as f: + for name in feature_names: + f.write(name + "\n") + + print("\nSaved:") + print(model_path) + print(scaler_path) + print(feature_names_path) + + +# ========================================================= +# main +# ========================================================= + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--workingDir", + required=True + ) + + parser.add_argument( + "--dataDir", + required=True + ) + + parser.add_argument( + "--trainSize", + required=True + ) + + parser.add_argument( + "--validSize", + required=True + ) + + args = parser.parse_args() + + TRAIN_SIZE = int(args.trainSize) + VALID_SIZE = int(args.validSize) + + train_model( + args.workingDir, + args.dataDir + ) \ No newline at end of file diff --git a/tsconfig.json b/tsconfig.json deleted file mode 100644 index 1b20c17..0000000 --- a/tsconfig.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "compilerOptions": { - "lib": [] - } -} \ No newline at end of file