license import { DataType , InferenceModel , ModelPredictConfig , ModelTensorInfo , NamedTensorMap , tensor , Tensor } from '@tensorflow/tfjs-core' ; import * as tfliteWebAPIClient from './tflite_web_api_client' ; import { ProfileItem , TFLiteDataType , TFLiteWebModelRunner , TFLiteWebModelRunnerOptions , TFLiteWebModelRunnerTensorInfo } from './types/tflite_web_model_runner' ; const TFHUB_SEARCH_PARAM = '?lite-format=tflite' ; doc export class TFLiteModel implements InferenceModel { constructor ( private readonly modelRunner ) { } get inputs ( ) : ModelTensorInfo [ ] { const modelInputs = this . modelRunner . getInputs ( ) ; return this . convertTFLiteTensorInfos ( modelInputs ) ; } get outputs ( ) : ModelTensorInfo [ ] { const modelOutputs = this . modelRunner . getOutputs ( ) ; return this . convertTFLiteTensorInfos ( modelOutputs ) ; } param inputs param config returns doc predict ( inputs , config ? ) { const modelInputs = this . modelRunner . getInputs ( ) ; const modelOutputs = this . modelRunner . getOutputs ( ) ; if ( inputs instanceof Tensor || Array . isArray ( inputs ) ) { let inputTensors ; if ( inputs instanceof Tensor ) { inputTensors = [ inputs ] ; } else { inputTensors = inputs ; } if ( modelInputs . length !== inputTensors . length ) { throw new Error ( `The size of TFLite model inputs (${ modelInputs . length }) does not match the size of the input tensors (${ inputTensors . length })` ) ; } for ( let i = 0 ; i < modelInputs . length ; i ++ ) { this . setModelInputFromTensor ( modelInputs [ i ] , inputTensors [ i ] ) ; } } else { const inputTensorNames = Object . keys ( inputs ) ; const modelInputMap = { } ; modelInputs . forEach ( modelInput => { modelInputMap [ modelInput . name ] = modelInput ; } ) ; const modelInputNames = Object . keys ( modelInputMap ) ; this . checkMapInputs ( inputTensorNames , modelInputNames ) ; for ( const name of inputTensorNames ) { this . setModelInputFromTensor ( modelInputMap [ name ] , inputs [ name ] ) ; } } const success = this . modelRunner . infer ( ) ; if ( ! success ) { throw new Error ( 'Failed running inference' ) ; } const outputTensors = { } ; for ( let i = 0 ; i < modelOutputs . length ; i ++ ) { const modelOutput = modelOutputs [ i ] ; let data = modelOutput . data ( ) ; switch ( modelOutput . dataType ) { case 'int8' : case 'int16' : case 'uint32' : data = Int32Array . from ( data ) ; break ; case 'float64' : console . warn ( `WARNING: converting output tensor from 'float64' to 'float32'` ) ; data = Float32Array . from ( data ) ; break ; default : break ; } const outputTensor = tensor ( data , this . getShapeFromTFLiteTensorInfo ( modelOutput ) ) ; outputTensors [ modelOutput . name ] = outputTensor ; } const names = Object . keys ( outputTensors ) ; return names . length === 1 ? outputTensors [ names [ 0 ] ] : outputTensors ; } param inputs param outputs returns execute ( inputs , outputs ) { throw new Error ( 'execute() of TFLiteModel is not supported yet.' ) ; } getProfilingResults ( ) { return this . modelRunner . getProfilingResults ( ) ; } getProfilingSummary ( ) { return this . modelRunner . getProfilingSummary ( ) ; } private setModelInputFromTensor ( modelInput , tensor ) { if ( tensor . dtype === 'string' || tensor . dtype === 'complex64' ) { throw new Error ( `Data type '${ tensor . dtype }' not supported.` ) ; } const modelInputShape = modelInput . shape . split ( ',' ) . map ( dim => Number ( dim ) ) ; if ( ! tensor . shape . every ( ( dim , index ) => modelInputShape [ index ] === - 1 || modelInputShape [ index ] === dim ) ) { throw new Error ( `Input tensor shape mismatch: expect '${ modelInput . shape }', got '${ tensor . shape . join ( ',' ) }'.` ) ; } switch ( modelInput . dataType ) { case 'bool' : case 'int8' : case 'uint8' : case 'int16' : case 'uint32' : case 'int32' : if ( tensor . dtype === 'float32' ) { throw this . getDataTypeMismatchError ( modelInput . dataType , tensor . dtype ) ; } else if ( modelInput . dataType !== tensor . dtype ) { console . warn ( `WARNING: converting '${ tensor . dtype }' to '${ modelInput . dataType }'` ) ; } break ; case 'float32' : case 'float64' : if ( modelInput . dataType !== tensor . dtype ) { console . warn ( `WARNING: converting '${ tensor . dtype }' to '${ modelInput . dataType }'` ) ; } break ; default : break ; } const modelInputBuffer = modelInput . data ( ) ; switch ( modelInput . dataType ) { case 'int8' : modelInputBuffer . set ( Int8Array . from ( tensor . dataSync ( ) ) ) ; break ; case 'uint8' : case 'bool' : modelInputBuffer . set ( Uint8Array . from ( tensor . dataSync ( ) ) ) ; break ; case 'int16' : modelInputBuffer . set ( Int16Array . from ( tensor . dataSync ( ) ) ) ; break ; case 'int32' : modelInputBuffer . set ( Int32Array . from ( tensor . dataSync ( ) ) ) ; break ; case 'uint32' : modelInputBuffer . set ( Uint32Array . from ( tensor . dataSync ( ) ) ) ; break ; case 'float32' : modelInputBuffer . set ( Float32Array . from ( tensor . dataSync ( ) ) ) ; break ; case 'float64' : modelInputBuffer . set ( Float64Array . from ( tensor . dataSync ( ) ) ) ; break ; default : break ; } } private convertTFLiteTensorInfos ( infos ) { return infos . map ( info => { const dtype = getDTypeFromTFLiteType ( info . dataType ) ; return { name : info . name , shape : this . getShapeFromTFLiteTensorInfo ( info ) , dtype , } ; } ) ; } private checkMapInputs ( inputTensorNames , modelInputNames ) { const notInModel = inputTensorNames . filter ( name => ! modelInputNames . includes ( name ) ) ; const notInInput = modelInputNames . filter ( name => ! inputTensorNames . includes ( name ) ) ; if ( notInModel . length === 0 && notInInput . length === 0 ) { return ; } const msgParts = [ 'The model input names don\'t match the model input names.' ] ; if ( notInModel . length > 0 ) { msgParts . push ( `Names in input but missing in model: [${ notInModel }].` ) ; } if ( notInInput . length > 0 ) { msgParts . push ( `Names in model but missing in inputs: [${ notInInput }].` ) ; } throw new Error ( msgParts . join ( ' ' ) ) ; } private getShapeFromTFLiteTensorInfo ( info ) { return info . shape . split ( ',' ) . map ( s => Number ( s ) ) ; } private getDataTypeMismatchError ( expected , got ) { return new Error ( `Data type mismatch: input tensor expects '${ expected }', got '${ got }'` ) ; } } param model param options doc export async function loadTFLiteModel ( model , options ? ) { if ( typeof model === 'string' && model . includes ( 'tfhub.dev' ) && model . includes ( 'lite-model' ) && ! model . endsWith ( TFHUB_SEARCH_PARAM ) ) { model = `${ model }${ TFHUB_SEARCH_PARAM }` ; } const tfliteModelRunner = await tfliteWebAPIClient . tfweb . TFLiteWebModelRunner . create ( model , options ) ; return new TFLiteModel ( tfliteModelRunner ) ; } param tfliteType doc export function getDTypeFromTFLiteType ( tfliteType ) { let dtype ; switch ( tfliteType ) { case 'float32' : case 'float64' : dtype = 'float32' ; break ; case 'int8' : case 'uint8' : case 'int16' : case 'int32' : case 'uint32' : dtype = 'int32' ; break ; case 'bool' : dtype = 'bool' ; break ; default : break ; } return dtype ; }