Add NonMaxSuppressionV4.

tfjs-backend-cpu/src/kernels/NonMaxSuppressionV4.ts

@@ -0,0 +1,46 @@
+/**
+ * @license
+ * Copyright 2020 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 {NonMaxSuppressionV4, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs} from '@tensorflow/tfjs-core';
+import {KernelConfig, TypedArray} from '@tensorflow/tfjs-core';
+import {kernel_impls} from '@tensorflow/tfjs-core';
+const nonMaxSuppressionV4Impl = kernel_impls.nonMaxSuppressionV4Impl;
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+export const nonMaxSuppressionV4Config: KernelConfig = {
+  kernelName: NonMaxSuppressionV4,
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend, attrs}) => {
+    const {boxes, scores} = inputs as NonMaxSuppressionV4Inputs;
+    const {maxOutputSize, iouThreshold, scoreThreshold, padToMaxOutputSize} =
+        attrs as unknown as NonMaxSuppressionV4Attrs;
+
+    const cpuBackend = backend as MathBackendCPU;
+
+    assertNotComplex(boxes, 'NonMaxSuppressionPadded');
+
+    const boxesVals = cpuBackend.data.get(boxes.dataId).values as TypedArray;
+    const scoresVals = cpuBackend.data.get(scores.dataId).values as TypedArray;
+
+    const {selectedIndices, validOutputs} = nonMaxSuppressionV4Impl(
+        boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold,
+        padToMaxOutputSize);
+
+    return [selectedIndices, validOutputs];
+  }
+};

tfjs-backend-cpu/src/register_all_kernels.ts

@@ -25,6 +25,7 @@ import {dilation2dBackpropInputConfig} from './kernels/Dilation2DBackpropInput';
 import {divConfig} from './kernels/Div';
 import {maxConfig} from './kernels/Max';
 import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax';
+import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4';
 import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5';
 import {rotateWithOffsetConfig} from './kernels/RotateWithOffset';
 import {squareConfig} from './kernels/Square';
@@ -35,8 +36,8 @@ import {transposeConfig} from './kernels/Transpose';
 const kernelConfigs: KernelConfig[] = [
   dilation2dConfig, dilation2dBackpropInputConfig,
   dilation2dBackpropFilterConfig, divConfig, maxPoolWithArgmaxConfig, maxConfig,
-  nonMaxSuppressionV5Config, rotateWithOffsetConfig, squareConfig,
-  squaredDifferenceConfig, transposeConfig
+  nonMaxSuppressionV4Config, nonMaxSuppressionV5Config, rotateWithOffsetConfig,
+  squareConfig, squaredDifferenceConfig, transposeConfig
 ];
 
 for (const kernelConfig of kernelConfigs) {

tfjs-backend-wasm/src/cc/BUILD

@@ -216,6 +216,7 @@ tfjs_cc_library(
         ":Multiply",
         ":Negate",
         ":NonMaxSuppressionV3",
+        ":NonMaxSuppressionV4",
         ":NonMaxSuppressionV5",
         ":NotEqual",
         ":OneHot",
@@ -655,6 +656,16 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "NonMaxSuppressionV4",
+    srcs = ["kernels/NonMaxSuppressionV4.cc"],
+    deps = [
+        ":backend",
+        ":non_max_suppression_impl",
+        ":util",
+    ],
+)
+
 tfjs_cc_library(
     name = "NonMaxSuppressionV5",
     srcs = ["kernels/NonMaxSuppressionV5.cc"],

tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV3.cc

@@ -35,11 +35,9 @@ EMSCRIPTEN_KEEPALIVE
 const NonMaxSuppressionResult* NonMaxSuppressionV3(
     const size_t boxes_id, const size_t scores_id, const size_t max_out_size,
     const float iou_threshold, const float score_threshold) {
-  const float dummy_soft_nms_sigma = 0.0;
-
-  return tfjs::wasm::non_max_suppression_impl(boxes_id, scores_id, max_out_size,
-                                              iou_threshold, score_threshold,
-                                              dummy_soft_nms_sigma);
+  return tfjs::wasm::non_max_suppression_impl(
+      boxes_id, scores_id, max_out_size, iou_threshold, score_threshold,
+      0.0 /* soft_nms_sigma */, false /* pad_to_max_output_size */);
 }
 
 }  // extern "C"

tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV4.cc

@@ -0,0 +1,48 @@
+/* Copyright 2020 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstring>
+#include <memory>
+#include <queue>
+#include <vector>
+
+#include "src/cc/non_max_suppression_impl.h"
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+const NonMaxSuppressionResult* NonMaxSuppressionV4(
+    const size_t boxes_id, const size_t scores_id, const size_t max_out_size,
+    const float iou_threshold, const float score_threshold,
+    const bool pad_to_max_output_size) {
+  return tfjs::wasm::non_max_suppression_impl(
+      boxes_id, scores_id, max_out_size, iou_threshold, score_threshold,
+      0.0 /* soft_nms_sigma */, pad_to_max_output_size);
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV5.cc

@@ -40,9 +40,9 @@ const NonMaxSuppressionResult* NonMaxSuppressionV5(const size_t boxes_id,
                                                    const float iou_threshold,
                                                    const float score_threshold,
                                                    const float soft_nms_sigma) {
-  return tfjs::wasm::non_max_suppression_impl(boxes_id, scores_id, max_out_size,
-                                              iou_threshold, score_threshold,
-                                              soft_nms_sigma);
+  return tfjs::wasm::non_max_suppression_impl(
+      boxes_id, scores_id, max_out_size, iou_threshold, score_threshold,
+      soft_nms_sigma, false /* pad_to_max_output_size */);
 }
 
 }  // extern "C"

tfjs-backend-wasm/src/cc/non_max_suppression_impl.cc

@@ -81,7 +81,7 @@ namespace wasm {
 const NonMaxSuppressionResult* non_max_suppression_impl(
     const size_t boxes_id, const size_t scores_id, const size_t max_out_size,
     const float iou_threshold, const float score_threshold,
-    const float soft_nms_sigma) {
+    const float soft_nms_sigma, const bool pad_to_max_output_size) {
   auto& boxes_info = backend::get_tensor_info(boxes_id);
   auto& scores_info = backend::get_tensor_info_out(scores_id);
   const float* boxes = boxes_info.f32();
@@ -175,6 +175,12 @@ const NonMaxSuppressionResult* non_max_suppression_impl(
     }
   }
 
+  size_t num_valid_outputs = selected_indices.size();
+  if (pad_to_max_output_size) {
+    selected_indices.resize(max_out_size, 0);
+    selected_scores.resize(max_out_size, 0.0);
+  }
+
   // Allocate memory on the heap for the results and copy the data from the
   // `selected_indices` and `selected_scores` vector since we can't "steal" the
   // data from the vector.
@@ -190,10 +196,16 @@ const NonMaxSuppressionResult* non_max_suppression_impl(
   std::memcpy(selected_scores_data, selected_scores.data(),
               selected_scores_data_size);
 
+  size_t valid_outputs_data_size = sizeof(size_t);
+  size_t* valid_outputs_data =
+      static_cast<size_t*>(malloc(valid_outputs_data_size));
+  *valid_outputs_data = num_valid_outputs;
+
   // Allocate the result of the method on the heap so it survives past this
   // function and we can read it in js.
-  return new NonMaxSuppressionResult{
-      selected_indices_data, selected_indices.size(), selected_scores_data};
+  return new NonMaxSuppressionResult{selected_indices_data,
+                                     selected_indices.size(),
+                                     selected_scores_data, valid_outputs_data};
 }
 
 }  // namespace wasm

tfjs-backend-wasm/src/cc/non_max_suppression_impl.h

@@ -27,12 +27,13 @@ struct NonMaxSuppressionResult {
   int32_t* selected_indices;
   size_t selected_size;
   float* selected_scores;
+  size_t* valid_outputs;
 };
 
 const NonMaxSuppressionResult* non_max_suppression_impl(
     const size_t boxes_id, const size_t scores_id, const size_t max_out_size,
     const float iou_threshold, const float score_threshold,
-    const float soft_nms_sigma);
+    const float soft_nms_sigma, const bool pad_to_max_output_size);
 
 }  // namespace wasm
 }  // namespace tfjs

tfjs-backend-wasm/src/kernels/NonMaxSuppressionV3.ts

@@ -64,11 +64,12 @@ function kernelFunc(args: {
   const resOffset =
       wasmFunc(boxesId, scoresId, maxOutputSize, iouThreshold, scoreThreshold);
 
-  const {pSelectedIndices, selectedSize, pSelectedScores} =
+  const {pSelectedIndices, selectedSize, pSelectedScores, pValidOutputs} =
       parseResultStruct(backend, resOffset);
 
   // Since we are not using scores for V3, we have to delete it from the heap.
   backend.wasm._free(pSelectedScores);
+  backend.wasm._free(pValidOutputs);
 
   const selectedIndicesTensor =
       backend.makeOutput([selectedSize], 'int32', pSelectedIndices);

tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts

@@ -0,0 +1,79 @@
+/**
+ * @license
+ * Copyright 2020 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 {KernelConfig, KernelFunc, NonMaxSuppressionV4, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+import {parseResultStruct} from './NonMaxSuppression_util';
+
+let wasmFunc: (
+    boxesId: number, scoresId: number, maxOutputSize: number,
+    iouThreshold: number, scoreThreshold: number,
+    padToMaxOutputSize: boolean) => number;
+
+function setup(backend: BackendWasm): void {
+  wasmFunc = backend.wasm.cwrap(
+      NonMaxSuppressionV4,
+      'number',  // Result*
+      [
+        'number',  // boxesId
+        'number',  // scoresId
+        'number',  // maxOutputSize
+        'number',  // iouThreshold
+        'number',  // scoreThreshold
+        'bool',    // padToMaxOutputSize
+      ]);
+}
+
+function nonMaxSuppressionV4(args: {
+  backend: BackendWasm,
+  inputs: NonMaxSuppressionV4Inputs,
+  attrs: NonMaxSuppressionV4Attrs
+}): TensorInfo[] {
+  const {backend, inputs, attrs} = args;
+  const {iouThreshold, maxOutputSize, scoreThreshold, padToMaxOutputSize} =
+      attrs;
+  const {boxes, scores} = inputs;
+
+  const boxesId = backend.dataIdMap.get(boxes.dataId).id;
+  const scoresId = backend.dataIdMap.get(scores.dataId).id;
+
+  const resOffset = wasmFunc(
+      boxesId, scoresId, maxOutputSize, iouThreshold, scoreThreshold,
+      padToMaxOutputSize);
+
+  const {pSelectedIndices, selectedSize, pSelectedScores, pValidOutputs} =
+      parseResultStruct(backend, resOffset);
+
+  // Since we are not using scores for V4, we have to delete it from the heap.
+  backend.wasm._free(pSelectedScores);
+
+  const selectedIndicesTensor =
+      backend.makeOutput([selectedSize], 'int32', pSelectedIndices);
+
+  const validOutputsTensor = backend.makeOutput([], 'int32', pValidOutputs);
+
+  return [selectedIndicesTensor, validOutputsTensor];
+}
+
+export const nonMaxSuppressionV4Config: KernelConfig = {
+  kernelName: NonMaxSuppressionV4,
+  backendName: 'wasm',
+  setupFunc: setup,
+  kernelFunc: nonMaxSuppressionV4 as {} as KernelFunc,
+};

tfjs-backend-wasm/src/kernels/NonMaxSuppressionV5.ts

@@ -68,11 +68,12 @@ function kernelFunc(args: {
       boxesId, scoresId, maxOutputSize, iouThreshold, scoreThreshold,
       softNmsSigma);
 
-  const {
-    pSelectedIndices,
-    selectedSize,
-    pSelectedScores,
-  } = parseResultStruct(backend, resOffset);
+  const {pSelectedIndices, selectedSize, pSelectedScores, pValidOutputs} =
+      parseResultStruct(backend, resOffset);
+
+  // Since we are not using validOutputs for V5, we have to delete it from the
+  // heap.
+  backend.wasm._free(pValidOutputs);
 
   const selectedIndicesTensor =
       backend.makeOutput([selectedSize], 'int32', pSelectedIndices);

tfjs-backend-wasm/src/kernels/NonMaxSuppression_util.ts

@@ -22,18 +22,20 @@ interface Result {
   pSelectedIndices: number;
   selectedSize: number;
   pSelectedScores: number;
+  pValidOutputs: number;
 }
 /**
  * Parse the result of the c++ method, which has the shape equivalent to
  * `Result`.
  */
 export function parseResultStruct(
     backend: BackendWasm, resOffset: number): Result {
-  const result = new Int32Array(backend.wasm.HEAPU8.buffer, resOffset, 3);
+  const result = new Int32Array(backend.wasm.HEAPU8.buffer, resOffset, 4);
   const pSelectedIndices = result[0];
   const selectedSize = result[1];
   const pSelectedScores = result[2];
+  const pValidOutputs = result[3];
   // Since the result was allocated on the heap, we have to delete it.
   backend.wasm._free(resOffset);
-  return {pSelectedIndices, selectedSize, pSelectedScores};
+  return {pSelectedIndices, selectedSize, pSelectedScores, pValidOutputs};
 }

tfjs-backend-wasm/src/register_all_kernels.ts

@@ -59,6 +59,7 @@ import {minimumConfig} from './kernels/Minimum';
 import {multiplyConfig} from './kernels/Multiply';
 import {negateConfig} from './kernels/Negate';
 import {nonMaxSuppressionV3Config} from './kernels/NonMaxSuppressionV3';
+import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4';
 import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5';
 import {notEqualConfig} from './kernels/NotEqual';
 import {oneHotConfig} from './kernels/OneHot';
@@ -132,6 +133,7 @@ const kernelConfigs: KernelConfig[] = [
   multiplyConfig,
   negateConfig,
   nonMaxSuppressionV3Config,
+  nonMaxSuppressionV4Config,
   nonMaxSuppressionV5Config,
   notEqualConfig,
   oneHotConfig,