image-dataset
Version:
Tool to build image dataset: collect, classify, review
260 lines (259 loc) • 9.91 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.getClassNames = getClassNames;
exports.loadModels = loadModels;
exports.saveClassifierModelMetadata = saveClassifierModelMetadata;
exports.initClassNames = initClassNames;
exports.getClassLabelsInfo = getClassLabelsInfo;
exports.updateClassLabels = updateClassLabels;
const fs_1 = require("fs");
const path_1 = require("path");
const fs_2 = require("@beenotung/tslib/fs");
const tensorflow_helpers_1 = require("tensorflow-helpers");
const config_1 = require("./config");
const npm_init_helper_1 = require("npm-init-helper");
const env_1 = require("./env");
const promises_1 = require("fs/promises");
function getClassNames(options) {
let modelFile = (0, path_1.join)(config_1.config.classifierModelDir, 'model.json');
if ((0, fs_1.existsSync)(modelFile)) {
let json = JSON.parse((0, fs_1.readFileSync)(modelFile, 'utf8'));
if (json.classNames) {
return json.classNames;
}
if (json.userDefinedMetadata?.classNames) {
return json.userDefinedMetadata.classNames;
}
}
let classNames = new Set();
function checkDir(dir) {
if (!(0, fs_1.existsSync)(dir)) {
return;
}
for (let className of (0, fs_2.getDirFilenamesSync)(dir)) {
classNames.add(className);
}
}
checkDir(config_1.config.datasetRootDir);
checkDir(config_1.config.classifiedRootDir);
if (classNames.size == 0) {
if (options?.fallback) {
return options.fallback;
}
throw new Error('no class names found in dataset nor classified directory. Example: others, cat, dog, both');
}
return Array.from(classNames);
}
function createClassNameDirectories(dir, classNames) {
let existingClassNames = (0, fs_1.existsSync)(dir) ? (0, fs_2.getDirFilenamesSync)(dir) : [];
for (let className of classNames) {
if (existingClassNames.includes(className)) {
continue;
}
(0, fs_1.mkdirSync)((0, path_1.join)(dir, className), { recursive: true });
}
for (let className of existingClassNames) {
if (classNames.includes(className)) {
continue;
}
let file = (0, path_1.join)(dir, className);
let stat = (0, fs_1.statSync)(file);
if (stat.isDirectory()) {
file = JSON.stringify(file);
console.warn(`Warning: extra directory ${file} does not exist in classNames`);
}
}
}
async function loadModels() {
let classNames = getClassNames();
let imageModelSpec = tensorflow_helpers_1.PreTrainedImageModels.mobilenet['mobilenet-v3-large-100'];
let { db } = await Promise.resolve().then(() => __importStar(require('./db')));
let has_embedding = db
.prepare(
/* sql */ `select (case when embedding is null then 0 else 1 end) as count from image where filename = ?`)
.pluck();
let select_embedding = db
.prepare(
/* sql */ `select embedding from image where filename = ?`)
.pluck();
let update_embedding = db.prepare(
/* sql */ `update image set embedding = :embedding where filename = :filename`);
let insert_embedding = db.prepare(
/* sql */ `insert into image (filename, embedding) values (:filename, :embedding)`);
let select_cached_images = db
.prepare(
/* sql */ `
select filename from image where embedding is not null
`)
.pluck();
let embeddingCache = {
keys() {
return select_cached_images.all();
},
has(filename) {
return has_embedding.get(filename) == 1;
},
get(filename) {
let embedding = select_embedding.get(filename);
if (!embedding)
return null;
return embedding.split(',').map(s => +s);
},
set(filename, values) {
let embedding = values.join(',');
if (update_embedding.run({ filename, embedding }).changes == 1) {
return;
}
insert_embedding.run({ filename, embedding });
},
};
let baseModel = await (0, tensorflow_helpers_1.loadImageModel)({
dir: config_1.config.baseModelDir,
spec: imageModelSpec,
cache: embeddingCache,
});
let classifierModel = await (0, tensorflow_helpers_1.loadImageClassifierModel)({
modelDir: config_1.config.classifierModelDir,
datasetDir: config_1.config.datasetRootDir,
baseModel,
classNames,
hiddenLayers: [
(0, tensorflow_helpers_1.calcHiddenLayerSize)({
inputSize: imageModelSpec.features,
outputSize: classNames.length,
difficulty: env_1.env.CLASSIFICATION_DIFFICULTY,
}),
],
});
let metadata = loadClassifierModelMetadata(config_1.config.classifierModelDir);
createClassNameDirectories(config_1.config.datasetRootDir, classNames);
createClassNameDirectories(config_1.config.classifiedRootDir, classNames);
return {
embeddingCache,
baseModel,
classifierModel,
metadata,
};
}
function loadClassifierModelMetadata(dir) {
let file = (0, path_1.join)(dir, 'metadata.json');
let epochs = 0;
try {
let text = (0, fs_1.readFileSync)(file, 'utf8');
let json = JSON.parse(text);
epochs = json.epochs || 0;
}
catch (error) {
// missing file or invalid json
}
return {
epochs,
};
}
async function saveClassifierModelMetadata(dir, metadata) {
let file = (0, path_1.join)(dir, 'metadata.json');
let text = JSON.stringify(metadata, null, 2);
await (0, promises_1.writeFile)(file, text, 'utf8');
}
async function initClassNames() {
let classNames = getClassNames({ fallback: [] });
if (classNames.length > 1) {
console.log('loaded class names:', classNames);
return;
}
console.log();
console.log('no class names found in dataset or classified directory.');
console.log('example: others, cat, dog, both');
classNames = await askClassNames('input class names: ');
for (let className of classNames) {
(0, fs_1.mkdirSync)((0, path_1.join)(config_1.config.datasetRootDir, className), { recursive: true });
}
}
async function askClassNames(question) {
while (true) {
let input = await (0, npm_init_helper_1.ask)(question);
let classNames = input.split(',').map(name => name.trim());
if (classNames.length > 1) {
return classNames;
}
console.log('warning: at least two class names are needed');
}
}
async function getClassLabelsInfo() {
let classNames = getClassNames({ fallback: [] });
return {
classNames,
complexity: env_1.env.CLASSIFICATION_DIFFICULTY,
};
}
async function updateClassLabels(body) {
let existingClassNames = getClassNames({ fallback: [] });
let newClassNames = body.classNames;
let classesToRemove = existingClassNames.filter(className => !newClassNames.includes(className));
// check if there are same samples in removed classes
let hasFiles = false;
let rootDirs = [config_1.config.datasetRootDir, config_1.config.classifiedRootDir];
for (let rootDir of rootDirs) {
for (let className of classesToRemove) {
let dir = (0, path_1.join)(rootDir, className);
let files = (0, fs_1.existsSync)(dir) ? (0, fs_1.readdirSync)(dir).length : 0;
if (files > 0) {
console.log(`warning: ${dir} is not empty`);
hasFiles = true;
}
}
}
if (hasFiles) {
throw new Error('Cannot update labels: Some classes to be removed still contain images');
}
// remove directories for removed class names
for (let rootDir of rootDirs) {
for (let className of classesToRemove) {
let dir = (0, path_1.join)(rootDir, className);
if ((0, fs_1.existsSync)(dir)) {
(0, fs_1.rmSync)(dir, { recursive: true });
}
}
}
// remove classifier model
if ((0, fs_1.existsSync)(config_1.config.classifierModelDir)) {
(0, fs_1.rmSync)(config_1.config.classifierModelDir, { recursive: true });
}
// create directories for new class names
for (let className of newClassNames) {
(0, fs_1.mkdirSync)((0, path_1.join)(config_1.config.classifiedRootDir, className), { recursive: true });
}
}