tensorflow-helpers/classifier.d.ts

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

import { ImageModel } from './model';
import * as tf from '@tensorflow/tfjs';
import { ClassificationResult } from './classifier-utils';
export * from './classifier-utils';
export type ClassifierModel = Awaited<ReturnType<typeof loadImageClassifierModel>>;
export declare function loadImageClassifierModel(options: {
    baseModel: ImageModel;
    hiddenLayers?: number[];
    modelDir: string;
    datasetDir: string;
    /** @description if not provided, will be auto scanned from datasetDir or load from the model.json */
    classNames?: string[];
}): Promise<{
    baseModel: {
        spec: import("./image-model").ImageModelSpec;
        model: tf.GraphModel<string | tf.io.IOHandler> & {
            classNames?: string[];
        };
        fileEmbeddingCache: Map<string, tf.Tensor<tf.Rank>> | null;
        checkCache: (file_or_filename: string) => tf.Tensor | void;
        loadImageCropped: (file: string, options?: {
            expandAnimations?: boolean;
        } | undefined) => Promise<tf.Tensor3D | tf.Tensor4D>;
        imageFileToEmbedding: (file: string, options?: {
            expandAnimations?: boolean;
        } | undefined) => Promise<tf.Tensor>;
        imageTensorToEmbedding: (imageTensor: tf.Tensor3D | tf.Tensor4D) => tf.Tensor;
    };
    classifierModel: (tf.Sequential & {
        classNames?: string[];
    }) | (tf.LayersModel & {
        classNames?: string[];
    });
    classNames: string[];
    classifyImageFile: (file: string) => Promise<ClassificationResult[]>;
    classifyImageTensor: (imageTensor: tf.Tensor3D | tf.Tensor4D) => Promise<ClassificationResult[]>;
    classifyImageEmbedding: (embedding: tf.Tensor) => Promise<ClassificationResult[]>;
    loadDatasetFromDirectory: () => Promise<{
        x: tf.Tensor<tf.Rank>;
        y: tf.Tensor<tf.Rank>;
        classCounts: number[];
    }>;
    compile: () => void;
    train: (options?: tf.ModelFitArgs & ({
        x: tf.Tensor<tf.Rank>;
        y: tf.Tensor<tf.Rank>;
        /** @description to calculate classWeight */
        classCounts?: number[];
    } | {})) => Promise<tf.History>;
    save: (dir?: string) => Promise<tf.io.SaveResult>;
}>;