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");
const tensor_1 = require("../tensor");
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');
    }
    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 classifyImageUrl(url) {
        let embedding = await baseModel.imageUrlToEmbedding(url);
        /* do not dispose embedding because it may be cached */
        return classifyImageEmbedding(embedding);
    }
    async function classifyImageFile(file) {
        let embedding = await baseModel.imageFileToEmbedding(file);
        /* do not dispose embedding because it may be cached */
        return classifyImageEmbedding(embedding);
    }
    async function classifyImage(image) {
        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);
    }
    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 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,
    };
}