license import { KernelConfig , NamedAttrMap , NamedTensorInfoMap , scalar , Tensor1D , Tensor2D , TensorInfo } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; interface NonMaxSuppressionWithScoreInputs extends NamedTensorInfoMap { boxes : TensorInfo ; scores : TensorInfo ; } interface NonMaxSuppressionWithScoreAttrs extends NamedAttrMap { maxOutputSize : number ; iouThreshold : number ; scoreThreshold : number ; softNmsSigma : number ; } export const nonMaxSuppressionV5Config = { kernelName : 'NonMaxSuppressionV5' , backendName : 'tensorflow' , kernelFunc : ( { inputs , backend , attrs } ) => { const { boxes , scores } = inputs as NonMaxSuppressionWithScoreInputs ; const { maxOutputSize , iouThreshold , scoreThreshold , softNmsSigma } = attrs as NonMaxSuppressionWithScoreAttrs ; const maxOutputSizeTensor = scalar ( maxOutputSize , 'int32' ) ; const iouThresholdTensor = scalar ( iouThreshold ) ; const scoreThresholdTensor = scalar ( scoreThreshold ) ; const softNmsSigmaTensor = scalar ( softNmsSigma ) ; const opAttrs = [ createTensorsTypeOpAttr ( 'T' , boxes . dtype ) ] ; const nodeBackend = backend as NodeJSKernelBackend ; const [ selectedIndices , selectedScores , validOutputs ] = nodeBackend . executeMultipleOutputs ( 'NonMaxSuppressionV5' , opAttrs , [ boxes as Tensor2D , scores as Tensor1D , maxOutputSizeTensor , iouThresholdTensor , scoreThresholdTensor , softNmsSigmaTensor ] , 3 ) ; maxOutputSizeTensor . dispose ( ) ; iouThresholdTensor . dispose ( ) ; scoreThresholdTensor . dispose ( ) ; softNmsSigmaTensor . dispose ( ) ; validOutputs . dispose ( ) ; return [ selectedIndices , selectedScores ] ; } } ;