license import * as tf from '@tensorflow/tfjs' ; export type MobileNetMultiplier = 0.50 | 0.75 | 1.0 | 1.01 ; export type ConvolutionType = 'conv2d' | 'separableConv' ; export type ConvolutionDefinition = [ ConvolutionType , number ] ; export type OutputStride = 32 | 16 | 8 ; const mobileNet100Architecture = [ [ 'conv2d' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] ] ; const mobileNet75Architecture = [ [ 'conv2d' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] ] ; const mobileNet50Architecture = [ [ 'conv2d' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 2 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] , [ 'separableConv' , 1 ] ] ; const VALID_OUTPUT_STRIDES = [ 8 , 16 , 32 ] ; export function assertValidOutputStride ( outputStride ) { tf . util . assert ( typeof outputStride === 'number' , 'outputStride is not a number' ) ; tf . util . assert ( VALID_OUTPUT_STRIDES . indexOf ( outputStride ) >= 0 , `outputStride of ${ outputStride } is invalid. ` + `It must be either 8, 16, or 32` ) ; } export function assertValidResolution ( resolution , outputStride ) { tf . util . assert ( typeof resolution === 'number' , 'resolution is not a number' ) ; tf . util . assert ( ( resolution - 1 ) % outputStride === 0 , `resolution of ${ resolution } is invalid for output stride ` + `${ outputStride }.` ) ; } export function assertValidScaleFactor ( imageScaleFactor ) { tf . util . assert ( typeof imageScaleFactor === 'number' , 'imageScaleFactor is not a number' ) ; tf . util . assert ( imageScaleFactor >= 0.2 && imageScaleFactor <= 1.0 , 'imageScaleFactor must be between 0.2 and 1.0' ) ; } export const mobileNetArchitectures = { 100 : mobileNet100Architecture , 75 : mobileNet75Architecture , 50 : mobileNet50Architecture } ; type Layer = { blockId : number ; stride : number ; outputStride : number ; convType : ConvolutionType ; rate : number ; } ; /**
 * Takes a mobilenet architectures' convolution definitions and converts them
 * into definitions for convolutional layers that will generate outputs with the
 * desired output stride. It does this by reducing the input stride in certain
 * layers and applying atrous convolution in subsequent layers. Raises an error
 * if the output stride is not possible with the architecture.
 */ function toOutputStridedLayers ( convolutionDefinition , outputStride ) { let currentStride = 1 ; let rate = 1 ; return convolutionDefinition . map ( ( [ convType , stride ] , blockId ) : Layer => { let layerStride , layerRate ; if ( currentStride === outputStride ) { layerStride = 1 ; layerRate = rate ; rate *= stride ; } else { layerStride = stride ; layerRate = 1 ; currentStride *= stride ; } return { blockId , convType , stride : layerStride , rate : layerRate , outputStride : currentStride } ; } ) ; } export class MobileNet { private variables : { [ varName ] : tf . Tensor ; } ; private convolutionDefinitions : ConvolutionDefinition [ ] ; private PREPROCESS_DIVISOR = tf . scalar ( 255.0 / 2 ) ; private ONE = tf . scalar ( 1 ) ; constructor ( variables , convolutionDefinitions ) { this . variables = variables ; this . convolutionDefinitions = convolutionDefinitions ; } predict ( input , outputStride ) { const preprocessedInput = tf . cast ( input , 'float32' ) . div ( this . PREPROCESS_DIVISOR ) . sub ( this . ONE ) as tf . Tensor3D ; const layers = toOutputStridedLayers ( this . convolutionDefinitions , outputStride ) ; return layers . reduce ( ( previousLayer , { blockId , stride , convType , rate } ) => { if ( convType === 'conv2d' ) { return this . conv ( previousLayer , stride , blockId ) ; } else if ( convType === 'separableConv' ) { return this . separableConv ( previousLayer , stride , blockId , rate ) ; } else { throw Error ( 'Unknown conv type of ' + convType ) ; } } , preprocessedInput ) ; } public convToOutput ( mobileNetOutput , outputLayerName ) { return mobileNetOutput . conv2d ( this . weights ( outputLayerName ) , 1 , 'same' ) . add ( this . biases ( outputLayerName ) ) as tf . Tensor3D ; } private conv ( inputs , stride , blockId ) { return inputs . conv2d ( this . weights ( `Conv2d_${ String ( blockId ) }` ) , stride , 'same' ) . add ( this . biases ( `Conv2d_${ String ( blockId ) }` ) ) . clipByValue ( 0 , 6 ) as tf . Tensor3D ; } private separableConv ( inputs , stride , blockID , dilations = 1 ) { const dwLayer = `Conv2d_${ String ( blockID ) }_depthwise` ; const pwLayer = `Conv2d_${ String ( blockID ) }_pointwise` ; const x1 = inputs . depthwiseConv2D ( this . depthwiseWeights ( dwLayer ) , stride , 'same' , 'NHWC' , dilations ) . add ( this . biases ( dwLayer ) ) . clipByValue ( 0 , 6 ) as tf . Tensor3D ; const x2 = x1 . conv2d ( this . weights ( pwLayer ) , [ 1 , 1 ] , 'same' ) . add ( this . biases ( pwLayer ) ) . clipByValue ( 0 , 6 ) as tf . Tensor3D ; return x2 ; } private weights ( layerName ) { return this . variables [ `MobilenetV1/${ layerName }/weights` ] as tf . Tensor4D ; } private biases ( layerName ) { return this . variables [ `MobilenetV1/${ layerName }/biases` ] as tf . Tensor1D ; } private depthwiseWeights ( layerName ) { return this . variables [ `MobilenetV1/${ layerName }/depthwise_weights` ] as tf . Tensor4D ; } dispose ( ) { for ( const varName in this . variables ) { this . variables [ varName ] . dispose ( ) ; } } }