Split NonMaxSuppressionV3 and NonMaxSuppresionV5 to two APIs

tfjs-core/src/backends/backend.ts

@@ -590,8 +590,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
 
   nonMaxSuppression(
       boxes: Tensor2D, scores: Tensor1D, maxOutputSize: number,
-      iouThreshold: number, scoreThreshold?: number,
-      softNmsSigma?: number): Tensor1D {
+      iouThreshold: number, scoreThreshold?: number): Tensor1D {
     return notYetImplemented('nonMaxSuppression');
   }
 

tfjs-core/src/backends/cpu/all_kernels.ts

@@ -19,3 +19,4 @@
 // global registry when we compile the library. A modular build would replace
 // the contents of this file and import only the kernels that are needed.
 import './square';
+import './non_max_suppression_v5';

tfjs-core/src/backends/cpu/backend_cpu.ts

@@ -42,7 +42,7 @@ import {getArrayFromDType, inferDtype, now, sizeFromShape} from '../../util';
 import {BackendTimingInfo, DataStorage, EPSILON_FLOAT32, KernelBackend} from '../backend';
 import * as backend_util from '../backend_util';
 import * as complex_util from '../complex_util';
-import {nonMaxSuppressionImpl} from '../non_max_suppression_impl';
+import {nonMaxSuppressionV3} from '../non_max_suppression_impl';
 import {split} from '../split_shared';
 import {tile} from '../tile_impl';
 import {topkImpl} from '../topk_impl';
@@ -3270,15 +3270,13 @@ export class MathBackendCPU extends KernelBackend {
 
   nonMaxSuppression(
       boxes: Tensor2D, scores: Tensor1D, maxOutputSize: number,
-      iouThreshold: number, scoreThreshold: number,
-      softNmsSigma: number): Tensor1D {
+      iouThreshold: number, scoreThreshold: number): Tensor1D {
     assertNotComplex(boxes, 'nonMaxSuppression');
 
     const boxesVals = this.readSync(boxes.dataId) as TypedArray;
     const scoresVals = this.readSync(scores.dataId) as TypedArray;
-    return nonMaxSuppressionImpl(
-        boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold,
-        softNmsSigma);
+    return nonMaxSuppressionV3(
+        boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold);
   }
 
   fft(x: Tensor2D): Tensor2D {

tfjs-core/src/backends/cpu/non_max_suppression_v5.ts

@@ -0,0 +1,63 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
+import {TypedArray} from '../../types';
+import {nonMaxSuppressionV5} from '../non_max_suppression_impl';
+
+import {MathBackendCPU} from './backend_cpu';
+import {assertNotComplex} from './cpu_util';
+
+interface NonMaxSuppressionWithScoreInputs extends NamedTensorInfoMap {
+  boxes: TensorInfo;
+  scores: TensorInfo;
+}
+
+interface NonMaxSuppressionWithScoreAttrs extends NamedAttrMap {
+  maxOutputSize: number;
+  iouThreshold: number;
+  scoreThreshold: number;
+  softNmsSigma: number;
+}
+
+registerKernel({
+  kernelName: 'NonMaxSuppressionV5',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend, attrs}) => {
+    const {boxes, scores} = inputs as NonMaxSuppressionWithScoreInputs;
+    const {maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma} =
+        attrs as NonMaxSuppressionWithScoreAttrs;
+
+    const cpuBackend = backend as MathBackendCPU;
+
+    assertNotComplex(boxes, 'NonMaxSuppressionWithScore');
+
+    const boxesVals = cpuBackend.data.get(boxes.dataId).values as TypedArray;
+    const scoresVals = cpuBackend.data.get(scores.dataId).values as TypedArray;
+
+    const maxOutputSizeVal = maxOutputSize;
+    const iouThresholdVal = iouThreshold;
+    const scoreThresholdVal = scoreThreshold;
+    const softNmsSigmaVal = softNmsSigma;
+
+    const {selectedIndices, selectedScores} = nonMaxSuppressionV5(
+        boxesVals, scoresVals, maxOutputSizeVal, iouThresholdVal,
+        scoreThresholdVal, softNmsSigmaVal);
+
+    return [selectedIndices, selectedScores];
+  }
+});

tfjs-core/src/backends/non_max_suppression_impl.ts

@@ -19,8 +19,9 @@
  * Implementation of the NonMaxSuppression kernel shared between webgl and cpu.
  */
 
-import {tensor1d} from '../ops/tensor_ops';
+import {scalar, tensor1d} from '../ops/tensor_ops';
 import {Tensor1D} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
 import {TypedArray} from '../types';
 
 import {binaryInsert} from './array_util';
@@ -31,10 +32,41 @@ interface Candidate {
   suppressBeginIndex: number;
 }
 
-export function nonMaxSuppressionImpl(
+export function nonMaxSuppressionV3(
+    boxes: TypedArray, scores: TypedArray, maxOutputSize: number,
+    iouThreshold: number, scoreThreshold: number): Tensor1D {
+  const dummySoftNmsSigma = 0.0;
+
+  return nonMaxSuppressionImpl_(
+             boxes, scores, maxOutputSize, iouThreshold, scoreThreshold,
+             dummySoftNmsSigma)
+             .selectedIndices as Tensor1D;
+}
+
+export function nonMaxSuppressionV5(
     boxes: TypedArray, scores: TypedArray, maxOutputSize: number,
     iouThreshold: number, scoreThreshold: number,
-    softNmsSigma: number): Tensor1D {
+    softNmsSigma: number): NamedTensorMap {
+  // For NonMaxSuppressionV5Op, we always return a second output holding
+  // corresponding scores.
+  const returnScoresTensor = true;
+
+  const result = nonMaxSuppressionImpl_(
+      boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma,
+      returnScoresTensor);
+
+  result.numValidOutputs.dispose();
+
+  return {
+    selectedIndices: result.selectedIndices,
+    selectedScores: result.selectedScores
+  };
+}
+
+function nonMaxSuppressionImpl_(
+    boxes: TypedArray, scores: TypedArray, maxOutputSize: number,
+    iouThreshold: number, scoreThreshold: number, softNmsSigma: number,
+    returnScoresTensor = false, padToMaxOutputSize = false): NamedTensorMap {
   // The list is sorted in ascending order, so that we can always pop the
   // candidate with the largest score in O(1) time.
   const candidates =
@@ -47,9 +79,10 @@ export function nonMaxSuppressionImpl(
   // before.
   const scale = softNmsSigma > 0 ? (-0.5 / softNmsSigma) : 0.0;
 
-  const selected: number[] = [];
+  const selectedIndices: number[] = [];
+  const selectedScores: number[] = [];
 
-  while (selected.length < maxOutputSize && candidates.length > 0) {
+  while (selectedIndices.length < maxOutputSize && candidates.length > 0) {
     const candidate = candidates.pop();
     const {score: originalScore, boxIndex, suppressBeginIndex} = candidate;
 
@@ -59,13 +92,13 @@ export function nonMaxSuppressionImpl(
 
     // Overlapping boxes are likely to have similar scores, therefore we
     // iterate through the previously selected boxes backwards in order to
-    // see if candidate's score should be suppressed. We use suppressBeginIndex
-    // to track and ensure a candidate can be suppressed by a selected box no
-    // more than once. Also, if the overlap exceeds iouThreshold, we simply
-    // ignore the candidate.
+    // see if candidate's score should be suppressed. We use
+    // suppressBeginIndex to track and ensure a candidate can be suppressed
+    // by a selected box no more than once. Also, if the overlap exceeds
+    // iouThreshold, we simply ignore the candidate.
     let ignoreCandidate = false;
-    for (let j = selected.length - 1; j >= suppressBeginIndex; --j) {
-      const iou = intersectionOverUnion(boxes, boxIndex, selected[j]);
+    for (let j = selectedIndices.length - 1; j >= suppressBeginIndex; --j) {
+      const iou = intersectionOverUnion(boxes, boxIndex, selectedIndices[j]);
 
       if (iou >= iouThreshold) {
         ignoreCandidate = true;
@@ -81,19 +114,20 @@ export function nonMaxSuppressionImpl(
     }
 
     // At this point, if `candidate.score` has not dropped below
-    // `scoreThreshold`, then we know that we went through all of the previous
-    // selections and can safely update `suppressBeginIndex` to the end of the
-    // selected array. Then we can re-insert the candidate with the updated
-    // score and suppressBeginIndex back in the candidate list. If on the other
-    // hand, `candidate.score` has dropped below the score threshold, we will
-    // not add it back to the candidates list.
-    candidate.suppressBeginIndex = selected.length;
+    // `scoreThreshold`, then we know that we went through all of the
+    // previous selections and can safely update `suppressBeginIndex` to the
+    // end of the selected array. Then we can re-insert the candidate with
+    // the updated score and suppressBeginIndex back in the candidate list.
+    // If on the other hand, `candidate.score` has dropped below the score
+    // threshold, we will not add it back to the candidates list.
+    candidate.suppressBeginIndex = selectedIndices.length;
 
     if (!ignoreCandidate) {
-      // Candidate has passed all the tests, and is not suppressed, so select
-      // the candidate.
+      // Candidate has passed all the tests, and is not suppressed, so
+      // select the candidate.
       if (candidate.score === originalScore) {
-        selected.push(boxIndex);
+        selectedIndices.push(boxIndex);
+        selectedScores.push(candidate.score);
       } else if (candidate.score > scoreThreshold) {
         // Candidate's score is suppressed but is still high enough to be
         // considered, so add back to the candidates list.
@@ -102,7 +136,18 @@ export function nonMaxSuppressionImpl(
     }
   }
 
-  return tensor1d(selected, 'int32');
+  // NonMaxSuppressionV4 feature: padding output to maxOutputSize.
+  const numValidOutputs = selectedIndices.length;
+  if (padToMaxOutputSize) {
+    selectedIndices.fill(0, numValidOutputs);
+    selectedScores.fill(0.0, numValidOutputs);
+  }
+
+  return {
+    selectedIndices: tensor1d(selectedIndices, 'int32'),
+    selectedScores: tensor1d(selectedScores, 'float32'),
+    numValidOutputs: scalar(numValidOutputs, 'int32')
+  };
 }
 
 function intersectionOverUnion(boxes: TypedArray, i: number, j: number) {

tfjs-core/src/backends/webgl/all_kernels.ts

@@ -20,3 +20,4 @@
 // the contents of this file and import only the kernels that are needed.
 import './square';
 import './fromPixels';
+import './non_max_suppression_v5';

tfjs-core/src/backends/webgl/backend_webgl.ts

@@ -45,7 +45,7 @@ import {getArrayFromDType, getTypedArrayFromDType, inferDtype, sizeFromShape} fr
 import {DataStorage, EPSILON_FLOAT16, EPSILON_FLOAT32, KernelBackend} from '../backend';
 import * as backend_util from '../backend_util';
 import {mergeRealAndImagArrays} from '../complex_util';
-import {nonMaxSuppressionImpl} from '../non_max_suppression_impl';
+import {nonMaxSuppressionV3} from '../non_max_suppression_impl';
 import {split} from '../split_shared';
 import {tile} from '../tile_impl';
 import {topkImpl} from '../topk_impl';
@@ -2244,16 +2244,14 @@ export class MathBackendWebGL extends KernelBackend {
 
   nonMaxSuppression(
       boxes: Tensor2D, scores: Tensor1D, maxOutputSize: number,
-      iouThreshold: number, scoreThreshold: number,
-      softNmsSigma: number): Tensor1D {
+      iouThreshold: number, scoreThreshold: number): Tensor1D {
     warn(
         'tf.nonMaxSuppression() in webgl locks the UI thread. ' +
         'Call tf.nonMaxSuppressionAsync() instead');
     const boxesVals = boxes.dataSync();
     const scoresVals = scores.dataSync();
-    return nonMaxSuppressionImpl(
-        boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold,
-        softNmsSigma);
+    return nonMaxSuppressionV3(
+        boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold);
   }
 
   cropAndResize(

tfjs-core/src/backends/webgl/non_max_suppression_v5.ts

@@ -0,0 +1,65 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
+import {warn} from '../../log';
+import {TypedArray} from '../../types';
+import {nonMaxSuppressionV5} from '../non_max_suppression_impl';
+
+import {MathBackendWebGL} from './backend_webgl';
+
+interface NonMaxSuppressionWithScoreInputs extends NamedTensorInfoMap {
+  boxes: TensorInfo;
+  scores: TensorInfo;
+}
+
+interface NonMaxSuppressionWithScoreAttrs extends NamedAttrMap {
+  maxOutputSize: number;
+  iouThreshold: number;
+  scoreThreshold: number;
+  softNmsSigma: number;
+}
+
+registerKernel({
+  kernelName: 'NonMaxSuppressionV5',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, backend, attrs}) => {
+    warn(
+        'tf.nonMaxSuppression() in webgl locks the UI thread. ' +
+        'Call tf.nonMaxSuppressionAsync() instead');
+
+    const {boxes, scores} = inputs as NonMaxSuppressionWithScoreInputs;
+    const {maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma} =
+        attrs as NonMaxSuppressionWithScoreAttrs;
+
+    const gpuBackend = backend as MathBackendWebGL;
+
+    const boxesVals = gpuBackend.readSync(boxes.dataId) as TypedArray;
+    const scoresVals = gpuBackend.readSync(scores.dataId) as TypedArray;
+
+    const maxOutputSizeVal = maxOutputSize;
+    const iouThresholdVal = iouThreshold;
+    const scoreThresholdVal = scoreThreshold;
+    const softNmsSigmaVal = softNmsSigma;
+
+    const {selectedIndices, selectedScores} = nonMaxSuppressionV5(
+        boxesVals, scoresVals, maxOutputSizeVal, iouThresholdVal,
+        scoreThresholdVal, softNmsSigmaVal);
+
+    return [selectedIndices, selectedScores];
+  }
+});