image-dataset
Version:
Tool to build image dataset: collect, classify, review
306 lines (305 loc) • 10.8 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.startServer = startServer;
const express_1 = __importDefault(require("express"));
const listening_on_1 = require("listening-on");
const train_1 = require("./train");
const retrain_1 = require("./retrain");
const classify_1 = require("./classify");
const unclassify_1 = require("./unclassify");
const rename_by_content_hash_1 = require("./rename-by-content-hash");
const path_1 = require("path");
const promises_1 = require("fs/promises");
const fs_1 = require("@beenotung/tslib/fs");
const cache_1 = require("./cache");
const tensorflow_helpers_1 = require("tensorflow-helpers");
const functional_1 = require("@beenotung/tslib/functional");
const env_1 = require("./env");
const file_1 = require("./file");
const fs_2 = require("fs");
const timer_1 = require("@beenotung/tslib/timer");
const time_1 = require("@beenotung/tslib/time");
const wait_1 = require("@beenotung/tslib/async/wait");
const compare_1 = require("@beenotung/tslib/compare");
const config_1 = require("./config");
const model_1 = require("./model");
let app = (0, express_1.default)();
(0, fs_2.mkdirSync)('downloaded', { recursive: true });
(0, fs_2.mkdirSync)('dataset', { recursive: true });
(0, fs_2.mkdirSync)('classified', { recursive: true });
(0, fs_2.mkdirSync)('unclassified', { recursive: true });
app.use('/data-template', express_1.default.static((0, file_1.resolveFile)('node_modules/data-template')));
app.use('/heatmap-helpers/bundle.js', express_1.default.static((0, file_1.resolveFile)('node_modules/heatmap-helpers/bundle.js')));
app.use('/ionicons', express_1.default.static((0, file_1.resolveFile)('node_modules/ionicons')));
app.use('/classified', express_1.default.static('classified'));
app.use('/unclassified', express_1.default.static('unclassified'));
app.use('/saved_models', express_1.default.static('saved_models'));
app.use(express_1.default.static((0, file_1.resolveFile)('public')));
app.use(express_1.default.json());
app.use(express_1.default.urlencoded({ extended: false }));
/***********/
/* actions */
/***********/
let actions = {
renameByContentHash: rename_by_content_hash_1.main,
async updateModelSettings(body) {
(0, classify_1.stopClassify)();
env_1.env.CLASSIFICATION_DIFFICULTY = body.complexity;
await (0, model_1.updateClassLabels)(body);
unclassifiedImageCache.clear();
return cache_1.modelsCache.load();
},
loadModel() {
unclassifiedImageCache.clear();
return cache_1.modelsCache.load();
},
loadDataset() {
(0, classify_1.stopClassify)();
return cache_1.datasetCache.load();
},
train() {
(0, classify_1.stopClassify)();
unclassifiedImageCache.clear();
return (0, train_1.main)();
},
retrain: retrain_1.main,
classify: classify_1.main,
unclassifyAll: unclassify_1.main,
};
app.get('/actions', (req, res) => {
res.json({ actions: Object.keys(actions) });
});
for (let [name, fn] of Object.entries(actions)) {
app.post('/' + name, async (req, res) => {
try {
await fn(req.body);
res.json({});
}
catch (error) {
res.json({ error: String(error) });
}
});
}
app.post('/unclassify', async (req, res) => {
try {
let dir = req.query.dir;
if (!dir || typeof dir != 'string') {
throw new Error('dir must be given in req.query');
}
(0, unclassify_1.unclassifyDir)(dir);
res.json({
total: await countDirFiles('unclassified'),
});
}
catch (error) {
res.json({ error: String(error) });
}
});
app.post('/unclassified/restore', (req, res) => {
try {
let json = (0, unclassify_1.restoreUnclassified)();
res.json(json);
}
catch (error) {
res.json({ error: String(error) });
}
});
/******************/
/* model settings */
/******************/
app.get('/getClassLabelsInfo', async (req, res) => {
try {
(0, classify_1.stopClassify)();
let info = await (0, model_1.getClassLabelsInfo)();
unclassifiedImageCache.clear();
res.json(info);
}
catch (error) {
res.json({ error: String(error) });
}
});
app.post('/updateModelSettings', async (req, res) => {
try {
let { complexity, classNames } = req.body;
if (!(complexity >= 1 && complexity <= 5)) {
throw new Error('Complexity must be between 1 and 5');
}
if (!Array.isArray(classNames) || classNames.length < 2) {
throw new Error('At least two class names are required');
}
await actions.updateModelSettings({ complexity, classNames });
res.json({});
}
catch (error) {
res.json({ error: String(error) });
}
});
/*********/
/* stats */
/*********/
async function countDirFiles(dir) {
let filenames = await (0, fs_1.getDirFilenames)(dir);
return filenames.length;
}
async function countDir2Files(dir) {
let dirnames = await (0, fs_1.getDirFilenames)(dir);
return Promise.all(dirnames.map(async (dirname) => ({
dirname,
count: await countDirFiles((0, path_1.join)(dir, dirname)),
})));
}
app.get('/stats', async (req, res) => {
res.json({
stats: {
classNames: (0, model_1.getClassNames)({ fallback: [] }),
dataset: await countDir2Files(config_1.config.datasetRootDir),
classified: await countDir2Files(config_1.config.classifiedRootDir),
downloaded: await countDir2Files(config_1.config.downloadedRootDir),
unclassified: await countDirFiles(config_1.config.unclassifiedRootDir),
},
});
});
/************/
/* classify */
/************/
async function scanDir2Files(dir) {
let classNames = await (0, fs_1.getDirFilenames)(dir);
return Promise.all(classNames.map(async (className) => {
return {
className,
filenames: await (0, fs_1.getDirFilenames)((0, path_1.join)(dir, className)),
};
}));
}
let isClassifying = false;
let unclassifiedImageCache = new Map();
app.get('/unclassified', async (req, res) => {
let hasSentResponse = false;
try {
let { baseModel, classifierModel, metadata } = await cache_1.modelsCache.get();
let epochs = metadata.epochs;
let images = [];
let filenames = await (0, fs_1.getDirFilenames)('unclassified');
filenames = filenames.filter(name => !name.endsWith('.txt') && !name.endsWith('.log'));
let newFilenames = [];
let deadlineTimeout = 3 * time_1.SECOND;
let deadline = Date.now() + deadlineTimeout;
let timer = (0, timer_1.startTimer)('load unclassified (cached)');
timer.setEstimateProgress(filenames.length);
for (let filename of filenames) {
if (metadata.epochs != epochs)
break;
let image = unclassifiedImageCache.get(filename);
if (image) {
images.push(image);
}
else {
newFilenames.push(filename);
}
timer.tick();
}
if (isClassifying) {
sendResponse();
return;
}
isClassifying = true;
timer.next('load unclassified (uncached)');
timer.setEstimateProgress(newFilenames.length);
for (let filename of newFilenames) {
if (metadata.epochs != epochs)
return;
let file = (0, path_1.join)('unclassified', filename);
let results = await classifierModel.classifyImageFile(file);
let result = (0, tensorflow_helpers_1.topClassifyResult)(results);
let image = {
filename,
label: result.label,
confidence: result.confidence,
results,
};
unclassifiedImageCache.set(filename, image);
images.push(image);
if (!hasSentResponse && Date.now() >= deadline) {
sendResponse();
}
timer.tick();
await (0, wait_1.later)(0);
}
timer.end();
sendResponse();
function sendResponse() {
if (hasSentResponse) {
return;
}
hasSentResponse = true;
res.json({
classes: Array.from((0, functional_1.groupBy)(image => image.label, images).entries(), ([className, images]) => ({
className,
images,
})).sort((a, b) => (0, compare_1.compare)(a.className, b.className)),
});
}
}
catch (error) {
console.error(error);
if (!hasSentResponse) {
res.json({ error: String(error) });
}
}
finally {
isClassifying = false;
}
});
app.get('/classified', async (req, res) => {
res.json({ classes: await scanDir2Files('classified') });
});
app.post('/correct', async (req, res) => {
try {
let { className, images } = req.body;
for (let srcFile of images) {
srcFile = (0, path_1.join)('.', srcFile);
if (!!srcFile.startsWith('classified/') &&
!!srcFile.startsWith('unclassified/')) {
continue;
}
let filename = (0, path_1.basename)(srcFile);
unclassifiedImageCache.delete(filename);
let destFile = (0, path_1.join)('dataset', className, filename);
await (0, promises_1.rename)(srcFile, destFile);
}
res.json({});
}
catch (error) {
res.json({ error: String(error) });
}
});
app.post('/undo', async (req, res) => {
try {
let { className, images } = req.body;
for (let destFile of images) {
destFile = (0, path_1.join)('.', destFile);
if (!!destFile.startsWith('classified/') &&
!!destFile.startsWith('unclassified/')) {
continue;
}
let filename = (0, path_1.basename)(destFile);
let srcFile = (0, path_1.join)('dataset', className, filename);
await (0, promises_1.rename)(srcFile, destFile);
}
res.json({});
}
catch (error) {
res.json({ error: String(error) });
}
});
function startServer(port) {
app.listen(port, () => {
(0, listening_on_1.print)(port);
});
}
if ((0, path_1.basename)(process.argv[1]) == (0, path_1.basename)(__filename)) {
startServer(env_1.env.PORT);
}