license import * as tf from '@tensorflow/tfjs' ; import * as fs from 'fs' ; import { dirname , join , resolve } from 'path' ; import { promisify } from 'util' ; import { getModelArtifactsInfoForJSON , toArrayBuffer } from './io_utils' ; const stat = promisify ( fs . stat ) ; const writeFile = promisify ( fs . writeFile ) ; const readFile = promisify ( fs . readFile ) ; const mkdir = promisify ( fs . mkdir ) ; function doesNotExistHandler ( name ) { return e => { switch ( e . code ) { case 'ENOENT' : throw new Error ( `${ name } ${ e . path } does not exist: loading failed` ) ; default : throw e ; } } ; } export class NodeFileSystem implements tf . io . IOHandler { static readonly URL_SCHEME = 'file://' ; protected readonly path : string | string [ ] ; readonly MODEL_JSON_FILENAME = 'model.json' ; readonly WEIGHTS_BINARY_FILENAME = 'weights.bin' ; readonly MODEL_BINARY_FILENAME = 'tensorflowjs.pb' ; param path constructor ( path ) { if ( Array . isArray ( path ) ) { tf . util . assert ( path . length === 2 , ( ) => 'file paths must have a length of 2, ' + `(actual length is ${ path . length }).` ) ; this . path = path . map ( p => resolve ( p ) ) ; } else { this . path = resolve ( path ) ; } } async save ( modelArtifacts ) { if ( Array . isArray ( this . path ) ) { throw new Error ( 'Cannot perform saving to multiple paths.' ) ; } await this . createOrVerifyDirectory ( ) ; if ( modelArtifacts . modelTopology instanceof ArrayBuffer ) { throw new Error ( 'NodeFileSystem.save() does not support saving model topology ' + 'in binary format yet.' ) ; } else { const weightsBinPath = join ( this . path , this . WEIGHTS_BINARY_FILENAME ) ; const weightsManifest = [ { paths : [ this . WEIGHTS_BINARY_FILENAME ] , weights : modelArtifacts . weightSpecs } ] ; const modelJSON = { modelTopology : modelArtifacts . modelTopology , weightsManifest , format : modelArtifacts . format , generatedBy : modelArtifacts . generatedBy , convertedBy : modelArtifacts . convertedBy } ; if ( modelArtifacts . trainingConfig != null ) { modelJSON . trainingConfig = modelArtifacts . trainingConfig ; } if ( modelArtifacts . signature != null ) { modelJSON . signature = modelArtifacts . signature ; } if ( modelArtifacts . userDefinedMetadata != null ) { modelJSON . userDefinedMetadata = modelArtifacts . userDefinedMetadata ; } const modelJSONPath = join ( this . path , this . MODEL_JSON_FILENAME ) ; await writeFile ( modelJSONPath , JSON . stringify ( modelJSON ) , 'utf8' ) ; await writeFile ( weightsBinPath , Buffer . from ( modelArtifacts . weightData ) , 'binary' ) ; return { modelArtifactsInfo : getModelArtifactsInfoForJSON ( modelArtifacts ) as any } ; } } async load ( ) { return Array . isArray ( this . path ) ? this . loadBinaryModel ( ) : this . loadJSONModel ( ) ; } protected async loadBinaryModel ( ) { const topologyPath = this . path [ 0 ] ; const weightManifestPath = this . path [ 1 ] ; const topology = await stat ( topologyPath ) . catch ( doesNotExistHandler ( 'Topology Path' ) ) ; const weightManifest = await stat ( weightManifestPath ) . catch ( doesNotExistHandler ( 'Weight Manifest Path' ) ) ; if ( ! topology . isFile ( ) ) { throw new Error ( 'File specified for topology is not a file!' ) ; } if ( ! weightManifest . isFile ( ) ) { throw new Error ( 'File specified for the weight manifest is not a file!' ) ; } const modelTopology = await readFile ( this . path [ 0 ] ) ; const weightsManifest = JSON . parse ( await readFile ( this . path [ 1 ] , 'utf8' ) ) ; const modelArtifacts = { modelTopology , } ; const [ weightSpecs , weightData ] = await this . loadWeights ( weightsManifest , this . path [ 1 ] ) ; modelArtifacts . weightSpecs = weightSpecs ; modelArtifacts . weightData = weightData ; return modelArtifacts ; } protected async loadJSONModel ( ) { const path = this . path as string ; const info = await stat ( path ) . catch ( doesNotExistHandler ( 'Path' ) ) ; if ( info . isFile ( ) ) { const modelJSON = JSON . parse ( await readFile ( path , 'utf8' ) ) ; return tf . io . getModelArtifactsForJSON ( modelJSON , ( weightsManifest ) => this . loadWeights ( weightsManifest , path ) ) ; } else { throw new Error ( 'The path to load from must be a file. Loading from a directory ' + 'is not supported.' ) ; } } private async loadWeights ( weightsManifest , path ) { const dirName = dirname ( path ) ; const buffers = [ ] ; const weightSpecs = [ ] ; for ( const group of weightsManifest ) { for ( const path of group . paths ) { const weightFilePath = join ( dirName , path ) ; const buffer = await readFile ( weightFilePath ) . catch ( doesNotExistHandler ( 'Weight file' ) ) ; buffers . push ( buffer ) ; } weightSpecs . push ( ... group . weights ) ; } return [ weightSpecs , toArrayBuffer ( buffers ) ] ; } /**
     * For each item in `this.path`, creates a directory at the path or verify
     * that the path exists as a directory.
     */ protected async createOrVerifyDirectory ( ) { const paths = Array . isArray ( this . path ) ? this . path : [ this . path ] ; for ( const path of paths ) { try { await mkdir ( path ) ; } catch ( e ) { if ( e . code === 'EEXIST' ) { if ( ( await stat ( path ) ) . isFile ( ) ) { throw new Error ( `Path ${ path } exists as a file. The path must be ` + `nonexistent or point to a directory.` ) ; } } else { throw e ; } } } } } export const nodeFileSystemRouter = ( url ) => { if ( Array . isArray ( url ) ) { if ( url . every ( urlElement => urlElement . startsWith ( NodeFileSystem . URL_SCHEME ) ) ) { return new NodeFileSystem ( url . map ( urlElement => urlElement . slice ( NodeFileSystem . URL_SCHEME . length ) ) ) ; } else { return null ; } } else { if ( url . startsWith ( NodeFileSystem . URL_SCHEME ) ) { return new NodeFileSystem ( url . slice ( NodeFileSystem . URL_SCHEME . length ) ) ; } else { return null ; } } } ; param path export function fileSystem ( path ) { return new NodeFileSystem ( path ) ; }