1.1.2 • Published 9 months ago

tensorflow-helpers v1.1.2

Weekly downloads
-
License
BSD-2-Clause
Repository
github
Last release
9 months ago

tensorflow-helpers

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

npm Package Version

Features

  • Support transfer learning and continuous learning
  • Custom image classifier using embedding features from pre-trained image model
  • Correctly save/load model on filesystem1
  • Load image file into tensor with resize and crop
  • List varies pre-trained models (url, image dimension, embedding size)
  • Support both nodejs and browser environment
  • Support caching model and image embedding
  • Typescript support
  • Works with plain Javascript, Typescript is not mandatory

1: The built-in tf.loadGraphModel() cannot load the model saved by model.save()

Installation

npm install tensorflow-helpers

You can also install tensorflow-helpers with pnpm, yarn, or slnpm

Usage Example

See model.test.ts and classifier.test.ts for complete examples.

Usage from browser:

import {
  loadImageClassifierModel,
  loadImageModel,
  toOneTensor,
} from 'tensorflow-helpers/browser'

declare var fileInput: HTMLInputElement

async function main() {
  let baseModel = await loadImageModel({
    url: 'saved_model/mobilenet-v3-large-100',
    cacheUrl: 'indexeddb://mobilenet-v3-large-100',
    checkForUpdates: false,
  })

  let classifier = await loadImageClassifierModel({
    baseModel,
    classNames: ['anime', 'real', 'others'],
    modelUrl: 'saved_model/emotion-classifier',
    cacheUrl: 'indexeddb://emotion-classifier',
  })

  fileInput.onchange = async () => {
    let file = fileInput.files?.[0]
    if (!file) return

    /* load from file directly */
    let result1 = await classifier.classifyImageFile(file)
    // result1 is Array<{ label: string, confidence: number }>

    /* load from embedding tensor */
    let embeddingTensor = await baseModel.imageFileToEmbedding(file)
    let embeddingFeatures = toOneTensor(embeddingTensor).dataSync()
    // embeddingFeatures is Float32Array

    let result1 = await classifier.classifyImageEmbedding(embeddingTensor)
  }
}
main().catch(e => console.error(e))

Usage from nodejs:

import {
  loadImageModel,
  loadImageClassifierModel,
  PreTrainedImageModels,
} from 'tensorflow-helpers'

// auto cache locally
let baseModel = await loadImageModel({
  spec: PreTrainedImageModels.mobilenet['mobilenet-v3-large-100'],
  dir: 'saved_model/base_model',
})
console.log('embedding features:', baseModel.spec.features)
// [print] embedding features: 1280

// restore or create new model
let classifier = await loadImageClassifierModel({
  baseModel,
  modelDir: 'saved_model/classifier_model',
  hiddenLayers: [128],
  datasetDir: 'dataset',
  // classNames: ['anime', 'real', 'others'], // auto scan from datasetDir
})

// auto load training dataset
let history = await classifier.trainAsync({
  epochs: 5,
  batchSize: 32,
})

// persist the parameters across restart
await classifier.save()

// auto load image from filesystem, resize and crop
let classes = await classifier.classifyAsync('image.jpg')
let topClass = topClassificationResult(classes)

console.log('result:', topClass)
// [print] result: { label: 'anime', confidence: 0.7991582155227661 }

Typescript Signature

Details see the type hints from IDE.

exported as 'tensorflow-helpers':

import * as tfjs from '@tensorflow/tfjs-node'

export let tensorflow: typeof tfjs
export let tf: typeof tfjs

exported as 'tensorflow-helpers/browser':

import * as tfjs from '@tensorflow/tfjs'

export let tensorflow: typeof tfjs
export let tf: typeof tfjs
export const PreTrainedImageModels: {
  mobilenet: {
    'mobilenet-v3-large-100': {
      url: 'https://www.kaggle.com/models/google/mobilenet-v3/TfJs/large-100-224-feature-vector/1'
      width: 224
      height: 224
      channels: 3
      features: 1280
    }
    // more models omitted ...
  }
}
export type Model = tf.GraphModel | tf.LayersModel

export function saveModel(options: {
  model: Model
  dir: string
}): Promise<SaveResult>

export function loadGraphModel(options: { dir: string }): Promise<tf.GraphModel>

export function loadLayersModel(options: {
  dir: string
}): Promise<tf.LayersModel>

export function cachedLoadGraphModel(options: {
  url: string
  dir: string
}): Promise<Model>

export function cachedLoadLayersModel(options: {
  url: string
  dir: string
}): Promise<Model>

export function loadImageModel(options: {
  spec: ImageModelSpec
  dir: string
  aspectRatio?: CropAndResizeAspectRatio
  cache?: EmbeddingCache | boolean
}): Promise<ImageModel>

export type EmbeddingCache = {
  get(filename: string): number[] | null | undefined
  set(filename: string, values: number[]): void
}

export type ImageModelSpec = {
  url: string
  width: number
  height: number
  channels: number
  features: number
}

export type ImageModel = {
  spec: ImageModelSpec
  model: Model

  fileEmbeddingCache: Map<string, tf.Tensor> | 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
}
export function loadImageFile(
  file: string,
  options?: {
    channels?: number
    dtype?: string
    expandAnimations?: boolean
    crop?: {
      width: number
      height: number
      aspectRatio?: CropAndResizeAspectRatio
    }
  },
): Promise<tf.Tensor3D | tf.Tensor4D>

export type ImageTensor = tf.Tensor3D | tf.Tensor4D

export function getImageTensorShape(imageTensor: tf.Tensor3D | tf.Tensor4D): {
  width: number
  height: number
}

export type Box = [top: number, left: number, bottom: number, right: number]

/**
 * @description calculate center-crop box
 * @returns [top,left,bottom,right], values range: 0..1
 */
export function calcCropBox(options: {
  sourceShape: { width: number; height: number }
  targetShape: { width: number; height: number }
}): Box

/**
 * @description default is 'rescale'
 *
 * 'rescale' -> scratch/transform to target shape;
 *
 * 'center-crop' -> crop the edges, maintain aspect ratio at center
 */
export type CropAndResizeAspectRatio = 'rescale' | 'center-crop'

export function cropAndResizeImageTensor(options: {
  imageTensor: tf.Tensor3D | tf.Tensor4D
  width: number
  height: number
  aspectRatio?: CropAndResizeAspectRatio
}): tf.Tensor4D

export function cropAndResizeImageFile(options: {
  srcFile: string
  destFile: string
  width: number
  height: number
  aspectRatio?: CropAndResizeAspectRatio
}): Promise<void>
export function disposeTensor(tensor: tf.Tensor | tf.Tensor[]): void

export function toOneTensor(
  tensor: tf.Tensor | tf.Tensor[] | tf.NamedTensorMap,
): tf.Tensor

export function toTensor4D(tensor: tf.Tensor3D | tf.Tensor4D): tf.Tensor4D

export function toTensor3D(tensor: tf.Tensor3D | tf.Tensor4D): tf.Tensor3D
export type ClassifierModelSpec = {
  embeddingFeatures: number
  hiddenLayers?: number[]
  classes: number
}

export function createImageClassifier(spec: ClassifierModelSpec): tf.Sequential

export type ClassificationResult = {
  label: string
  /** @description between 0 to 1 */
  confidence: number
}

export type ClassifierModel = {
  baseModel: {
    spec: ImageModelSpec
    model: Model
    loadImageAsync: (file: string) => Promise<tf.Tensor4D>
    loadImageSync: (file: string) => tf.Tensor4D
    loadAnimatedImageAsync: (file: string) => Promise<tf.Tensor4D>
    loadAnimatedImageSync: (file: string) => tf.Tensor4D
    inferEmbeddingAsync: (
      file_or_image_tensor: string | tf.Tensor,
    ) => Promise<tf.Tensor>
    inferEmbeddingSync: (file_or_image_tensor: string | tf.Tensor) => tf.Tensor
  }
  classifierModel: tf.LayersModel | tf.Sequential
  classNames: string[]
  classifyAsync: (
    file_or_image_tensor: string | tf.Tensor,
  ) => Promise<ClassificationResult[]>
  classifySync: (
    file_or_image_tensor: string | tf.Tensor,
  ) => ClassificationResult[]
  loadDatasetFromDirectoryAsync: () => Promise<{
    x: tf.Tensor<tf.Rank>
    y: tf.Tensor<tf.Rank>
  }>
  compile: () => void
  trainAsync: (options?: tf.ModelFitArgs) => Promise<tf.History>
  save: (dir?: string) => Promise<SaveResult>
}

export function loadImageClassifierModel(options: {
  baseModel: ImageModel
  hiddenLayers?: number[]
  modelDir: string
  datasetDir: string
  classNames?: string[]
}): Promise<ClassifierModel>

export function topClassifyResult(
  items: ClassificationResult[],
): ClassificationResult

/**
 * @description the values is returned as is.
 * It should has be applied softmax already
 * */
export function mapWithClassName(
  classNames: string[],
  values: ArrayLike<number>,
  options?: {
    sort?: boolean
  },
): ClassificationResult[]
/**
 * @description
 * - rename filename to content hash + extname;
 * - return list of (renamed) filenames
 */
export async function scanDir(dir: string): Promise<string[]>

export function isContentHash(file_or_filename: string): boolean

export async function saveFile(args: {
  dir: string
  content: Buffer
  mimeType: string
}): Promise<void>

export function hashContent(
  content: Buffer,
  encoding: BufferEncoding = 'hex',
): string

/** @returns new filename with content hash and extname */
export async function renameFileByContentHash(file: string): Promise<string>
/**
 * @example `loadGraphModel({ url: 'saved_model/mobilenet-v3-large-100' })`
 */
export function loadGraphModel(options: { url: string }): Promise<tf.GraphModel>

/**
 * @example `loadGraphModel({ url: 'saved_model/emotion-classifier' })`
 */
export function loadLayersModel(options: {
  url: string
}): Promise<tf.LayersModel>

/**
 * @example ```
 * cachedLoadGraphModel({
 *   url: 'saved_model/mobilenet-v3-large-100',
 *   cacheUrl: 'indexeddb://mobilenet-v3-large-100',
 * })
 * ```
 */
export function cachedLoadGraphModel(options: {
  url: string
  cacheUrl: string
  checkForUpdates?: boolean
}): Promise<tf.GraphModel<string | tf.io.IOHandler>>

/**
 * @example ```
 * cachedLoadLayersModel({
 *   url: 'saved_model/emotion-classifier',
 *   cacheUrl: 'indexeddb://emotion-classifier',
 * })
 * ```
 */
export function cachedLoadLayersModel(options: {
  url: string
  cacheUrl: string
  checkForUpdates?: boolean
}): Promise<tf.LayersModel>
export type ImageModel = {
  spec: ImageModelSpec
  model: tf.GraphModel<string | tf.io.IOHandler>
  fileEmbeddingCache: Map<string, tf.Tensor<tf.Rank>> | null
  checkCache: (url: string) => tf.Tensor | void
  loadImageCropped: (url: string) => Promise<tf.Tensor4D & tf.Tensor<tf.Rank>>
  imageUrlToEmbedding: (url: string) => Promise<tf.Tensor>
  imageFileToEmbedding: (file: File) => Promise<tf.Tensor>
  imageTensorToEmbedding: (imageTensor: ImageTensor) => tf.Tensor
}

/**
 * @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(url: string): number[] | null | undefined
  set(url: string, values: number[]): void
}

export function loadImageModel<Cache extends EmbeddingCache>(options: {
  url: string
  cacheUrl?: string
  checkForUpdates?: boolean
  aspectRatio?: CropAndResizeAspectRatio
  cache?: Cache | boolean
}): Promise<ImageModel>
export type ClassifierModel = {
  baseModel: ImageModel
  classifierModel: tf.LayersModel | tf.Sequential
  classNames: string[]
  classifyImageUrl(url: string): Promise<ClassificationResult[]>
  classifyImageFile(file: File): Promise<ClassificationResult[]>
  classifyImageTensor(
    imageTensor: tf.Tensor3D | tf.Tensor4D,
  ): Promise<ClassificationResult[]>
  classifyImage(
    image: Parameters<typeof tf.browser.fromPixels>[0],
  ): Promise<ClassificationResult[]>
  classifyImageEmbedding(embedding: tf.Tensor): Promise<ClassificationResult[]>
  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>
}

export function loadImageClassifierModel(options: {
  baseModel: ImageModel
  hiddenLayers?: number[]
  modelUrl?: string
  cacheUrl?: string
  checkForUpdates?: boolean
  classNames: string[]
}): Promise<ClassifierModel>
1.1.1

9 months ago

1.1.0

9 months ago

1.0.1

9 months ago

1.0.0

9 months ago

0.9.11

9 months ago

1.1.2

9 months ago

0.9.10

11 months ago

0.9.9

11 months ago

0.9.8

11 months ago

0.9.7

11 months ago

0.9.6

11 months ago

0.9.5

11 months ago

0.9.4

11 months ago

0.9.3

11 months ago

0.9.2

11 months ago

0.9.1

11 months ago

0.9.0

11 months ago

0.8.2

11 months ago

0.8.1

11 months ago

0.8.0

11 months ago

0.7.2

11 months ago

0.7.1

11 months ago

0.7.0

11 months ago

0.6.0

11 months ago

0.5.1

11 months ago

0.5.0

11 months ago

0.4.0

11 months ago

0.3.1

11 months ago

0.3.0

11 months ago

0.2.0

11 months ago

0.1.2

11 months ago

0.1.1

11 months ago

0.1.0

11 months ago

0.0.0

11 months ago