tensorflow-helpers/image-features.js

Helper functions to use tensorflow in nodejs for transfer learning, image classification, and more

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.getImageFeatures = getImageFeatures;
const spatial_utils_1 = require("./spatial-utils");
function getName(node) {
    return typeof node == 'string' ? node : node.name;
}
async function getImageFeatures(options) {
    let { tf, imageModel, image } = options;
    let model = imageModel.model;
    let spatialNodes = options.spatialNodes || options.spatialNode || (0, spatial_utils_1.getLastSpatialNodeName)(model);
    let outputNode = options.outputNode || 'Identity:0';
    let names = Array.isArray(spatialNodes)
        ? [...spatialNodes.map(getName), outputNode]
        : [getName(spatialNodes), outputNode];
    let input = typeof image == 'string' ? await imageModel.loadImageCropped(image) : image;
    let output = tf.tidy(() => model.execute(input, names));
    if (typeof image == 'string') {
        input.dispose();
    }
    let spatialFeatures = Array.isArray(spatialNodes)
        ? output.slice(0, output.length - 1)
        : output[0];
    let pooledFeatures = output[output.length - 1];
    return {
        spatialFeatures,
        pooledFeatures,
    };
}