film-poster-genres/src/subcommands/train/train.mjs

37 lines
1.1 KiB
JavaScript

"use strict";
import fs from 'fs';
import DataPreprocessor from '../../lib/DataPreprocessor.mjs';
import FilmPredictor from '../../lib/FilmPredictor.mjs';
export default async function(settings) {
if(typeof settings.input !== "string") {
console.error("Error: No input directory specified");
process.exit(1);
}
if(typeof settings.output !== "string") {
console.error("Error: No output directory specified");
process.exit(1);
}
if(!fs.existsSync(settings.input)) {
console.error(`Error: The input directory '${settings.input}' doesn't exist (did you type it correctly?)`);
process.exit(1);
}
if(!fs.existsSync(settings.output))
await fs.promises.mkdir(settings.output, { recursive: true, mode: 0o755 });
let preprocessor = new DataPreprocessor(settings.input);
let model = new FilmPredictor(settings, preprocessor.cats);
model.image_size = settings.image_size;
model.use_crossentropy = settings.cross_entropy;
await model.init(); // We're training a new model here
await model.train(
preprocessor.dataset_train(),
preprocessor.dataset_validate()
);
}