Files
fumen-measure-analyze/script/train_feature.ts
2026-04-25 16:51:07 +09:00

91 lines
2.5 KiB
TypeScript

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"
},
script: {
type: "string"
},
trainSize: {
type: 'string'
},
validSize: {
type: 'string'
}
},
allowPositionals: true,
});
if (!values.dataDir || !values.workingDir || !values.script) {
console.error("--workingDir --dataDir --trainDir");
process.exit(1);
}
generateFeatures();
const child = spawn("python3", [values.script,
"--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')
}