feat: predict_tja script

2026-04-25 16:19:30 +09:00
parent df82515a33
commit 4f6ce49704
5 changed files with 201 additions and 38 deletions
--- a/script/predict_tja.ts
+++ b/script/predict_tja.ts
@@ -0,0 +1,94 @@
+import Bun from 'bun';
+import path from 'node:path';
+import { parseArgs } from 'node:util';
+import fs from 'node:fs';
+import { featurize } from '../preprocess/featurize';
+import { parseTja } from '../preprocess/parse';
+
+const { values } = parseArgs({
+    args: Bun.argv,
+    options: {
+        tja: {
+            type: "string"
+        },
+        workingDir: {
+            type: "string"
+        },
+        script: {
+            type: "string"
+        }
+    },
+    allowPositionals: true,
+});
+
+if (!values.tja || !values.workingDir || !values.script) {
+    console.error("Usage: bun script/predict_tja.ts --tja <path_to_tja> [--workingDir <dir>] [--script <python_script>]");
+    process.exit(1);
+}
+
+const tjaPath = values.tja;
+if (!fs.existsSync(tjaPath)) {
+    console.error(`File not found: ${tjaPath}`);
+    process.exit(1);
+}
+
+const workingDir = values.workingDir!;
+const pythonScript = values.script!;
+
+if (!fs.existsSync(pythonScript)) {
+    console.error(`Python script not found: ${pythonScript}`);
+    process.exit(1);
+}
+
+// 1. TJA 파싱 및 피처 추출
+const tjaContent = fs.readFileSync(tjaPath, 'utf-8');
+const songno = path.basename(tjaPath, '.tja');
+const parsed = parseTja(tjaContent);
+
+const features: any[] = [];
+if (parsed.oni) {
+    features.push({
+        songno,
+        difficulty: 'oni',
+        ...featurize(parsed.oni)
+    });
+}
+if (parsed.edit) {
+    features.push({
+        songno,
+        difficulty: 'ura',
+        ...featurize(parsed.edit)
+    });
+}
+
+if (features.length === 0) {
+    console.error("No Oni or Ura difficulty found in the TJA file.");
+    process.exit(1);
+}
+
+// 2. 임시 피처 파일 저장
+const tempFeaturePath = path.join(workingDir, `temp_predict_${Date.now()}.json`);
+fs.writeFileSync(tempFeaturePath, JSON.stringify(features, null, 2), 'utf-8');
+
+try {
+    // 3. Python 예측 스크립트 실행
+    const proc = Bun.spawnSync([
+        "python3",
+        pythonScript,
+        "--workingDir", workingDir,
+        "--feature", tempFeaturePath,
+        "--songno", songno
+    ]);
+
+    if (proc.success) {
+        console.log(proc.stdout.toString());
+    } else {
+        console.error("Prediction failed:");
+        console.error(proc.stderr.toString());
+    }
+} finally {
+    // 4. 임시 파일 삭제
+    if (fs.existsSync(tempFeaturePath)) {
+        fs.unlinkSync(tempFeaturePath);
+    }
+}