imgrecog/lib/tensorflow.js

Node.js tool to parse and act on images, using the Google Vision and Sightengine APIs.

"use strict";
// TENSORFLOW
Object.defineProperty(exports, "__esModule", { value: true });
exports.TensowFlow = void 0;
const utils_1 = require("./utils");
const fs = require("fs");
const mobilenet = require("@tensorflow-models/mobilenet");
const tfnode = require("@tensorflow/tfjs-node");
/**
 * TensorFlow wrapper.
 */
class TensowFlow {
    constructor() {
        /**
         * Prepare the TensorFlow MobileNet model.
         * @param options Program options.
         */
        this.prepare = async (options) => {
            if (!this.tfModel) {
                require("@tensorflow/tfjs");
                this.tfModel = await mobilenet.load();
                utils_1.logDebug(options, `Loaded TensorFlow / MobileNet model`);
            }
        };
        /**
         * Parse the image against the
         * @param options Program options.
         * @param filepath Image file to be scanned.
         */
        this.parse = async (options, filepath) => {
            try {
                const buffer = fs.readFileSync(filepath);
                const tfImage = tfnode.node.decodeImage(buffer);
                const results = await this.tfModel.classify(tfImage);
                const logtext = [];
                const tags = {};
                // Iterate results and set individual tags.
                for (let prediction of results) {
                    const keywords = prediction.className.split(",");
                    for (let kw of keywords) {
                        const key = utils_1.normalizeTag(kw);
                        const score = utils_1.normalizeScore(prediction.probability);
                        if (score) {
                            logtext.push(`${key}:${score}`);
                            tags[key] = score;
                        }
                    }
                }
                const details = logtext.length > 0 ? logtext.join(", ") : "NONE";
                const logDetails = `${filepath}: tags - ${details}`;
                utils_1.logInfo(options, logDetails);
                return {
                    file: filepath,
                    tags: tags
                };
            }
            catch (ex) {
                utils_1.logError(options, `${filepath} - error parsing`, ex);
            }
        };
    }
}
exports.TensowFlow = TensowFlow;
// Exports...
exports.default = new TensowFlow();