image-dataset/dist/train.js

Tool to build image dataset: collect, classify, review

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.main = main;
const cache_1 = require("./cache");
const config_1 = require("./config");
const model_1 = require("./model");
async function main() {
    let { classifierModel, metadata } = await cache_1.modelsCache.get();
    let { x, y, classCounts } = await cache_1.datasetCache.get();
    let epochs = 10;
    let batchSize = 32;
    let total_epochs = metadata.epochs + epochs;
    let total_batches = Math.ceil(x.shape[0] / batchSize);
    console.log(`\nTraining ${epochs} epochs...`);
    await classifierModel.train({
        x,
        y,
        classCounts,
        epochs,
        batchSize,
        verbose: 0,
        callbacks: [
            {
                onEpochBegin(epoch, logs) {
                    metadata.epochs++;
                    process.stdout.write(`Epoch: ${metadata.epochs}/${total_epochs}, Batch: 1/${total_batches}`);
                },
                onBatchEnd: (batch, logs) => {
                    let accuracy = formatNumber(logs.categoricalAccuracy);
                    let loss = formatNumber(logs.loss);
                    process.stdout.write(`\rEpoch: ${metadata.epochs}/${total_epochs}, Batch: ${batch + 1}/${total_batches}, Accuracy: ${accuracy}, Loss: ${loss}`);
                },
                onEpochEnd: (epoch, logs) => {
                    process.stdout.write(`\n`);
                },
            },
        ],
    });
    await classifierModel.save();
    await (0, model_1.saveClassifierModelMetadata)(config_1.config.classifierModelDir, metadata);
}
function formatNumber(x) {
    if (x >= 1) {
        return x.toFixed(2);
    }
    if (x >= 0.1) {
        return x.toFixed(3);
    }
    if (x >= 0.01) {
        return x.toFixed(4);
    }
    return x.toExponential(2);
}
if (process.argv[1] == __filename) {
    main();
}