tensorflow-helpers
Version:
Helper functions to use tensorflow in nodejs for transfer learning, image classification, and more
194 lines (193 loc) • 7.32 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;
const model_1 = require("./model");
const tf = __importStar(require("@tensorflow/tfjs"));
const fs_1 = require("fs");
const path_1 = require("path");
const tensor_1 = require("./tensor");
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;
if (!classNames && (0, fs_1.existsSync)(datasetDir)) {
classNames = (0, fs_2.getDirFilenamesSync)(datasetDir);
}
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) {
throw new Error('classNames not provided');
}
return classNames.length;
},
classNames,
});
classNames = classifierModel.classNames;
if (!classNames) {
throw new Error('classNames not provided');
}
let classCount = classNames.length;
if (classCount < 2) {
throw new Error('expect at least 2 classes');
}
let compiled = false;
function compile() {
compiled = true;
classifierModel.compile({
optimizer: 'adam',
loss: tf.metrics.categoricalCrossentropy,
metrics: [tf.metrics.categoricalAccuracy],
});
}
async function classifyImageFile(file) {
let embedding = await baseModel.imageFileToEmbedding(file);
/* do not dispose embedding because it may be cached */
return classifyImageEmbedding(embedding);
}
async function classifyImageTensor(imageTensor) {
let embedding = baseModel.imageTensorToEmbedding(imageTensor);
let results = await classifyImageEmbedding(embedding);
embedding.dispose();
return results;
}
async function classifyImageEmbedding(embedding) {
let outputs = tf.tidy(() => {
let outputs = classifierModel.predict(embedding);
return (0, tensor_1.toOneTensor)(outputs);
});
let values = await outputs.data();
(0, tensor_1.disposeTensor)(outputs);
return (0, classifier_utils_1.mapWithClassName)(classNames, values);
}
async function loadDatasetFromDirectory() {
let xs = [];
let classIndices = [];
let classCounts = new Array(classCount).fill(0);
let total = 0;
let classes = await Promise.all(classNames.map(async (className, classIdx) => {
let dir = (0, path_1.join)(datasetDir, className);
let filenames = await (0, fs_2.getDirFilenames)(dir);
total += filenames.length;
return { classIdx, dir, filenames };
}));
let timer = (0, timer_1.startTimer)('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);
xs.push(embedding);
classIndices.push(classIdx);
classCounts[classIdx]++;
timer.tick();
}
}
timer.end();
let x = tf.concat(xs);
if (!baseModel.fileEmbeddingCache) {
for (let x of xs) {
x.dispose();
}
}
let y = tf.tidy(() => tf.oneHot(tf.tensor1d(classIndices, 'int32'), classCount));
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,
};
}