license import * as tfconv from '@tensorflow/tfjs-converter' ; import * as tf from '@tensorflow/tfjs-core' ; import { CLASSES } from './classes' ; const BASE_PATH = 'https://storage.googleapis.com/tfjs-models/savedmodel/' ; export { version } from './version' ; docinline export type ObjectDetectionBaseModel = 'mobilenet_v1' | 'mobilenet_v2' | 'lite_mobilenet_v2' ; export interface DetectedObject { bbox : [ number , number , number , number ] ; class : string ; score : number ; } /**
 * Coco-ssd model loading is configurable using the following config dictionary.
 */ export interface ModelConfig { /**
     * It determines wich object detection architecture to load. The supported
     * architectures are: 'mobilenet_v1', 'mobilenet_v2' and 'lite_mobilenet_v2'.
     * It is default to 'lite_mobilenet_v2'.
     */ base ? : ObjectDetectionBaseModel ; /**
     *
     * An optional string that specifies custom url of the model. This is useful
     * for area/countries that don't have access to the model hosted on GCP.
     */ modelUrl ? : string ; } export async function load ( config = { } ) { 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.` ) ; } const base = config . base || 'lite_mobilenet_v2' ; const modelUrl = config . modelUrl ; if ( [ 'mobilenet_v1' , 'mobilenet_v2' , 'lite_mobilenet_v2' ] . indexOf ( base ) === - 1 ) { throw new Error ( `ObjectDetection constructed with invalid base model ` + `${ base }. Valid names are 'mobilenet_v1',` + ` 'mobilenet_v2' and 'lite_mobilenet_v2'.` ) ; } const objectDetection = new ObjectDetection ( base , modelUrl ) ; await objectDetection . load ( ) ; return objectDetection ; } export class ObjectDetection { private modelPath : string ; private model : tfconv . GraphModel ; constructor ( base , modelUrl ? ) { this . modelPath = modelUrl || `${ BASE_PATH }${ this . getPrefix ( base ) }/model.json` ; } private getPrefix ( base ) { return base === 'lite_mobilenet_v2' ? `ssd${ base }` : `ssd_${ base }` ; } async load ( ) { this . model = await tfconv . loadGraphModel ( this . modelPath ) ; const zeroTensor = tf . zeros ( [ 1 , 300 , 300 , 3 ] , 'int32' ) ; const result = await this . model . executeAsync ( zeroTensor ) as tf . Tensor [ ] ; await Promise . all ( result . map ( t => t . data ( ) ) ) ; result . map ( t => t . dispose ( ) ) ; zeroTensor . dispose ( ) ; } param img param maxNumBoxes param minScore private async infer ( img , maxNumBoxes , minScore ) { const batched = tf . tidy ( ( ) => { if ( ! ( img instanceof tf . Tensor ) ) { img = tf . browser . fromPixels ( img ) ; } return tf . expandDims ( img ) ; } ) ; const height = batched . shape [ 1 ] ; const width = batched . shape [ 2 ] ; const result = await this . model . executeAsync ( batched ) as tf . Tensor [ ] ; const scores = result [ 0 ] . dataSync ( ) as Float32Array ; const boxes = result [ 1 ] . dataSync ( ) as Float32Array ; batched . dispose ( ) ; tf . dispose ( result ) ; const [ maxScores , classes ] = this . calculateMaxScores ( scores , result [ 0 ] . shape [ 1 ] , result [ 0 ] . shape [ 2 ] ) ; const prevBackend = tf . getBackend ( ) ; if ( tf . getBackend ( ) === 'webgl' ) { tf . setBackend ( 'cpu' ) ; } const indexTensor = tf . tidy ( ( ) => { const boxes2 = tf . tensor2d ( boxes , [ result [ 1 ] . shape [ 1 ] , result [ 1 ] . shape [ 3 ] ] ) ; return tf . image . nonMaxSuppression ( boxes2 , maxScores , maxNumBoxes , minScore , minScore ) ; } ) ; const indexes = indexTensor . dataSync ( ) as Float32Array ; indexTensor . dispose ( ) ; if ( prevBackend !== tf . getBackend ( ) ) { tf . setBackend ( prevBackend ) ; } return this . buildDetectedObjects ( width , height , boxes , maxScores , indexes , classes ) ; } private buildDetectedObjects ( width , height , boxes , scores , indexes , classes ) { const count = indexes . length ; const objects = [ ] ; for ( let i = 0 ; i < count ; i ++ ) { const bbox = [ ] ; for ( let j = 0 ; j < 4 ; j ++ ) { bbox [ j ] = boxes [ indexes [ i ] * 4 + j ] ; } const minY = bbox [ 0 ] * height ; const minX = bbox [ 1 ] * width ; const maxY = bbox [ 2 ] * height ; const maxX = bbox [ 3 ] * width ; bbox [ 0 ] = minX ; bbox [ 1 ] = minY ; bbox [ 2 ] = maxX - minX ; bbox [ 3 ] = maxY - minY ; objects . push ( { bbox : bbox as [ number , number , number , number ] , class : CLASSES [ classes [ indexes [ i ] ] + 1 ] . displayName , score : scores [ indexes [ i ] ] } ) ; } return objects ; } private calculateMaxScores ( scores , numBoxes , numClasses ) { const maxes = [ ] ; const classes = [ ] ; for ( let i = 0 ; i < numBoxes ; i ++ ) { let max = Number . MIN_VALUE ; let index = - 1 ; for ( let j = 0 ; j < numClasses ; j ++ ) { if ( scores [ i * numClasses + j ] > max ) { max = scores [ i * numClasses + j ] ; index = j ; } } maxes [ i ] = max ; classes [ i ] = index ; } return [ maxes , classes ] ; } param img param maxNumBoxes param minScore async detect ( img , maxNumBoxes = 20 , minScore = 0.5 ) { return this . infer ( img , maxNumBoxes , minScore ) ; } /**
     * Dispose the tensors allocated by the model. You should call this when you
     * are done with the model.
     */ dispose ( ) { if ( this . model != null ) { this . model . dispose ( ) ; } } }