xgboost

predict/predict_xgboost.py

@@ -39,15 +39,17 @@ def safe_float(value):
 
 def predict(
     working_dir: str,
-    songno: str
+    songno: str,
+    feature: str = None
 ):
     # =====================================================
     # 경로
     # =====================================================
 
-    features_path = os.path.join(
+    features_path = (
-        working_dir,
+        os.path.join(working_dir, FEATURES_FILENAME)
-        FEATURES_FILENAME
+        if feature is None
+        else feature
     )
 
     model_path = os.path.join(
@@ -142,6 +144,11 @@ if __name__ == "__main__":
         required=True
     )
     
+    parser.add_argument(
+        "--feature",
+        required=False
+    )
+
     parser.add_argument(
         "--songno",
         required=True

script/compare.ts

@@ -0,0 +1,25 @@
+import Bun from 'bun';
+import { parseArgs } from 'node:util';
+
+const { values } = parseArgs({
+    args: Bun.argv,
+    options: {
+        workingDir: {
+            type: "string"
+        },
+        tja: {
+            type: "string"
+        },
+        predictScript: {
+            type: "string"
+        },
+    },
+    allowPositionals: true,
+});
+
+if (!values.tja || !values.workingDir || !values.predictScript) {
+    console.error("--workingDir --dataDir --trainDir");
+    process.exit(1);
+}
+
+const songno = "temp";

script/preprocess.ts

@@ -2,14 +2,13 @@ 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: {
+        workingDir: {
             type: "string"
         },
         dataDir: {
@@ -20,13 +19,13 @@ const { values } = parseArgs({
     allowPositionals: true,
 })
 
-if (!values.dataDir || !values.outputDir) {
+if (!values.dataDir || !values.workingDir) {
-    console.error("--outputDir --dataDir");
+    console.error("--workingDir --dataDir");
     process.exit(1);
 }
 
-const outputDir = values.outputDir ?? '';
+const workingDir = values.workingDir ?? '';
-if (!fs.existsSync(outputDir)) mkdirSync(outputDir)
+if (!fs.existsSync(workingDir)) mkdirSync(workingDir)
 const dataDir = values.dataDir ?? '';
 
 const tjaDir = path.join(dataDir, 'tja');
@@ -61,5 +60,5 @@ for (const file of files) {
     }
 }
 
-const featurePath = path.join(outputDir, 'features.json');
+const featurePath = path.join(workingDir, 'features.json');
 fs.writeFileSync(featurePath, JSON.stringify(features, null, 2), 'utf-8');

script/train.ts

@@ -0,0 +1,91 @@
+import Bun from 'bun';
+import { execSync, spawn, spawnSync } from 'node:child_process';
+import { parseArgs } from 'node:util';
+import fs from 'fs';
+import path from 'path';
+import { featurize } from '../preprocess/featurize';
+import { parseTja } from '../preprocess/parse'
+
+const { values } = parseArgs({
+    args: Bun.argv,
+    options: {
+        workingDir: {
+            type: "string"
+        },
+        dataDir: {
+            type: "string"
+        },
+        trainScript: {
+            type: "string"
+        },
+        trainSize: {
+            type: 'string'
+        },
+        validSize: {
+            type: 'string'
+        }
+    },
+    allowPositionals: true,
+});
+
+if (!values.dataDir || !values.workingDir || !values.trainScript) {
+    console.error("--workingDir --dataDir --trainDir");
+    process.exit(1);
+}
+
+generateFeatures();
+const child = spawn("python3", [values.trainScript,
+    "--workingDir", values.workingDir,
+    "--dataDir", values.dataDir,
+    "--trainSize", (Number(values.trainSize) || 500).toString(),
+    "--validSize", (Number(values.validSize) || 100).toString(),
+]);
+child.stdout.pipe(process.stdout);
+child.stderr.pipe(process.stderr);
+child.on("close", () => {
+    process.exit()
+})
+
+// funcs
+function generateFeatures() {
+    const workingDir = values.workingDir ?? '';
+    if (!fs.existsSync(workingDir)) fs.mkdirSync(workingDir);
+    const featurePath = path.join(workingDir, 'features.json');
+    if (fs.existsSync(featurePath)) return;
+    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);
+        }
+    }
+
+    fs.writeFileSync(featurePath, JSON.stringify(features, null, 2), 'utf-8');
+    console.log('features.json generated')
+}

test/features.txt

@@ -0,0 +1,8 @@
+bpm_avg
+bpm_change
+color_complexity
+density_avg
+density_peak
+note_count
+rhythm_complexity
+scroll_change

train/train_xgboost.py

@@ -3,11 +3,11 @@ import csv
 import json
 import math
 import os
+import random
 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
 
@@ -15,7 +15,9 @@ from sklearn.metrics import mean_absolute_error
 # Hyper Parameters
 # =========================================================
 
-TEST_SIZE = 0.2
+TRAIN_SIZE = 0
+VALID_SIZE = 0
+
 RANDOM_STATE = 42
 
 N_ESTIMATORS = 500
@@ -27,7 +29,7 @@ COLSAMPLE_BYTREE = 0.8
 CONTINUE_TRAINING = True
 
 # 예측 성공으로 간주할 허용 오차
-ERROR_TOLERANCE = 0.2
+ERROR_TOLERANCE = 0.1
 
 # =========================================================
 # 파일명
@@ -50,6 +52,10 @@ IGNORE_KEYS = {
 }
 
 
+# =========================================================
+# safe float
+# =========================================================
+
 def safe_float(value):
     if value is None:
         return 0.0
@@ -62,10 +68,16 @@ def safe_float(value):
     return x
 
 
+# =========================================================
+# train
+# =========================================================
+
 def train_model(
     working_dir: str,
     data_dir: str
 ):
+    random.seed(RANDOM_STATE)
+
     # =====================================================
     # path
     # =====================================================
@@ -129,14 +141,14 @@ def train_model(
     ])
 
     # =====================================================
-    # measure.csv
+    # dataset build
     # =====================================================
 
-    X = []
+    dataset = []
-    y = []
 
     with open(measure_path, "r", encoding="utf-8") as f:
         reader = csv.reader(f)
+
         next(reader, None)
 
         for row in reader:
@@ -163,29 +175,59 @@ def train_model(
                 for k in feature_names
             ]
 
-            X.append(features)
+            dataset.append((
-            y.append(measure)
+                features,
+                measure
+            ))
 
-    if len(X) == 0:
+    # =====================================================
-        raise ValueError("No training data")
+    # shuffle
+    # =====================================================
 
-    X = np.array(X, dtype=np.float32)
+    random.shuffle(dataset)
-    y = np.array(y, dtype=np.float32)
 
-    print(f"Dataset Size: {len(X)}")
+    required_size = TRAIN_SIZE + VALID_SIZE
-    print(f"Feature Count: {len(feature_names)}")
+
+    if len(dataset) < required_size:
+        raise ValueError(
+            f"Not enough dataset "
+            f"({len(dataset)} < {required_size})"
+        )
 
     # =====================================================
     # split
     # =====================================================
 
-    X_train, X_valid, y_train, y_valid = train_test_split(
+    train_dataset = dataset[:TRAIN_SIZE]
-        X,
+    valid_dataset = dataset[
-        y,
+        TRAIN_SIZE:
-        test_size=TEST_SIZE,
+        TRAIN_SIZE + VALID_SIZE
-        random_state=RANDOM_STATE
+    ]
+
+    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
     # =====================================================
@@ -251,6 +293,7 @@ def train_model(
     accuracy = correct / len(y_valid)
 
     print(f"\nMAE: {mae:.4f}")
+
     print(
         f"Accuracy "
         f"(±{ERROR_TOLERANCE}): "
@@ -305,8 +348,21 @@ if __name__ == "__main__":
         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

train/train_xgboost_full.py

@@ -3,11 +3,11 @@ import csv
 import json
 import math
 import os
-import random
 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
 
@@ -15,9 +15,7 @@ from sklearn.metrics import mean_absolute_error
 # Hyper Parameters
 # =========================================================
 
-TRAIN_SIZE = 0
+TEST_SIZE = 0.2
-VALID_SIZE = 0
-
 RANDOM_STATE = 42
 
 N_ESTIMATORS = 500
@@ -29,7 +27,7 @@ COLSAMPLE_BYTREE = 0.8
 CONTINUE_TRAINING = True
 
 # 예측 성공으로 간주할 허용 오차
-ERROR_TOLERANCE = 0.5
+ERROR_TOLERANCE = 0.1
 
 # =========================================================
 # 파일명
@@ -52,10 +50,6 @@ IGNORE_KEYS = {
 }
 
 
-# =========================================================
-# safe float
-# =========================================================
-
 def safe_float(value):
     if value is None:
         return 0.0
@@ -68,16 +62,10 @@ def safe_float(value):
     return x
 
 
-# =========================================================
-# train
-# =========================================================
-
 def train_model(
     working_dir: str,
     data_dir: str
 ):
-    random.seed(RANDOM_STATE)
-
     # =====================================================
     # path
     # =====================================================
@@ -141,14 +129,14 @@ def train_model(
     ])
 
     # =====================================================
-    # dataset build
+    # measure.csv
     # =====================================================
 
-    dataset = []
+    X = []
+    y = []
 
     with open(measure_path, "r", encoding="utf-8") as f:
         reader = csv.reader(f)
-
         next(reader, None)
 
         for row in reader:
@@ -175,59 +163,29 @@ def train_model(
                 for k in feature_names
             ]
 
-            dataset.append((
+            X.append(features)
-                features,
+            y.append(measure)
-                measure
-            ))
 
-    # =====================================================
+    if len(X) == 0:
-    # shuffle
+        raise ValueError("No training data")
-    # =====================================================
 
-    random.shuffle(dataset)
+    X = np.array(X, dtype=np.float32)
+    y = np.array(y, dtype=np.float32)
 
-    required_size = TRAIN_SIZE + VALID_SIZE
+    print(f"Dataset Size: {len(X)}")
-
+    print(f"Feature Count: {len(feature_names)}")
-    if len(dataset) < required_size:
-        raise ValueError(
-            f"Not enough dataset "
-            f"({len(dataset)} < {required_size})"
-        )
 
     # =====================================================
     # split
     # =====================================================
 
-    train_dataset = dataset[:TRAIN_SIZE]
+    X_train, X_valid, y_train, y_valid = train_test_split(
-    valid_dataset = dataset[
+        X,
-        TRAIN_SIZE:
+        y,
-        TRAIN_SIZE + VALID_SIZE
+        test_size=TEST_SIZE,
-    ]
+        random_state=RANDOM_STATE
-
-    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
     # =====================================================
@@ -293,7 +251,6 @@ def train_model(
     accuracy = correct / len(y_valid)
 
     print(f"\nMAE: {mae:.4f}")
-
     print(
         f"Accuracy "
         f"(±{ERROR_TOLERANCE}): "
@@ -348,21 +305,8 @@ if __name__ == "__main__":
         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