UNPKG

tensorflow-helpers

Version:

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

57 lines (56 loc) 2.91 kB
import { ImageModel } from './model'; import * as tf from '@tensorflow/tfjs'; import { ClassificationOptions, 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> & { getArtifacts: () => import("./model-artifacts").PatchedModelArtifacts; classNames?: string[]; }; fileEmbeddingCache: Map<string, tf.Tensor<tf.Rank>> | null; checkCache: (file_or_filename: string, options?: import("./model-utils").ImageEmbeddingOptions) => tf.Tensor | void; loadImageCropped: (file: string, options?: { expandAnimations?: boolean; } | undefined) => Promise<tf.Tensor3D | tf.Tensor4D>; imageFileToEmbedding: (file: string, options?: import("./model-utils").ImageEmbeddingOptions) => Promise<tf.Tensor>; imageTensorToEmbedding: (imageTensor: tf.Tensor3D | tf.Tensor4D, options?: import("./model-utils").ImageEmbeddingOptions) => tf.Tensor; spatialNodes: import("./spatial-utils").SpatialNode[]; spatialNodesWithUniqueShapes: import("./spatial-utils").SpatialNode[]; lastSpatialNode: import("./spatial-utils").SpatialNode | undefined; }; classifierModel: (tf.Sequential & { getArtifacts: () => import("./model-artifacts").PatchedModelArtifacts; classNames?: string[]; }) | (tf.LayersModel & { getArtifacts: () => import("./model-artifacts").PatchedModelArtifacts; classNames?: string[]; }); classNames: string[]; classifyImageFile: (file: string, options?: ClassificationOptions) => Promise<ClassificationResult[]>; classifyImageTensor: (imageTensor: tf.Tensor3D | tf.Tensor4D, options?: ClassificationOptions) => Promise<ClassificationResult[]>; classifyImageEmbedding: (embedding: tf.Tensor, options?: ClassificationOptions) => 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>; }>; export declare function getClassesFromDatasetDir(datasetDir: string): string[];