license import { io } from '@tensorflow/tfjs-core' ; import { fromByteArray , toByteArray } from 'base64-js' ; import { getModelArtifactsInfoForJSON } from './model_artifacts' ; type StorageKeys = { info : string ; modelArtifactsWithoutWeights : string ; weightData : string ; weightDataChunkSize : string ; } ; const PATH_SEPARATOR = '/' ; const PATH_PREFIX = 'tensorflowjs_models' ; const INFO_SUFFIX = 'info' ; const MODEL_SUFFIX = 'model_without_weight' ; const WEIGHT_DATA_SUFFIX = 'weight_data' ; const WEIGHT_DATA_SIZE_SUFFIX = 'weight_data_size' ; function getModelKeys ( path ) { return { info : [ PATH_PREFIX , path , INFO_SUFFIX ] . join ( PATH_SEPARATOR ) , modelArtifactsWithoutWeights : [ PATH_PREFIX , path , MODEL_SUFFIX ] . join ( PATH_SEPARATOR ) , weightData : [ PATH_PREFIX , path , WEIGHT_DATA_SUFFIX ] . join ( PATH_SEPARATOR ) , weightDataChunkSize : [ PATH_PREFIX , path , WEIGHT_DATA_SIZE_SUFFIX ] . join ( PATH_SEPARATOR ) } ; } param str param size function splitString ( str , size ) { return str . match ( new RegExp ( `.{1,${ size }}` , 'g' ) ) ; } function getStorage ( key ) { return new Promise ( ( resolve , reject ) => { wx . getStorage ( { key , success : ( res ) => resolve ( res . data ) , fail : ( res ) => reject ( res ) } ) ; } ) ; } class LocalStorageHandler implements io . IOHandler { protected readonly keys : StorageKeys ; constructor ( protected readonly modelPath ) { if ( modelPath == null || ! modelPath ) { throw new Error ( 'modelPath must not be null, undefined or empty.' ) ; } this . keys = getModelKeys ( this . modelPath ) ; } param modelArtifacts returns async save ( modelArtifacts ) { if ( modelArtifacts . modelTopology instanceof ArrayBuffer ) { throw new Error ( 'AsyncStorageHandler.save() does not support saving model topology ' + 'in binary format.' ) ; } else { const modelArtifactsInfo = getModelArtifactsInfoForJSON ( modelArtifacts ) ; const { weightData , ... modelArtifactsWithoutWeights } = modelArtifacts ; const weights = splitString ( fromByteArray ( new Uint8Array ( weightData ) ) , 800 * 1024 ) ; try { wx . setStorageSync ( this . keys . info , JSON . stringify ( modelArtifactsInfo ) ) ; wx . setStorageSync ( this . keys . modelArtifactsWithoutWeights , JSON . stringify ( modelArtifactsWithoutWeights ) ) ; wx . setStorageSync ( this . keys . weightDataChunkSize , weights . length ) ; weights . forEach ( ( value , index ) => { wx . setStorageSync ( `${ this . keys . weightData }:${ index }` , value ) ; } ) ; return { modelArtifactsInfo } ; } catch ( err ) { wx . removeStorageSync ( this . keys . info ) ; weights . forEach ( ( value , index ) => { wx . removeStorageSync ( `${ this . keys . weightData }:${ index }` ) ; } ) ; wx . removeStorageSync ( this . keys . weightDataChunkSize ) ; wx . removeStorageSync ( this . keys . modelArtifactsWithoutWeights ) ; throw new Error ( `Failed to save model '${ this . modelPath }' to local storage.
            Error info ${ err }` ) ; } } } returns async load ( ) { const info = JSON . parse ( await getStorage ( this . keys . info ) ) as io . ModelArtifactsInfo ; if ( info == null ) { throw new Error ( `In local storage, there is no model with name '${ this . modelPath }'` ) ; } if ( info . modelTopologyType !== 'JSON' ) { throw new Error ( 'BrowserLocalStorage does not support loading non-JSON model ' + 'topology yet.' ) ; } const modelArtifacts = JSON . parse ( await getStorage ( this . keys . modelArtifactsWithoutWeights ) ) ; const weightDataSize = await getStorage ( this . keys . weightDataChunkSize ) ; let weightDataBase64 = '' ; for ( let i = 0 ; i < weightDataSize ; i ++ ) { const partialData = await getStorage ( `${ this . keys . weightData }:${ i }` ) ; if ( partialData == null ) { throw new Error ( `In local storage, the binary weight values of model ` + `'${ this . modelPath }' are missing.` ) ; } weightDataBase64 += partialData ; } modelArtifacts . weightData = toByteArray ( weightDataBase64 ) . buffer ; return modelArtifacts ; } } param modelPath returns export function localStorageIO ( modelPath ) { return new LocalStorageHandler ( modelPath ) ; }