import * as tf from '@tensorflow/tfjs-core' ; import { BoundingBox , Dimensions , NetInput , NeuralNetwork , nonMaxSuppression , normalize , ObjectDetection , sigmoid , TNetInput , toNetInput , } from 'tfjs-image-recognition-base' ; import { convLayer } from '../common' ; import { ConvParams , SeparableConvParams } from '../common/types' ; import { TinyYolov2Config , validateConfig } from './config' ; import { DEFAULT_FILTER_SIZES } from './const' ; import { convWithBatchNorm } from './convWithBatchNorm' ; import { depthwiseSeparableConv } from './depthwiseSeparableConv' ; import { extractParams } from './extractParams' ; import { extractParamsFromWeigthMap } from './extractParamsFromWeigthMap' ; import { leaky } from './leaky' ; import { ITinyYolov2Options , TinyYolov2Options } from './TinyYolov2Options' ; import { DefaultTinyYolov2NetParams , MobilenetParams , TinyYolov2NetParams } from './types' ; export class TinyYolov2 extends NeuralNetwork < TinyYolov2NetParams > { private _config : TinyYolov2Config ; constructor ( config ) { super ( 'TinyYolov2' ) ; validateConfig ( config ) ; this . _config = config ; } public get config ( ) : TinyYolov2Config { return this . _config ; } public get withClassScores ( ) : boolean { return this . config . withClassScores || this . config . classes . length > 1 ; } public get boxEncodingSize ( ) : number { return 5 + ( this . withClassScores ? this . config . classes . length : 0 ) ; } public runTinyYolov2 ( x , params ) { let out = convWithBatchNorm ( x , params . conv0 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = convWithBatchNorm ( out , params . conv1 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = convWithBatchNorm ( out , params . conv2 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = convWithBatchNorm ( out , params . conv3 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = convWithBatchNorm ( out , params . conv4 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = convWithBatchNorm ( out , params . conv5 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 1 , 1 ] , 'same' ) ; out = convWithBatchNorm ( out , params . conv6 ) ; out = convWithBatchNorm ( out , params . conv7 ) ; return convLayer ( out , params . conv8 , 'valid' , false ) ; } public runMobilenet ( x , params ) { let out = this . config . isFirstLayerConv2d ? leaky ( convLayer ( x , params . conv0 as ConvParams , 'valid' , false ) ) : depthwiseSeparableConv ( x , params . conv0 as SeparableConvParams ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = depthwiseSeparableConv ( out , params . conv1 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = depthwiseSeparableConv ( out , params . conv2 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = depthwiseSeparableConv ( out , params . conv3 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = depthwiseSeparableConv ( out , params . conv4 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 2 , 2 ] , 'same' ) ; out = depthwiseSeparableConv ( out , params . conv5 ) ; out = tf . maxPool ( out , [ 2 , 2 ] , [ 1 , 1 ] , 'same' ) ; out = params . conv6 ? depthwiseSeparableConv ( out , params . conv6 ) : out ; out = params . conv7 ? depthwiseSeparableConv ( out , params . conv7 ) : out ; return convLayer ( out , params . conv8 , 'valid' , false ) ; } public forwardInput ( input , inputSize ) { const { params } = this ; if ( ! params ) { throw new Error ( 'TinyYolov2 - load model before inference' ) ; } return tf . tidy ( ( ) => { let batchTensor = input . toBatchTensor ( inputSize , false ) . toFloat ( ) ; batchTensor = this . config . meanRgb ? normalize ( batchTensor , this . config . meanRgb ) : batchTensor ; batchTensor = batchTensor . div ( tf . scalar ( 256 ) ) as tf . Tensor4D ; return this . config . withSeparableConvs ? this . runMobilenet ( batchTensor , params as MobilenetParams ) : this . runTinyYolov2 ( batchTensor , params as DefaultTinyYolov2NetParams ) ; } ) ; } public async forward ( input , inputSize ) { return await this . forwardInput ( await toNetInput ( input ) , inputSize ) ; } public async detect ( input , forwardParams = { } ) { const { inputSize , scoreThreshold } = new TinyYolov2Options ( forwardParams ) ; const netInput = await toNetInput ( input ) ; const out = await this . forwardInput ( netInput , inputSize ) ; const out0 = tf . tidy ( ( ) => tf . unstack ( out ) [ 0 ] . expandDims ( ) ) as tf . Tensor4D ; const inputDimensions = { width : netInput . getInputWidth ( 0 ) , height : netInput . getInputHeight ( 0 ) } ; const results = this . extractBoxes ( out0 , netInput . getReshapedInputDimensions ( 0 ) , scoreThreshold ) ; out . dispose ( ) ; out0 . dispose ( ) ; const boxes = results . map ( res => res . box ) ; const scores = results . map ( res => res . score ) ; const classScores = results . map ( res => res . classScore ) ; const classNames = results . map ( res => this . config . classes [ res . label ] ) ; const indices = nonMaxSuppression ( boxes . map ( box => box . rescale ( inputSize ) ) , scores , this . config . iouThreshold , true ) ; const detections = indices . map ( idx => new ObjectDetection ( scores [ idx ] , classScores [ idx ] , classNames [ idx ] , boxes [ idx ] , inputDimensions ) ) ; return detections ; } protected getDefaultModelName ( ) { return '' ; } protected extractParamsFromWeigthMap ( weightMap ) { return extractParamsFromWeigthMap ( weightMap , this . config ) ; } protected extractParams ( weights ) { const filterSizes = this . config . filterSizes || DEFAULT_FILTER_SIZES ; const numFilters = filterSizes ? filterSizes . length : undefined ; if ( numFilters !== 7 && numFilters !== 8 && numFilters !== 9 ) { throw new Error ( `TinyYolov2 - expected 7 | 8 | 9 convolutional filters, but found ${ numFilters } filterSizes in config` ) ; } return extractParams ( weights , this . config , this . boxEncodingSize , filterSizes ) ; } protected extractBoxes ( outputTensor , inputBlobDimensions , scoreThreshold ? ) { const { width , height } = inputBlobDimensions ; const inputSize = Math . max ( width , height ) ; const correctionFactorX = inputSize / width ; const correctionFactorY = inputSize / height ; const numCells = outputTensor . shape [ 1 ] ; const numBoxes = this . config . anchors . length ; const [ boxesTensor , scoresTensor , classScoresTensor ] = tf . tidy ( ( ) => { const reshaped = outputTensor . reshape ( [ numCells , numCells , numBoxes , this . boxEncodingSize ] ) ; const boxes = reshaped . slice ( [ 0 , 0 , 0 , 0 ] , [ numCells , numCells , numBoxes , 4 ] ) ; const scores = reshaped . slice ( [ 0 , 0 , 0 , 4 ] , [ numCells , numCells , numBoxes , 1 ] ) ; const classScores = this . withClassScores ? tf . softmax ( reshaped . slice ( [ 0 , 0 , 0 , 5 ] , [ numCells , numCells , numBoxes , this . config . classes . length ] ) , 3 ) : tf . scalar ( 0 ) ; return [ boxes , scores , classScores ] ; } ) ; const results = [ ] ; for ( let row = 0 ; row < numCells ; row ++ ) { for ( let col = 0 ; col < numCells ; col ++ ) { for ( let anchor = 0 ; anchor < numBoxes ; anchor ++ ) { const score = sigmoid ( scoresTensor . get ( row , col , anchor , 0 ) ) ; if ( ! scoreThreshold || score > scoreThreshold ) { const ctX = ( ( col + sigmoid ( boxesTensor . get ( row , col , anchor , 0 ) ) ) / numCells ) * correctionFactorX ; const ctY = ( ( row + sigmoid ( boxesTensor . get ( row , col , anchor , 1 ) ) ) / numCells ) * correctionFactorY ; const width = ( ( Math . exp ( boxesTensor . get ( row , col , anchor , 2 ) ) * this . config . anchors [ anchor ] . x ) / numCells ) * correctionFactorX ; const height = ( ( Math . exp ( boxesTensor . get ( row , col , anchor , 3 ) ) * this . config . anchors [ anchor ] . y ) / numCells ) * correctionFactorY ; const x = ( ctX - ( width / 2 ) ) ; const y = ( ctY - ( height / 2 ) ) ; const pos = { row , col , anchor } ; const { classScore , label } = this . withClassScores ? this . extractPredictedClass ( classScoresTensor as tf . Tensor4D , pos ) : { classScore : 1 , label : 0 } ; results . push ( { box : new BoundingBox ( x , y , x + width , y + height ) , score : score , classScore : score * classScore , label , ... pos } ) ; } } } } boxesTensor . dispose ( ) ; scoresTensor . dispose ( ) ; classScoresTensor . dispose ( ) ; return results ; } private extractPredictedClass ( classesTensor , pos ) { const { row , col , anchor } = pos ; return Array ( this . config . classes . length ) . fill ( 0 ) . map ( ( _ , i ) => classesTensor . get ( row , col , anchor , i ) ) . map ( ( classScore , label ) => ( { classScore , label } ) ) . reduce ( ( max , curr ) => max . classScore > curr . classScore ? max : curr ) ; } }