license import * as tfconv from '@tensorflow/tfjs-converter' ; import * as tf from '@tensorflow/tfjs-core' ; import { DeepLabInput , DeepLabOutput , ModelArchitecture , ModelConfig , PredictionConfig } from './types' ; import { getColormap , getLabels , getURL , toInputTensor , toSegmentationImage } from './utils' ; export { version } from './version' ; export { getColormap , getLabels , getURL , ModelConfig , PredictionConfig , toSegmentationImage } ; param input param config return export async function load ( modelConfig = { base : 'pascal' , quantizationBytes : 2 } ) { if ( tf == null ) { throw new Error ( `Cannot find TensorFlow.js.` + ` If you are using a <script> tag, please ` + `also include @tensorflow/tfjs on the page before using this model.` ) ; } if ( modelConfig . base ) { if ( [ 'pascal' , 'cityscapes' , 'ade20k' ] . indexOf ( modelConfig . base ) === - 1 ) { throw new Error ( `SemanticSegmentation cannot be constructed ` + `with an invalid base model ${ modelConfig . base }. ` + `Try one of 'pascal', 'cityscapes' and 'ade20k'.` ) ; } if ( [ 1 , 2 , 4 ] . indexOf ( modelConfig . quantizationBytes ) === - 1 ) { throw new Error ( `Only quantization to 1, 2 or 4 bytes is supported.` ) ; } } else if ( ! modelConfig . modelUrl ) { throw new Error ( `SemanticSegmentation can be constructed either by passing ` + `the weights URL or one of the supported base model names from ` + `'pascal', 'cityscapes' and 'ade20k',` + `together with the degree of quantization (either 1, 2 or 4).` + `Aborting, since neither has been provided.` ) ; } const graphModel = await tfconv . loadGraphModel ( modelConfig . modelUrl || getURL ( modelConfig . base , modelConfig . quantizationBytes ) ) ; const deeplab = new SemanticSegmentation ( graphModel , modelConfig . base ) ; return deeplab ; } export class SemanticSegmentation { readonly model : tfconv . GraphModel ; readonly base : ModelArchitecture ; constructor ( graphModel , base ? ) { this . model = graphModel ; this . base = base ; } param input return public predict ( input ) { return tf . tidy ( ( ) => { const data = tf . cast ( toInputTensor ( input ) , 'int32' ) ; return tf . squeeze ( this . model . execute ( data ) as tf . Tensor ) ; } ) ; } param image param config returns public async segment ( input , config = { } ) { if ( ! ( ( config . colormap && config . labels ) || this . base ) ) { throw new Error ( `Calling the 'segment' method requires either the 'base'` + ` attribute to be defined ` + `(e.g. 'pascal', 'cityscapes' or'ade20k'),` + ` or 'colormap' and 'labels' options to be set. ` + `Aborting, since neither has been provided.` ) ; } else if ( ! ( config . colormap && config . labels ) ) { config . colormap = getColormap ( this . base ) ; config . labels = getLabels ( this . base ) ; } const { colormap , labels , canvas } = config ; const rawSegmentationMap = tf . tidy ( ( ) => this . predict ( input ) ) ; const [ height , width ] = rawSegmentationMap . shape ; const { legend , segmentationMap } = await toSegmentationImage ( colormap , labels , rawSegmentationMap , canvas ) ; tf . dispose ( rawSegmentationMap ) ; return { legend , height , width , segmentationMap } ; } /**
     * Dispose of the tensors allocated by the model.
     * You should call this when you are done with the model.
     */ public async dispose ( ) { if ( this . model ) { this . model . dispose ( ) ; } } }