tensorflow-helpers/model.d.ts

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

import * as tf from '@tensorflow/tfjs-node';
import { CropAndResizeAspectRatio } from './image';
import { ImageModelSpec } from './image-model';
import { SaveResult } from './classifier-utils';
export { ImageModelSpec, PreTrainedImageModels } from './image-model';
export type Model = tf.GraphModel | tf.LayersModel;
export declare function saveModel(options: {
    model: Model;
    dir: string;
    classNames?: string[];
}): Promise<SaveResult>;
export declare function loadGraphModel(options: {
    dir: string;
    classNames?: string[];
}): Promise<tf.GraphModel<string | import("@tensorflow/tfjs-core/dist/io/types").IOHandler> & {
    classNames?: string[];
}>;
export declare function loadLayersModel(options: {
    dir: string;
    classNames?: string[];
}): Promise<tf.LayersModel & {
    classNames?: string[];
}>;
export declare function cachedLoadGraphModel(options: {
    url: string;
    dir: string;
    classNames?: string[];
}): Promise<tf.GraphModel<string | import("@tensorflow/tfjs-core/dist/io/types").IOHandler> & {
    classNames?: string[];
}>;
export declare function cachedLoadLayersModel(options: {
    url: string;
    dir: string;
    classNames?: string[];
}): Promise<tf.LayersModel & {
    classNames?: string[];
}>;
export type ImageModel = Awaited<ReturnType<typeof loadImageModel>>;
/**
 * @description cache image embedding keyed by filename.
 * The dirname is ignored.
 * The filename is expected to be content hash (w/wo extname)
 */
export type EmbeddingCache = {
    get(filename: string): number[] | null | undefined;
    set(filename: string, values: number[]): void;
};
export declare function loadImageModel<Cache extends EmbeddingCache>(options: {
    spec: ImageModelSpec;
    dir: string;
    aspectRatio?: CropAndResizeAspectRatio;
    cache?: Cache | boolean;
}): Promise<{
    spec: ImageModelSpec;
    model: tf.GraphModel<string | import("@tensorflow/tfjs-core/dist/io/types").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;
    }) => Promise<tf.Tensor3D | tf.Tensor4D>;
    imageFileToEmbedding: (file: string, options?: {
        expandAnimations?: boolean;
    }) => Promise<tf.Tensor>;
    imageTensorToEmbedding: (imageTensor: tf.Tensor3D | tf.Tensor4D) => tf.Tensor;
}>;