tensorflow-helpers/browser/classifier.js

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

"use strict";
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.loadImageClassifierModel = loadImageClassifierModel;
const tf = __importStar(require("@tensorflow/tfjs"));
const model_1 = require("./model");
const classifier_utils_1 = require("../classifier-utils");
async function loadImageClassifierModel(options) {
    let { baseModel, classNames } = options;
    async function loadClassifierModel() {
        let { modelUrl: url, cacheUrl, checkForUpdates } = options;
        if (url && cacheUrl) {
            try {
                let model = await (0, model_1.cachedLoadLayersModel)({
                    url,
                    cacheUrl,
                    checkForUpdates,
                    classNames,
                });
                return model;
            }
            catch (error) {
                if (!String(error).includes('file not found')) {
                    throw error;
                }
            }
        }
        return (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,
        });
    }
    let classifierModel = await loadClassifierModel();
    classNames = classifierModel.classNames;
    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 classifyImageUrl(url, options) {
        let embedding = await baseModel.imageUrlToEmbedding(url);
        /* do not dispose embedding because it may be cached */
        return classifyImageEmbedding(embedding, options);
    }
    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 classifyImage(image, options) {
        let imageTensor = await tf.browser.fromPixelsAsync(image);
        let embedding = baseModel.imageTensorToEmbedding(imageTensor);
        imageTensor.dispose();
        /* 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 train(options) {
        if (!compiled) {
            compile();
        }
        let { x, y, classCounts, ...rest } = options;
        let classWeight = options.classWeight ||
            (classCounts
                ? (0, classifier_utils_1.calcClassWeight)({
                    classes: classCount,
                    classCounts,
                })
                : undefined);
        let history = await classifierModel.fit(x, y, {
            ...options,
            shuffle: true,
            classWeight,
        });
        return history;
    }
    return {
        baseModel,
        classifierModel,
        classNames,
        classifyImageUrl,
        classifyImageFile,
        classifyImageTensor,
        classifyImage,
        classifyImageEmbedding,
        compile,
        train,
    };
}