tensorflow-helpers
Version:
Helper functions to use tensorflow in nodejs for transfer learning, image classification, and more
244 lines (243 loc) • 9.23 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;
};
})();
var __exportStar = (this && this.__exportStar) || function(m, exports) {
for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.loadImageClassifierModel = loadImageClassifierModel;
exports.getClassesFromDatasetDir = getClassesFromDatasetDir;
const model_1 = require("./model");
const tf = __importStar(require("@tensorflow/tfjs"));
const fs_1 = require("fs");
const path_1 = require("path");
const timer_1 = require("@beenotung/tslib/timer");
const fs_2 = require("@beenotung/tslib/fs");
const classifier_utils_1 = require("./classifier-utils");
__exportStar(require("./classifier-utils"), exports);
async function loadImageClassifierModel(options) {
let { baseModel, datasetDir, modelDir, classNames } = options;
let classifierModel = (0, fs_1.existsSync)(modelDir)
? await (0, model_1.loadLayersModel)({ dir: modelDir, classNames })
: (0, classifier_utils_1.createImageClassifier)({
embeddingFeatures: baseModel.spec.features,
hiddenLayers: options.hiddenLayers,
get classes() {
if (!classNames && !datasetDir) {
throw new Error('classNames or datasetDir must be provided');
}
if (!classNames) {
classNames = getClassesFromDatasetDir(datasetDir);
}
return classNames.length;
},
classNames,
});
classNames = classifierModel.classNames || classNames;
if (!classNames && !datasetDir) {
throw new Error('classNames or datasetDir must be provided');
}
if (!classNames && datasetDir && (0, fs_1.existsSync)(datasetDir)) {
classNames = getClassesFromDatasetDir(datasetDir);
}
if (!classNames) {
throw new Error('classNames not provided');
}
if (classNames.length < 2) {
throw new Error('expect at least 2 classes');
}
let classCount = (0, classifier_utils_1.getClassCount)(classifierModel.outputShape);
if (classNames.length !== classCount) {
throw new Error(`classNames length mismatch: given ${classNames.length} classes, but the model has ${classCount} outputs`);
}
let compiled = false;
function compile() {
compiled = true;
classifierModel.compile({
optimizer: 'adam',
loss: tf.losses.softmaxCrossEntropy,
metrics: [tf.metrics.categoricalAccuracy],
});
}
async function classifyImageFile(file, options) {
let embedding = await baseModel.imageFileToEmbedding(file);
/* do not dispose embedding because it may be cached */
return classifyImageEmbedding(embedding, options);
}
async function classifyImageTensor(imageTensor, options) {
let embedding = baseModel.imageTensorToEmbedding(imageTensor);
let results = await classifyImageEmbedding(embedding, options);
embedding.dispose();
return results;
}
async function classifyImageEmbedding(embedding, options) {
let outputs = tf.tidy(() => {
if (embedding.rank === 1) {
embedding = tf.expandDims(embedding, 0);
}
let y = classifierModel.predict(embedding);
if (options?.squeeze) {
y = tf.squeeze(y, [0]);
}
if (options?.applySoftmax !== false) {
y = tf.softmax(y);
}
return y;
});
let values = await outputs.data();
outputs.dispose();
return (0, classifier_utils_1.mapWithClassName)(classNames, values);
}
async function loadDatasetFromDirectory() {
if (!datasetDir) {
throw new Error('datasetDir not provided');
}
let xs = [];
let classIndices = [];
let classCounts = new Array(classCount).fill(0);
let total = 0;
let classes = [];
let timer = (0, timer_1.startTimer)('scan dataset');
timer.setEstimateProgress(classNames.length);
for (let i = 0; i < classNames.length; i++) {
let className = classNames[i];
let dir = (0, path_1.join)(datasetDir, className);
let filenames = await (0, fs_2.getDirFilenames)(dir);
total += filenames.length;
classes.push({
classIdx: i,
dir,
filenames,
});
timer.tick();
}
timer.next('load dataset');
timer.setEstimateProgress(total);
for (let { classIdx, dir, filenames } of classes) {
for (let filename of filenames) {
let file = (0, path_1.join)(dir, filename);
let embedding = await baseModel.imageFileToEmbedding(file, {
squeeze: true,
});
xs.push(embedding);
classIndices.push(classIdx);
classCounts[classIdx]++;
timer.tick();
}
}
timer.next('stack embeddings');
let x = tf.stack(xs);
if (!baseModel.fileEmbeddingCache) {
timer.next('dispose individual embeddings');
timer.setEstimateProgress(xs.length);
for (let x of xs) {
x.dispose();
timer.tick();
}
}
timer.next('prepare one-hot labels');
let y = tf.tidy(() => tf.oneHot(tf.tensor1d(classIndices, 'int32'), classCount));
timer.end();
return { x, y, classCounts };
}
async function train(options) {
if (!compiled) {
compile();
}
let next = async (x, y, classWeight) => {
let history = await classifierModel.fit(x, y, {
...options,
shuffle: true,
classWeight,
});
return history;
};
if (options && 'x' in options && 'y' in options) {
let { x, y, classCounts, ...rest } = options;
options = rest;
let classWeight = options.classWeight ||
(classCounts
? (0, classifier_utils_1.calcClassWeight)({
classes: classCount,
classCounts,
})
: undefined);
return next(x, y, classWeight);
}
else {
let { x, y, classCounts } = await loadDatasetFromDirectory();
let classWeight = options?.classWeight ||
(0, classifier_utils_1.calcClassWeight)({
classes: classCount,
classCounts,
});
let history = await next(x, y, classWeight);
x.dispose();
y.dispose();
return history;
}
}
async function save(dir = modelDir) {
return await (0, model_1.saveModel)({
model: classifierModel,
dir,
classNames,
});
}
return {
baseModel,
classifierModel,
classNames,
classifyImageFile,
classifyImageTensor,
classifyImageEmbedding,
loadDatasetFromDirectory,
compile,
train,
save,
};
}
function getClassesFromDatasetDir(datasetDir) {
if (!(0, fs_1.existsSync)(datasetDir)) {
throw new Error('datasetDir not exists');
}
let classNames = (0, fs_2.getDirFilenamesSync)(datasetDir).filter(filename => {
let file = (0, path_1.join)(datasetDir, filename);
return (0, fs_1.statSync)(file).isDirectory();
});
if (classNames.length === 0) {
throw new Error('no image folders found in datasetDir');
}
return classNames;
}