image-dataset/dist/classify.js

Tool to build image dataset: collect, classify, review

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.stopClassify = stopClassify;
exports.main = main;
const fs_1 = require("fs");
const promises_1 = require("fs/promises");
const fs_2 = require("@beenotung/tslib/fs");
const path_1 = require("path");
const tensorflow_helpers_1 = require("tensorflow-helpers");
const timer_1 = require("@beenotung/tslib/timer");
const cache_1 = require("./cache");
let unclassifiedDir = 'unclassified';
let classifiedDir = 'classified';
(0, fs_1.mkdirSync)(unclassifiedDir, { recursive: true });
let running = false;
function stopClassify() {
    running = false;
}
async function main() {
    let timer = (0, timer_1.startTimer)('load models');
    let { classifierModel } = await cache_1.modelsCache.get();
    timer.next('init directories');
    for (let className of classifierModel.classNames) {
        (0, fs_1.mkdirSync)((0, path_1.join)(classifiedDir, className), { recursive: true });
    }
    timer.next('load file list');
    let filenames = await (0, fs_2.getDirFilenames)(unclassifiedDir);
    timer.next('classify images');
    timer.setEstimateProgress(filenames.length);
    running = true;
    for (let filename of filenames) {
        if (!running) {
            break;
        }
        let src = (0, path_1.join)(unclassifiedDir, filename);
        let result = await classifierModel.classifyImageFile(src);
        let className = (0, tensorflow_helpers_1.topClassifyResult)(result).label;
        let dest = (0, path_1.join)(classifiedDir, className, filename);
        await (0, promises_1.rename)(src, dest);
        timer.tick();
    }
    timer.end();
}
if (process.argv[1] == __filename) {
    main();
}