From c5814a98a770b15c04f0987ae2f8427bf9af9c86 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 13 Jul 2020 22:52:49 -0700 Subject: [PATCH 1/7] Resolve merge conflict. --- .../src/kernels/NonMaxSuppressionV4.ts | 46 ++++++++++ tfjs-backend-cpu/src/register_all_kernels.ts | 5 +- .../src/cc/kernels/NonMaxSuppressionV3.cc | 8 +- .../src/cc/kernels/NonMaxSuppressionV4.cc | 48 ++++++++++ .../src/cc/kernels/NonMaxSuppressionV5.cc | 6 +- .../src/cc/non_max_suppression_impl.cc | 18 +++- .../src/cc/non_max_suppression_impl.h | 3 +- .../src/kernels/NonMaxSuppressionV3.ts | 3 +- .../src/kernels/NonMaxSuppressionV4.ts | 79 +++++++++++++++++ .../src/kernels/NonMaxSuppressionV5.ts | 11 +-- .../src/kernels/NonMaxSuppression_util.ts | 6 +- tfjs-backend-wasm/src/register_all_kernels.ts | 2 + .../src/kernels/NonMaxSuppressionV4.ts | 45 ++++++++++ .../src/register_all_kernels.ts | 6 +- tfjs-core/src/backends/kernel_impls.ts | 2 +- .../src/backends/non_max_suppression_impl.ts | 76 +++++++++------- tfjs-core/src/index.ts | 3 +- tfjs-core/src/kernel_names.ts | 10 +++ .../src/ops/non_max_suppression_padded.ts | 87 +++++++++++++++++++ .../ops/non_max_suppression_padded_test.ts | 60 +++++++++++++ tfjs-core/src/ops/ops.ts | 4 +- tfjs-core/src/tests.ts | 1 + tfjs-node/binding/utils.h | 2 +- tfjs-node/src/kernels/NonMaxSuppressionV4.ts | 59 +++++++++++++ tfjs-node/src/nodejs_kernel_backend.ts | 11 ++- tfjs-node/src/register_all_kernels.ts | 5 +- 26 files changed, 542 insertions(+), 64 deletions(-) create mode 100644 tfjs-backend-cpu/src/kernels/NonMaxSuppressionV4.ts create mode 100644 tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV4.cc create mode 100644 tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts create mode 100644 tfjs-backend-webgl/src/kernels/NonMaxSuppressionV4.ts create mode 100644 tfjs-core/src/ops/non_max_suppression_padded.ts create mode 100644 tfjs-core/src/ops/non_max_suppression_padded_test.ts create mode 100644 tfjs-node/src/kernels/NonMaxSuppressionV4.ts diff --git a/tfjs-backend-cpu/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-cpu/src/kernels/NonMaxSuppressionV4.ts new file mode 100644 index 00000000000..55fb52f9a89 --- /dev/null +++ b/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]; + } +}; diff --git a/tfjs-backend-cpu/src/register_all_kernels.ts b/tfjs-backend-cpu/src/register_all_kernels.ts index 0d51e27d49a..82e4d076651 100644 --- a/tfjs-backend-cpu/src/register_all_kernels.ts +++ b/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) { diff --git a/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV3.cc b/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV3.cc index e975b6bfc77..98fb90ad11c 100644 --- a/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV3.cc +++ b/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" diff --git a/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV4.cc b/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV4.cc new file mode 100644 index 00000000000..cf2ec5256fa --- /dev/null +++ b/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 +#endif + +#include +#include +#include +#include +#include +#include +#include + +#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 diff --git a/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV5.cc b/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV5.cc index 3e23fcd72a4..bed63cc36e6 100644 --- a/tfjs-backend-wasm/src/cc/kernels/NonMaxSuppressionV5.cc +++ b/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" diff --git a/tfjs-backend-wasm/src/cc/non_max_suppression_impl.cc b/tfjs-backend-wasm/src/cc/non_max_suppression_impl.cc index dd3f3b8383a..1d976fa3bc8 100644 --- a/tfjs-backend-wasm/src/cc/non_max_suppression_impl.cc +++ b/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(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 diff --git a/tfjs-backend-wasm/src/cc/non_max_suppression_impl.h b/tfjs-backend-wasm/src/cc/non_max_suppression_impl.h index bb3168599a7..71cc2190ef0 100644 --- a/tfjs-backend-wasm/src/cc/non_max_suppression_impl.h +++ b/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 diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV3.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV3.ts index ee6a1beb0c6..70a423801d0 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV3.ts +++ b/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); diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts new file mode 100644 index 00000000000..e3b149f4410 --- /dev/null +++ b/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, 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 kernelFunc(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, +}; diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV5.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV5.ts index 157b5a44078..b92b803be7b 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV5.ts +++ b/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); diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppression_util.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppression_util.ts index c7088d92a0c..2b2fc6debcc 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppression_util.ts +++ b/tfjs-backend-wasm/src/kernels/NonMaxSuppression_util.ts @@ -22,6 +22,7 @@ interface Result { pSelectedIndices: number; selectedSize: number; pSelectedScores: number; + pValidOutputs: number; } /** * Parse the result of the c++ method, which has the shape equivalent to @@ -29,11 +30,12 @@ interface 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}; } diff --git a/tfjs-backend-wasm/src/register_all_kernels.ts b/tfjs-backend-wasm/src/register_all_kernels.ts index a3055f20459..9675466ba42 100644 --- a/tfjs-backend-wasm/src/register_all_kernels.ts +++ b/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, diff --git a/tfjs-backend-webgl/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-webgl/src/kernels/NonMaxSuppressionV4.ts new file mode 100644 index 00000000000..df604b4ffda --- /dev/null +++ b/tfjs-backend-webgl/src/kernels/NonMaxSuppressionV4.ts @@ -0,0 +1,45 @@ +/** + * @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 {backend_util, kernel_impls, KernelConfig, NonMaxSuppressionV4, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, TypedArray} from '@tensorflow/tfjs-core'; +const nonMaxSuppressionV4Impl = kernel_impls.nonMaxSuppressionV4Impl; +import {MathBackendWebGL} from '../backend_webgl'; + +export const nonMaxSuppressionV4Config: KernelConfig = { + kernelName: NonMaxSuppressionV4, + backendName: 'webgl', + kernelFunc: ({inputs, backend, attrs}) => { + backend_util.warn( + 'tf.nonMaxSuppression() in webgl locks the UI thread. ' + + 'Call tf.nonMaxSuppressionAsync() instead'); + + const {boxes, scores} = inputs as NonMaxSuppressionV4Inputs; + const {maxOutputSize, iouThreshold, scoreThreshold, padToMaxOutputSize} = + attrs as unknown as NonMaxSuppressionV4Attrs; + + const gpuBackend = backend as MathBackendWebGL; + + const boxesVals = gpuBackend.readSync(boxes.dataId) as TypedArray; + const scoresVals = gpuBackend.readSync(scores.dataId) as TypedArray; + + const {selectedIndices, validOutputs} = nonMaxSuppressionV4Impl( + boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold, + padToMaxOutputSize); + + return [selectedIndices, validOutputs]; + } +}; diff --git a/tfjs-backend-webgl/src/register_all_kernels.ts b/tfjs-backend-webgl/src/register_all_kernels.ts index b55ba1ef6f1..18ad9aed90c 100644 --- a/tfjs-backend-webgl/src/register_all_kernels.ts +++ b/tfjs-backend-webgl/src/register_all_kernels.ts @@ -21,6 +21,7 @@ import {fromPixelsConfig} from './kernels/FromPixels'; import {maxConfig} from './kernels/Max'; import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax'; import {nonMaxSuppressionV3Config} from './kernels/NonMaxSuppressionV3'; +import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4'; import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5'; import {rotateWithOffsetConfig} from './kernels/RotateWithOffset'; import {squareConfig} from './kernels/Square'; @@ -30,8 +31,9 @@ import {transposeConfig} from './kernels/Transpose'; // List all kernel configs here const kernelConfigs: KernelConfig[] = [ maxConfig, fromPixelsConfig, divConfig, maxPoolWithArgmaxConfig, - nonMaxSuppressionV3Config, nonMaxSuppressionV5Config, rotateWithOffsetConfig, - squareConfig, squaredDifferenceConfig, transposeConfig + nonMaxSuppressionV3Config, nonMaxSuppressionV4Config, + nonMaxSuppressionV5Config, rotateWithOffsetConfig, squareConfig, + squaredDifferenceConfig, transposeConfig ]; for (const kernelConfig of kernelConfigs) { diff --git a/tfjs-core/src/backends/kernel_impls.ts b/tfjs-core/src/backends/kernel_impls.ts index 2b28dd4b65d..093f1f7bd0d 100644 --- a/tfjs-core/src/backends/kernel_impls.ts +++ b/tfjs-core/src/backends/kernel_impls.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -export {nonMaxSuppressionV3Impl, nonMaxSuppressionV5Impl} from './non_max_suppression_impl'; +export {nonMaxSuppressionV3Impl, nonMaxSuppressionV4Impl, nonMaxSuppressionV5Impl} from './non_max_suppression_impl'; export {split} from './split_shared'; export {tile} from './tile_impl'; export {topkImpl} from './topk_impl'; diff --git a/tfjs-core/src/backends/non_max_suppression_impl.ts b/tfjs-core/src/backends/non_max_suppression_impl.ts index 9acea57483b..3ccf8429c1e 100644 --- a/tfjs-core/src/backends/non_max_suppression_impl.ts +++ b/tfjs-core/src/backends/non_max_suppression_impl.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2018 Google LLC. All Rights Reserved. + * 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 @@ -36,40 +36,43 @@ interface Candidate { export function nonMaxSuppressionV3Impl( boxes: TypedArray, scores: TypedArray, maxOutputSize: number, iouThreshold: number, scoreThreshold: number): Tensor1D { - const dummySoftNmsSigma = 0.0; - - const result = nonMaxSuppressionImpl_( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, - dummySoftNmsSigma); - result.selectedScores.dispose(); - result.numValidOutputs.dispose(); - return result.selectedIndices as Tensor1D; + return nonMaxSuppressionImpl_( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + 0 /* softNmsSigma */) + .selectedIndices as Tensor1D; +} + +export function nonMaxSuppressionV4Impl( + boxes: TypedArray, scores: TypedArray, maxOutputSize: number, + iouThreshold: number, scoreThreshold: number, + padToMaxOutputSize: boolean): NamedTensorMap { + return nonMaxSuppressionImpl_( + boxes, + scores, + maxOutputSize, + iouThreshold, + scoreThreshold, + 0, /* softNmsSigma */ + false, /* returnScoresTensor */ + padToMaxOutputSize, /* padToMaxOutputSize */ + true, + /* returnValidOutputs */); } export function nonMaxSuppressionV5Impl( boxes: TypedArray, scores: TypedArray, maxOutputSize: number, iouThreshold: number, scoreThreshold: number, softNmsSigma: number): NamedTensorMap { - // For NonMaxSuppressionV5Op, we always return a second output holding - // corresponding scores. - const returnScoresTensor = true; - - const result = nonMaxSuppressionImpl_( + return nonMaxSuppressionImpl_( boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma, - returnScoresTensor); - - result.numValidOutputs.dispose(); - - return { - selectedIndices: result.selectedIndices, - selectedScores: result.selectedScores - }; + true /* returnScoresTensor */); } function nonMaxSuppressionImpl_( boxes: TypedArray, scores: TypedArray, maxOutputSize: number, iouThreshold: number, scoreThreshold: number, softNmsSigma: number, - returnScoresTensor = false, padToMaxOutputSize = false): NamedTensorMap { + returnScoresTensor = false, padToMaxOutputSize = false, + returnValidOutputs = 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 = []; @@ -144,17 +147,26 @@ function nonMaxSuppressionImpl_( } // NonMaxSuppressionV4 feature: padding output to maxOutputSize. - const numValidOutputs = selectedIndices.length; - if (padToMaxOutputSize) { - selectedIndices.fill(0, numValidOutputs); - selectedScores.fill(0.0, numValidOutputs); + const validOutputs = selectedIndices.length; + const elemsToPad = maxOutputSize - validOutputs; + + if (padToMaxOutputSize && elemsToPad > 0) { + selectedIndices.push(...new Array(elemsToPad).fill(0)); + selectedScores.push(...new Array(elemsToPad).fill(0.0)); + } + + const result: + NamedTensorMap = {selectedIndices: tensor1d(selectedIndices, 'int32')}; + + if (returnScoresTensor) { + result['selectedScores'] = tensor1d(selectedScores, 'float32'); + } + + if (returnValidOutputs) { + result['validOutputs'] = scalar(validOutputs, 'int32'); } - return { - selectedIndices: tensor1d(selectedIndices, 'int32'), - selectedScores: tensor1d(selectedScores, 'float32'), - numValidOutputs: scalar(numValidOutputs, 'int32') - }; + return result; } function intersectionOverUnion(boxes: TypedArray, i: number, j: number) { diff --git a/tfjs-core/src/index.ts b/tfjs-core/src/index.ts index 2dd95b94200..49a95123436 100644 --- a/tfjs-core/src/index.ts +++ b/tfjs-core/src/index.ts @@ -29,7 +29,6 @@ import './engine'; // Register backend-agnostic flags. import './flags'; - // Register all the gradients. import './register_all_gradients'; import './platforms/platform_browser'; @@ -61,7 +60,7 @@ export {RMSPropOptimizer} from './optimizers/rmsprop_optimizer'; export {SGDOptimizer} from './optimizers/sgd_optimizer'; export {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, Variable} from './tensor'; export {GradSaveFunc, NamedTensorMap, TensorContainer, TensorContainerArray, TensorContainerObject} from './tensor_types'; -export {BackendValues, DataType, DataTypeMap, DataValues, NumericDataType, PixelData, Rank, RecursiveArray, ShapeMap, sumOutType, TensorLike, TypedArray, upcastType} from './types'; +export {BackendValues, DataType, DataTypeMap, DataValues, NumericDataType, PixelData, Rank, RecursiveArray, ScalarLike, ShapeMap, sumOutType, TensorLike, TypedArray, upcastType} from './types'; export * from './ops/ops'; export {Reduction} from './ops/loss_ops_utils'; diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts index b81086e8443..6cd52eaa22c 100644 --- a/tfjs-core/src/kernel_names.ts +++ b/tfjs-core/src/kernel_names.ts @@ -513,6 +513,16 @@ export interface NonMaxSuppressionV3Attrs { scoreThreshold: number; } +export const NonMaxSuppressionV4 = 'NonMaxSuppressionV4'; +export type NonMaxSuppressionV4Inputs = + Pick; +export interface NonMaxSuppressionV4Attrs { + maxOutputSize: number; + iouThreshold: number; + scoreThreshold: number; + padToMaxOutputSize: boolean; +} + export const NonMaxSuppressionV5 = 'NonMaxSuppressionV5'; export type NonMaxSuppressionV5Inputs = Pick; diff --git a/tfjs-core/src/ops/non_max_suppression_padded.ts b/tfjs-core/src/ops/non_max_suppression_padded.ts new file mode 100644 index 00000000000..59a48e60039 --- /dev/null +++ b/tfjs-core/src/ops/non_max_suppression_padded.ts @@ -0,0 +1,87 @@ +/** + * @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 {ENGINE} from '../engine'; +import {NonMaxSuppressionV4, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs} from '../kernel_names'; +import {NamedAttrMap} from '../kernel_registry'; +import {Tensor, Tensor1D, Tensor2D} from '../tensor'; +import {NamedTensorMap} from '../tensor_types'; +import {convertToTensor} from '../tensor_util_env'; +import {TensorLike} from '../types'; + +import {nonMaxSuppSanityCheck} from './nonmax_util'; +import {op} from './operation'; + +/** + * Performs non maximum suppression of bounding boxes based on + * iou (intersection over union). + * + * This op also supports a Soft-NMS mode (c.f. + * Bodla et al, https://arxiv.org/abs/1704.04503) where boxes reduce the score + * of other overlapping boxes, therefore favoring different regions of the image + * with high scores. To enable this Soft-NMS mode, set the `softNmsSigma` + * parameter to be larger than 0. + * + * @param boxes a 2d tensor of shape `[numBoxes, 4]`. Each entry is + * `[y1, x1, y2, x2]`, where `(y1, x1)` and `(y2, x2)` are the corners of + * the bounding box. + * @param scores a 1d tensor providing the box scores of shape `[numBoxes]`. + * @param maxOutputSize The maximum number of boxes to be selected. + * @param iouThreshold A float representing the threshold for deciding whether + * boxes overlap too much with respect to IOU. Must be between [0, 1]. + * Defaults to 0.5 (50% box overlap). + * @param scoreThreshold A threshold for deciding when to remove boxes based + * on score. Defaults to -inf, which means any score is accepted. + * @param padToMaxOutputSize Defalts to false. If true, size of output + * `selectedIndices` is padded to maxOutputSize. + * @return A map with the following properties: + * - selectedIndices: A 1D tensor with the selected box indices. + * - validOutputs: A scalar denoting how many elements in `selectedIndices` + * are valid. Valid elements occur first, then padding. + */ +/** @doc {heading: 'Operations', subheading: 'Images', namespace: 'image'} */ +function nonMaxSuppressionPadded_( + boxes: Tensor2D|TensorLike, scores: Tensor1D|TensorLike, + maxOutputSize: number, iouThreshold = 0.5, + scoreThreshold = Number.NEGATIVE_INFINITY, + padToMaxOutputSize = false): NamedTensorMap { + const $boxes = convertToTensor(boxes, 'boxes', 'nonMaxSuppression'); + const $scores = convertToTensor(scores, 'scores', 'nonMaxSuppression'); + + const params = nonMaxSuppSanityCheck( + $boxes, $scores, maxOutputSize, iouThreshold, scoreThreshold, + null /* softNmsSigma */); + const $maxOutputSize = params.maxOutputSize; + const $iouThreshold = params.iouThreshold; + const $scoreThreshold = params.scoreThreshold; + + const inputs: NonMaxSuppressionV4Inputs = {boxes: $boxes, scores: $scores}; + const attrs: NonMaxSuppressionV4Attrs = { + maxOutputSize: $maxOutputSize, + iouThreshold: $iouThreshold, + scoreThreshold: $scoreThreshold, + padToMaxOutputSize + }; + + const result = ENGINE.runKernel( + NonMaxSuppressionV4, inputs as {} as NamedTensorMap, + attrs as {} as NamedAttrMap) as Tensor[]; + + return {selectedIndices: result[0], validOutputs: result[1]}; +} + +export const nonMaxSuppressionPadded = op({nonMaxSuppressionPadded_}); diff --git a/tfjs-core/src/ops/non_max_suppression_padded_test.ts b/tfjs-core/src/ops/non_max_suppression_padded_test.ts new file mode 100644 index 00000000000..c174286d343 --- /dev/null +++ b/tfjs-core/src/ops/non_max_suppression_padded_test.ts @@ -0,0 +1,60 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysEqual} from '../test_util'; + +describeWithFlags('nonMaxSuppressionPadded', ALL_ENVS, () => { + it('select from three clusters with pad five.', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 5; + const iouThreshold = 0.5; + const scoreThreshold = 0.0; + + const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); + + expectArraysEqual(await selectedIndices.data(), [3, 0, 5, 0, 0]); + expectArraysEqual(await validOutputs.data(), 3); + }); + + it('select from three clusters with pad five and score threshold.', + async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 6; + const iouThreshold = 0.5; + const scoreThreshold = 0.4; + + const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); + + expectArraysEqual(await selectedIndices.data(), [3, 0, 0, 0, 0, 0]); + expectArraysEqual(await validOutputs.data(), 2); + }); +}); diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts index ffb252b2810..d02b2f2fc89 100644 --- a/tfjs-core/src/ops/ops.ts +++ b/tfjs-core/src/ops/ops.ts @@ -239,6 +239,7 @@ import {nonMaxSuppression} from './non_max_suppression'; import {nonMaxSuppressionAsync} from './non_max_suppression_async'; import {nonMaxSuppressionWithScore} from './non_max_suppression_with_score'; import {nonMaxSuppressionWithScoreAsync} from './non_max_suppresion_with_score_async'; +import {nonMaxSuppressionPadded} from './non_max_suppression_padded'; import {resizeBilinear} from './resize_bilinear'; import {resizeNearestNeighbor} from './resize_nearest_neighbor'; const image = { @@ -249,7 +250,8 @@ const image = { nonMaxSuppression, nonMaxSuppressionAsync, nonMaxSuppressionWithScore, - nonMaxSuppressionWithScoreAsync + nonMaxSuppressionWithScoreAsync, + nonMaxSuppressionPadded }; // linalg namespace diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index 0e08d5965d2..9bb3a4d3f8a 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -148,6 +148,7 @@ import './ops/multi_rnn_cell_test'; import './ops/multinomial_test'; import './ops/neg_test'; import './ops/non_max_suppression_async_test'; +import './ops/non_max_suppression_padded_test'; import './ops/non_max_suppression_test'; import './ops/norm_test'; import './ops/not_equal_test'; diff --git a/tfjs-node/binding/utils.h b/tfjs-node/binding/utils.h index a6f547f0e99..b620845a9e5 100644 --- a/tfjs-node/binding/utils.h +++ b/tfjs-node/binding/utils.h @@ -150,7 +150,7 @@ inline bool EnsureValueIsNumber(napi_env env, napi_value value, ENSURE_NAPI_OK_RETVAL(env, napi_typeof(env, value, &type), false); bool is_number = type == napi_number; if (!is_number) { - NapiThrowError(env, file, line_number, "Argument is not a string!"); + NapiThrowError(env, file, line_number, "Argument is not a number!"); } return is_number; } diff --git a/tfjs-node/src/kernels/NonMaxSuppressionV4.ts b/tfjs-node/src/kernels/NonMaxSuppressionV4.ts new file mode 100644 index 00000000000..e98e15f1c11 --- /dev/null +++ b/tfjs-node/src/kernels/NonMaxSuppressionV4.ts @@ -0,0 +1,59 @@ +/** + * @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 {KernelConfig, NonMaxSuppressionV4, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, scalar, Tensor1D, Tensor2D} from '@tensorflow/tfjs'; + +import {createTensorsTypeOpAttr, NodeJSKernelBackend} from '../nodejs_kernel_backend'; + +// TODO(nsthorat, dsmilkov): Remove dependency on tensors, use dataId. +export const nonMaxSuppressionV4Config: KernelConfig = { + kernelName: NonMaxSuppressionV4, + backendName: 'tensorflow', + kernelFunc: ({inputs, backend, attrs}) => { + const {boxes, scores} = inputs as NonMaxSuppressionV4Inputs; + const {maxOutputSize, iouThreshold, scoreThreshold, padToMaxOutputSize} = + attrs as {} as NonMaxSuppressionV4Attrs; + const maxOutputSizeTensor = scalar(maxOutputSize, 'int32'); + const iouThresholdTensor = scalar(iouThreshold, 'float32'); + const scoreThresholdTensor = scalar(scoreThreshold, 'float32'); + + const nodeBackend = backend as NodeJSKernelBackend; + + const opAttrs = [ + createTensorsTypeOpAttr('T', boxes.dtype), + createTensorsTypeOpAttr('T_threshold', 'float32'), { + name: 'pad_to_max_output_size', + type: nodeBackend.binding.TF_ATTR_BOOL, + value: padToMaxOutputSize + } + ]; + + const [selectedIndices, validOutputs] = nodeBackend.executeMultipleOutputs( + 'NonMaxSuppressionV4', opAttrs, + [ + boxes as Tensor2D, scores as Tensor1D, maxOutputSizeTensor, + iouThresholdTensor, scoreThresholdTensor + ], + 2); + + maxOutputSizeTensor.dispose(); + iouThresholdTensor.dispose(); + scoreThresholdTensor.dispose(); + + return [selectedIndices, validOutputs]; + } +}; diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 231ece96df1..ba0f64331ff 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -16,7 +16,7 @@ */ import * as tf from '@tensorflow/tfjs'; -import {backend_util, BackendTimingInfo, DataId, DataType, fill, KernelBackend, ones, Rank, rsqrt, Scalar, scalar, ShapeMap, Tensor, Tensor1D, tensor1d, Tensor2D, tensor2d, Tensor3D, Tensor4D, Tensor5D, TensorInfo, tidy, util} from '@tensorflow/tfjs'; +import {backend_util, BackendTimingInfo, DataId, DataType, fill, KernelBackend, ones, Rank, rsqrt, Scalar, scalar, ScalarLike, ShapeMap, Tensor, Tensor1D, tensor1d, Tensor2D, tensor2d, Tensor3D, Tensor4D, Tensor5D, TensorInfo, tidy, util} from '@tensorflow/tfjs'; import {isArray, isNullOrUndefined} from 'util'; import {Int64Scalar} from './int64_tensors'; @@ -2053,6 +2053,15 @@ export function createTensorsTypeOpAttr( }; } +export function createOpAttr( + attrName: string, tensorsOrDtype: tf.Tensor|tf.Tensor[]|tf.DataType, + value: ScalarLike): TFEOpAttr { + if (isNullOrUndefined(tensorsOrDtype)) { + throw new Error('Invalid input tensors value.'); + } + return {name: attrName, type: nodeBackend().binding.TF_BOOL, value}; +} + /** Returns the dtype number for a single or list of input Tensors. */ function getTFDTypeForInputs(tensors: tf.Tensor|tf.Tensor[]): number { if (isNullOrUndefined(tensors)) { diff --git a/tfjs-node/src/register_all_kernels.ts b/tfjs-node/src/register_all_kernels.ts index f6d7727d29a..83312479b44 100644 --- a/tfjs-node/src/register_all_kernels.ts +++ b/tfjs-node/src/register_all_kernels.ts @@ -22,6 +22,7 @@ import {KernelConfig, registerKernel} from '@tensorflow/tfjs-core'; import {dilation2dConfig} from './kernels/Dilation2D'; import {dilation2dBackpropFilterConfig} from './kernels/Dilation2DBackpropFilter'; import {dilation2dBackpropInputConfig} from './kernels/Dilation2DBackpropInput'; +import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4'; import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5'; import {softmaxConfig} from './kernels/Softmax'; import {squaredDifferenceConfig} from './kernels/SquaredDifference'; @@ -29,8 +30,8 @@ import {squaredDifferenceConfig} from './kernels/SquaredDifference'; // List all kernel configs here const kernelConfigs: KernelConfig[] = [ dilation2dConfig, dilation2dBackpropInputConfig, - dilation2dBackpropFilterConfig, nonMaxSuppressionV5Config, softmaxConfig, - squaredDifferenceConfig + dilation2dBackpropFilterConfig, nonMaxSuppressionV4Config, + nonMaxSuppressionV5Config, softmaxConfig, squaredDifferenceConfig ]; for (const kernelConfig of kernelConfigs) { From 78b064ce66300c3e163d5c6b7f599a737d5adf74 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 14 Jul 2020 10:09:01 -0700 Subject: [PATCH 2/7] Add async version. --- .../src/kernels/NonMaxSuppressionV4.ts | 4 +- .../src/ops/non_max_suppression_async_test.ts | 84 ++-- .../src/ops/non_max_suppression_padded.ts | 10 +- .../ops/non_max_suppression_padded_async.ts | 80 ++++ .../non_max_suppression_padded_async_test.ts | 49 +++ .../ops/non_max_suppression_padded_test.ts | 20 + tfjs-core/src/ops/non_max_suppression_test.ts | 380 ++++++++---------- ...> non_max_suppression_with_score_async.ts} | 0 ...n_max_suppression_with_score_async_test.ts | 53 +++ .../non_max_suppression_with_score_test.ts | 45 +++ tfjs-core/src/ops/ops.ts | 6 +- tfjs-core/src/tests.ts | 3 + 12 files changed, 458 insertions(+), 276 deletions(-) create mode 100644 tfjs-core/src/ops/non_max_suppression_padded_async.ts create mode 100644 tfjs-core/src/ops/non_max_suppression_padded_async_test.ts rename tfjs-core/src/ops/{non_max_suppresion_with_score_async.ts => non_max_suppression_with_score_async.ts} (100%) create mode 100644 tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts create mode 100644 tfjs-core/src/ops/non_max_suppression_with_score_test.ts diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts index e3b149f4410..632bfe0a9f2 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts +++ b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -import {KernelConfig, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, TensorInfo} from '@tensorflow/tfjs-core'; +import {KernelConfig, KernelFunc, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, TensorInfo} from '@tensorflow/tfjs-core'; import {BackendWasm} from '../backend_wasm'; @@ -75,5 +75,5 @@ export const nonMaxSuppressionV4Config: KernelConfig = { kernelName: 'NonMaxSuppressionV4', backendName: 'wasm', setupFunc: setup, - kernelFunc, + kernelFunc as {} as KernelFunc, }; diff --git a/tfjs-core/src/ops/non_max_suppression_async_test.ts b/tfjs-core/src/ops/non_max_suppression_async_test.ts index c5e453b4fe8..0ea530ac022 100644 --- a/tfjs-core/src/ops/non_max_suppression_async_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_async_test.ts @@ -16,69 +16,33 @@ */ import * as tf from '../index'; import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose, expectArraysEqual} from '../test_util'; +import {expectArraysEqual} from '../test_util'; describeWithFlags('nonMaxSuppressionAsync', ALL_ENVS, () => { - describe('NonMaxSuppressionAsync basic', () => { - it('select from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = await tf.image.nonMaxSuppressionAsync( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - }); - - it('accepts a tensor-like object', async () => { - const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; - const scores = [1, 2]; - const indices = await tf.image.nonMaxSuppressionAsync(boxes, scores, 10); - expect(indices.shape).toEqual([2]); - expect(indices.dtype).toEqual('int32'); - expectArraysEqual(await indices.data(), [1, 0]); - }); + it('select from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = await tf.image.nonMaxSuppressionAsync( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); }); - describe('NonMaxSuppressionWithScoreAsync', () => { - it('select from three clusters with SoftNMS', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 6; - const iouThreshold = 1.0; - const scoreThreshold = 0; - const softNmsSigma = 0.5; - - const numTensorsBefore = tf.memory().numTensors; - - const {selectedIndices, selectedScores} = - await tf.image.nonMaxSuppressionWithScoreAsync( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, - softNmsSigma); - - const numTensorsAfter = tf.memory().numTensors; - - expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); - - expectArraysClose( - await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); - - // The number of tensors should increase by the number of tensors - // returned (i.e. selectedIndices and selectedScores). - expect(numTensorsAfter).toEqual(numTensorsBefore + 2); - }); + it('accepts a tensor-like object', async () => { + const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; + const scores = [1, 2]; + const indices = await tf.image.nonMaxSuppressionAsync(boxes, scores, 10); + expect(indices.shape).toEqual([2]); + expect(indices.dtype).toEqual('int32'); + expectArraysEqual(await indices.data(), [1, 0]); }); }); diff --git a/tfjs-core/src/ops/non_max_suppression_padded.ts b/tfjs-core/src/ops/non_max_suppression_padded.ts index 59a48e60039..dbe69d6256a 100644 --- a/tfjs-core/src/ops/non_max_suppression_padded.ts +++ b/tfjs-core/src/ops/non_max_suppression_padded.ts @@ -27,14 +27,8 @@ import {nonMaxSuppSanityCheck} from './nonmax_util'; import {op} from './operation'; /** - * Performs non maximum suppression of bounding boxes based on - * iou (intersection over union). - * - * This op also supports a Soft-NMS mode (c.f. - * Bodla et al, https://arxiv.org/abs/1704.04503) where boxes reduce the score - * of other overlapping boxes, therefore favoring different regions of the image - * with high scores. To enable this Soft-NMS mode, set the `softNmsSigma` - * parameter to be larger than 0. + * Asynchronously performs non maximum suppression of bounding boxes based on + * iou (intersection over union), with an option to pad results. * * @param boxes a 2d tensor of shape `[numBoxes, 4]`. Each entry is * `[y1, x1, y2, x2]`, where `(y1, x1)` and `(y2, x2)` are the corners of diff --git a/tfjs-core/src/ops/non_max_suppression_padded_async.ts b/tfjs-core/src/ops/non_max_suppression_padded_async.ts new file mode 100644 index 00000000000..9dc8149e3f8 --- /dev/null +++ b/tfjs-core/src/ops/non_max_suppression_padded_async.ts @@ -0,0 +1,80 @@ +/** + * @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 {nonMaxSuppressionV4Impl} from '../backends/non_max_suppression_impl'; +import {Tensor1D, Tensor2D} from '../tensor'; +import {NamedTensorMap} from '../tensor_types'; +import {convertToTensor} from '../tensor_util_env'; +import {TensorLike} from '../types'; +import {nonMaxSuppSanityCheck} from './nonmax_util'; + +/** + * Asynchronously performs non maximum suppression of bounding boxes based on + * iou (intersection over union), with an option to pad results. + * + * @param boxes a 2d tensor of shape `[numBoxes, 4]`. Each entry is + * `[y1, x1, y2, x2]`, where `(y1, x1)` and `(y2, x2)` are the corners of + * the bounding box. + * @param scores a 1d tensor providing the box scores of shape `[numBoxes]`. + * @param maxOutputSize The maximum number of boxes to be selected. + * @param iouThreshold A float representing the threshold for deciding whether + * boxes overlap too much with respect to IOU. Must be between [0, 1]. + * Defaults to 0.5 (50% box overlap). + * @param scoreThreshold A threshold for deciding when to remove boxes based + * on score. Defaults to -inf, which means any score is accepted. + * @param padToMaxOutputSize Defalts to false. If true, size of output + * `selectedIndices` is padded to maxOutputSize. + * @return A map with the following properties: + * - selectedIndices: A 1D tensor with the selected box indices. + * - validOutputs: A scalar denoting how many elements in `selectedIndices` + * are valid. Valid elements occur first, then padding. + */ +/** @doc {heading: 'Operations', subheading: 'Images', namespace: 'image'} */ +async function nonMaxSuppressionPaddedAsync_( + boxes: Tensor2D|TensorLike, scores: Tensor1D|TensorLike, + maxOutputSize: number, iouThreshold = 0.5, + scoreThreshold = Number.NEGATIVE_INFINITY, + padToMaxOutputSize = false): Promise { + const $boxes = convertToTensor(boxes, 'boxes', 'nonMaxSuppressionAsync'); + const $scores = convertToTensor(scores, 'scores', 'nonMaxSuppressionAsync'); + + const params = nonMaxSuppSanityCheck( + $boxes, $scores, maxOutputSize, iouThreshold, scoreThreshold, + null /* softNmsSigma */); + const $maxOutputSize = params.maxOutputSize; + const $iouThreshold = params.iouThreshold; + const $scoreThreshold = params.scoreThreshold; + + const [boxesVals, scoresVals] = + await Promise.all([$boxes.data(), $scores.data()]); + + // We call a cpu based impl directly with the typedarray data here rather + // than a kernel because all kernels are synchronous (and thus cannot await + // .data()). + const res = nonMaxSuppressionV4Impl( + boxesVals, scoresVals, $maxOutputSize, $iouThreshold, $scoreThreshold, + padToMaxOutputSize); + + if ($boxes !== boxes) { + $boxes.dispose(); + } + if ($scores !== scores) { + $scores.dispose(); + } + return res; +} + +export const nonMaxSuppressionPaddedAsync = nonMaxSuppressionPaddedAsync_; diff --git a/tfjs-core/src/ops/non_max_suppression_padded_async_test.ts b/tfjs-core/src/ops/non_max_suppression_padded_async_test.ts new file mode 100644 index 00000000000..03e9680dce4 --- /dev/null +++ b/tfjs-core/src/ops/non_max_suppression_padded_async_test.ts @@ -0,0 +1,49 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysEqual} from '../test_util'; + +describeWithFlags('nonMaxSuppressionPaddedAsync', ALL_ENVS, () => { + it('select from three clusters with pad five.', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 5; + const iouThreshold = 0.5; + const scoreThreshold = 0.0; + + const before = tf.memory().numTensors; + + const {selectedIndices, validOutputs} = + await tf.image.nonMaxSuppressionPaddedAsync( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); + + const after = tf.memory().numTensors; + + expectArraysEqual(await selectedIndices.data(), [3, 0, 5, 0, 0]); + expectArraysEqual(await validOutputs.data(), 3); + + // The number of tensors should increase by the number of tensors + // returned (i.e. selectedIndices and selectedScores). + expect(after).toEqual(before + 2); + }); +}); diff --git a/tfjs-core/src/ops/non_max_suppression_padded_test.ts b/tfjs-core/src/ops/non_max_suppression_padded_test.ts index c174286d343..ad702de99e9 100644 --- a/tfjs-core/src/ops/non_max_suppression_padded_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_padded_test.ts @@ -57,4 +57,24 @@ describeWithFlags('nonMaxSuppressionPadded', ALL_ENVS, () => { expectArraysEqual(await selectedIndices.data(), [3, 0, 0, 0, 0, 0]); expectArraysEqual(await validOutputs.data(), 2); }); + + it('select from three clusters with no padding when pad option is false.', + async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 5; + const iouThreshold = 0.5; + const scoreThreshold = 0.0; + + const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, false); + + expectArraysEqual(await selectedIndices.data(), [3, 0, 5]); + expectArraysEqual(await validOutputs.data(), 3); + }); }); diff --git a/tfjs-core/src/ops/non_max_suppression_test.ts b/tfjs-core/src/ops/non_max_suppression_test.ts index e1c0ffb0b6f..e6e28bfd9ec 100644 --- a/tfjs-core/src/ops/non_max_suppression_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_test.ts @@ -16,213 +16,185 @@ */ import * as tf from '../index'; import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose, expectArraysEqual} from '../test_util'; +import {expectArraysEqual} from '../test_util'; describeWithFlags('nonMaxSuppression', ALL_ENVS, () => { - describe('NonMaxSuppression Basic', () => { - it('select from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - }); - - it('select from three clusters flipped coordinates', async () => { - const boxes = tf.tensor2d( - [ - 1, 1, 0, 0, 0, 0.1, 1, 1.1, 0, .9, 1, -0.1, - 0, 10, 1, 11, 1, 10.1, 0, 11.1, 1, 101, 0, 100 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - }); - - it('select at most two boxes from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 2; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([2]); - expectArraysEqual(await indices.data(), [3, 0]); - }); - - it('select at most thirty boxes from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 30; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - }); - - it('select single box', async () => { - const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); - const scores = tf.tensor1d([0.9]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([1]); - expectArraysEqual(await indices.data(), [0]); - }); - - it('select from ten identical boxes', async () => { - const numBoxes = 10; - const corners = new Array(numBoxes) - .fill(0) - .map(_ => [0, 0, 1, 1]) - .reduce((arr, curr) => arr.concat(curr)); - const boxes = tf.tensor2d(corners, [numBoxes, 4]); - const scores = tf.tensor1d(Array(numBoxes).fill(0.9)); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([1]); - expectArraysEqual(await indices.data(), [0]); - }); - - it('inconsistent box and score shapes', () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5]); - const maxOutputSize = 30; - const iouThreshold = 0.5; - const scoreThreshold = 0; - expect( - () => tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) - .toThrowError(/scores has incompatible shape with boxes/); - }); - - it('invalid iou threshold', () => { - const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); - const scores = tf.tensor1d([0.9]); - const maxOutputSize = 3; - const iouThreshold = 1.2; - const scoreThreshold = 0; - expect( - () => tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) - .toThrowError(/iouThreshold must be in \[0, 1\]/); - }); - - it('empty input', async () => { - const boxes = tf.tensor2d([], [0, 4]); - const scores = tf.tensor1d([]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([0]); - expectArraysEqual(await indices.data(), []); - }); - - it('accepts a tensor-like object', async () => { - const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; - const scores = [1, 2]; - const indices = tf.image.nonMaxSuppression(boxes, scores, 10); - expect(indices.shape).toEqual([2]); - expect(indices.dtype).toEqual('int32'); - expectArraysEqual(await indices.data(), [1, 0]); - }); - - it('works when inputs are not explicitly initialized on the CPU', - async () => { - // This test ensures that asynchronous backends work with NMS, which - // requires inputs to reside on the CPU. - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const a = tf.tensor1d([0, 1, -2, -4, 4, -4]); - const b = tf.tensor1d([0.15, 0.2, 0.25, 0.5, 0.7, 1.2]); - const scores = a.div(b); - const maxOutputSize = 2; - const iouThreshold = 0.5; - const scoreThreshold = 0; - await scores.data(); - const indices = tf.image.nonMaxSuppression( - boxes, scores as tf.Tensor1D, maxOutputSize, iouThreshold, - scoreThreshold); - - expect(indices.shape).toEqual([2]); - expectArraysEqual(await indices.data(), [4, 1]); - }); + it('select from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); }); - describe('NonMaxSuppressionWithScore', () => { - it('select from three clusters with SoftNMS', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 6; - const iouThreshold = 1.0; - const scoreThreshold = 0; - const softNmsSigma = 0.5; - - const {selectedIndices, selectedScores} = - tf.image.nonMaxSuppressionWithScore( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, - softNmsSigma); - - expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); - - expectArraysClose( - await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); - }); + it('select from three clusters flipped coordinates', async () => { + const boxes = tf.tensor2d( + [ + 1, 1, 0, 0, 0, 0.1, 1, 1.1, 0, .9, 1, -0.1, + 0, 10, 1, 11, 1, 10.1, 0, 11.1, 1, 101, 0, 100 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); }); + + it('select at most two boxes from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 2; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([2]); + expectArraysEqual(await indices.data(), [3, 0]); + }); + + it('select at most thirty boxes from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 30; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); + }); + + it('select single box', async () => { + const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); + const scores = tf.tensor1d([0.9]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([1]); + expectArraysEqual(await indices.data(), [0]); + }); + + it('select from ten identical boxes', async () => { + const numBoxes = 10; + const corners = new Array(numBoxes) + .fill(0) + .map(_ => [0, 0, 1, 1]) + .reduce((arr, curr) => arr.concat(curr)); + const boxes = tf.tensor2d(corners, [numBoxes, 4]); + const scores = tf.tensor1d(Array(numBoxes).fill(0.9)); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([1]); + expectArraysEqual(await indices.data(), [0]); + }); + + it('inconsistent box and score shapes', () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5]); + const maxOutputSize = 30; + const iouThreshold = 0.5; + const scoreThreshold = 0; + expect( + () => tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) + .toThrowError(/scores has incompatible shape with boxes/); + }); + + it('invalid iou threshold', () => { + const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); + const scores = tf.tensor1d([0.9]); + const maxOutputSize = 3; + const iouThreshold = 1.2; + const scoreThreshold = 0; + expect( + () => tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) + .toThrowError(/iouThreshold must be in \[0, 1\]/); + }); + + it('empty input', async () => { + const boxes = tf.tensor2d([], [0, 4]); + const scores = tf.tensor1d([]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([0]); + expectArraysEqual(await indices.data(), []); + }); + + it('accepts a tensor-like object', async () => { + const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; + const scores = [1, 2]; + const indices = tf.image.nonMaxSuppression(boxes, scores, 10); + expect(indices.shape).toEqual([2]); + expect(indices.dtype).toEqual('int32'); + expectArraysEqual(await indices.data(), [1, 0]); + }); + + it('works when inputs are not explicitly initialized on the CPU', + async () => { + // This test ensures that asynchronous backends work with NMS, which + // requires inputs to reside on the CPU. + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const a = tf.tensor1d([0, 1, -2, -4, 4, -4]); + const b = tf.tensor1d([0.15, 0.2, 0.25, 0.5, 0.7, 1.2]); + const scores = a.div(b); + const maxOutputSize = 2; + const iouThreshold = 0.5; + const scoreThreshold = 0; + await scores.data(); + const indices = tf.image.nonMaxSuppression( + boxes, scores as tf.Tensor1D, maxOutputSize, iouThreshold, + scoreThreshold); + + expect(indices.shape).toEqual([2]); + expectArraysEqual(await indices.data(), [4, 1]); + }); }); diff --git a/tfjs-core/src/ops/non_max_suppresion_with_score_async.ts b/tfjs-core/src/ops/non_max_suppression_with_score_async.ts similarity index 100% rename from tfjs-core/src/ops/non_max_suppresion_with_score_async.ts rename to tfjs-core/src/ops/non_max_suppression_with_score_async.ts diff --git a/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts b/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts new file mode 100644 index 00000000000..75abb7251fe --- /dev/null +++ b/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts @@ -0,0 +1,53 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose, expectArraysEqual} from '../test_util'; + +describeWithFlags('nonMaxSuppressionWithScoreAsync', ALL_ENVS, () => { + it('select from three clusters with SoftNMS', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 6; + const iouThreshold = 1.0; + const scoreThreshold = 0; + const softNmsSigma = 0.5; + + const numTensorsBefore = tf.memory().numTensors; + + const {selectedIndices, selectedScores} = + await tf.image.nonMaxSuppressionWithScoreAsync( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + softNmsSigma); + + const numTensorsAfter = tf.memory().numTensors; + + expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); + + expectArraysClose( + await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); + + // The number of tensors should increase by the number of tensors + // returned (i.e. selectedIndices and selectedScores). + expect(numTensorsAfter).toEqual(numTensorsBefore + 2); + }); +}); diff --git a/tfjs-core/src/ops/non_max_suppression_with_score_test.ts b/tfjs-core/src/ops/non_max_suppression_with_score_test.ts new file mode 100644 index 00000000000..4f1269b6e67 --- /dev/null +++ b/tfjs-core/src/ops/non_max_suppression_with_score_test.ts @@ -0,0 +1,45 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose, expectArraysEqual} from '../test_util'; + +describeWithFlags('nonMaxSuppressionWithScore', ALL_ENVS, () => { + it('select from three clusters with SoftNMS.', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 6; + const iouThreshold = 1.0; + const scoreThreshold = 0; + const softNmsSigma = 0.5; + + const {selectedIndices, selectedScores} = + tf.image.nonMaxSuppressionWithScore( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + softNmsSigma); + + expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); + + expectArraysClose( + await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); + }); +}); diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts index d02b2f2fc89..916ae12846e 100644 --- a/tfjs-core/src/ops/ops.ts +++ b/tfjs-core/src/ops/ops.ts @@ -238,8 +238,9 @@ import {rotateWithOffset} from './rotate_with_offset'; import {nonMaxSuppression} from './non_max_suppression'; import {nonMaxSuppressionAsync} from './non_max_suppression_async'; import {nonMaxSuppressionWithScore} from './non_max_suppression_with_score'; -import {nonMaxSuppressionWithScoreAsync} from './non_max_suppresion_with_score_async'; +import {nonMaxSuppressionWithScoreAsync} from './non_max_suppression_with_score_async'; import {nonMaxSuppressionPadded} from './non_max_suppression_padded'; +import {nonMaxSuppressionPaddedAsync} from './non_max_suppression_padded_async'; import {resizeBilinear} from './resize_bilinear'; import {resizeNearestNeighbor} from './resize_nearest_neighbor'; const image = { @@ -251,7 +252,8 @@ const image = { nonMaxSuppressionAsync, nonMaxSuppressionWithScore, nonMaxSuppressionWithScoreAsync, - nonMaxSuppressionPadded + nonMaxSuppressionPadded, + nonMaxSuppressionPaddedAsync }; // linalg namespace diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index 9bb3a4d3f8a..fc267e37e7c 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -148,8 +148,11 @@ import './ops/multi_rnn_cell_test'; import './ops/multinomial_test'; import './ops/neg_test'; import './ops/non_max_suppression_async_test'; +import './ops/non_max_suppression_padded_async_test'; import './ops/non_max_suppression_padded_test'; import './ops/non_max_suppression_test'; +import './ops/non_max_suppression_with_score_async_test'; +import './ops/non_max_suppression_with_score_test'; import './ops/norm_test'; import './ops/not_equal_test'; import './ops/one_hot_test'; From dbe589e471b70a01b573113fbf82128a9d45e18c Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 14 Jul 2020 10:29:09 -0700 Subject: [PATCH 3/7] Fix lint. --- tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts | 2 +- tfjs-core/src/ops/non_max_suppression_async_test.ts | 4 ++++ tfjs-core/src/ops/non_max_suppression_padded_test.ts | 3 +++ tfjs-core/src/ops/non_max_suppression_test.ts | 4 ++++ .../src/ops/non_max_suppression_with_score_async_test.ts | 6 +++--- tfjs-core/src/ops/non_max_suppression_with_score_test.ts | 4 ++++ 6 files changed, 19 insertions(+), 4 deletions(-) diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts index 632bfe0a9f2..c1f4e2cf6e6 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts +++ b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts @@ -75,5 +75,5 @@ export const nonMaxSuppressionV4Config: KernelConfig = { kernelName: 'NonMaxSuppressionV4', backendName: 'wasm', setupFunc: setup, - kernelFunc as {} as KernelFunc, + kernelFunc: kernelFunc as {} as KernelFunc, }; diff --git a/tfjs-core/src/ops/non_max_suppression_async_test.ts b/tfjs-core/src/ops/non_max_suppression_async_test.ts index 0ea530ac022..85c00ba9a5c 100644 --- a/tfjs-core/src/ops/non_max_suppression_async_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_async_test.ts @@ -30,11 +30,15 @@ describeWithFlags('nonMaxSuppressionAsync', ALL_ENVS, () => { const maxOutputSize = 3; const iouThreshold = 0.5; const scoreThreshold = 0; + + const before = tf.memory().numTensors; const indices = await tf.image.nonMaxSuppressionAsync( boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + const after = tf.memory().numTensors; expect(indices.shape).toEqual([3]); expectArraysEqual(await indices.data(), [3, 0, 5]); + expect(after).toEqual(before + 1); }); it('accepts a tensor-like object', async () => { diff --git a/tfjs-core/src/ops/non_max_suppression_padded_test.ts b/tfjs-core/src/ops/non_max_suppression_padded_test.ts index ad702de99e9..5fdfb17d7ff 100644 --- a/tfjs-core/src/ops/non_max_suppression_padded_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_padded_test.ts @@ -51,11 +51,14 @@ describeWithFlags('nonMaxSuppressionPadded', ALL_ENVS, () => { const iouThreshold = 0.5; const scoreThreshold = 0.4; + const before = tf.memory().numTensors; const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); + const after = tf.memory().numTensors; expectArraysEqual(await selectedIndices.data(), [3, 0, 0, 0, 0, 0]); expectArraysEqual(await validOutputs.data(), 2); + expect(after).toEqual(before + 2); }); it('select from three clusters with no padding when pad option is false.', diff --git a/tfjs-core/src/ops/non_max_suppression_test.ts b/tfjs-core/src/ops/non_max_suppression_test.ts index e6e28bfd9ec..6988318937e 100644 --- a/tfjs-core/src/ops/non_max_suppression_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_test.ts @@ -30,11 +30,15 @@ describeWithFlags('nonMaxSuppression', ALL_ENVS, () => { const maxOutputSize = 3; const iouThreshold = 0.5; const scoreThreshold = 0; + + const before = tf.memory().numTensors; const indices = tf.image.nonMaxSuppression( boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + const after = tf.memory().numTensors; expect(indices.shape).toEqual([3]); expectArraysEqual(await indices.data(), [3, 0, 5]); + expect(after).toEqual(before + 1); }); it('select from three clusters flipped coordinates', async () => { diff --git a/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts b/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts index 75abb7251fe..ab0cfe665d6 100644 --- a/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts @@ -32,14 +32,14 @@ describeWithFlags('nonMaxSuppressionWithScoreAsync', ALL_ENVS, () => { const scoreThreshold = 0; const softNmsSigma = 0.5; - const numTensorsBefore = tf.memory().numTensors; + const before = tf.memory().numTensors; const {selectedIndices, selectedScores} = await tf.image.nonMaxSuppressionWithScoreAsync( boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma); - const numTensorsAfter = tf.memory().numTensors; + const after = tf.memory().numTensors; expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); @@ -48,6 +48,6 @@ describeWithFlags('nonMaxSuppressionWithScoreAsync', ALL_ENVS, () => { // The number of tensors should increase by the number of tensors // returned (i.e. selectedIndices and selectedScores). - expect(numTensorsAfter).toEqual(numTensorsBefore + 2); + expect(after).toEqual(before + 2); }); }); diff --git a/tfjs-core/src/ops/non_max_suppression_with_score_test.ts b/tfjs-core/src/ops/non_max_suppression_with_score_test.ts index 4f1269b6e67..a1b4372e93e 100644 --- a/tfjs-core/src/ops/non_max_suppression_with_score_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_with_score_test.ts @@ -32,14 +32,18 @@ describeWithFlags('nonMaxSuppressionWithScore', ALL_ENVS, () => { const scoreThreshold = 0; const softNmsSigma = 0.5; + const before = tf.memory().numTensors; const {selectedIndices, selectedScores} = tf.image.nonMaxSuppressionWithScore( boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma); + const after = tf.memory().numTensors; expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); expectArraysClose( await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); + + expect(after).toEqual(before + 2); }); }); From e91dd0f0e223f0c9a6b01b961c570bfd3ade9524 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 14 Jul 2020 11:06:21 -0700 Subject: [PATCH 4/7] Add build rule. --- tfjs-backend-wasm/src/cc/BUILD | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD index fa8c56fb933..f1ab1a2101d 100644 --- a/tfjs-backend-wasm/src/cc/BUILD +++ b/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"], From dd285dc2ce6036c4205088c53aaf8836d098dd85 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 14 Jul 2020 16:54:06 -0700 Subject: [PATCH 5/7] Fix test. --- .../src/kernels/NonMaxSuppressionV4.ts | 8 +- .../src/ops/non_max_suppression_async_test.ts | 118 +++-- .../non_max_suppression_padded_async_test.ts | 49 -- .../ops/non_max_suppression_padded_test.ts | 83 ---- tfjs-core/src/ops/non_max_suppression_test.ts | 437 +++++++++++------- ...n_max_suppression_with_score_async_test.ts | 53 --- .../non_max_suppression_with_score_test.ts | 49 -- tfjs-core/src/tests.ts | 2 - 8 files changed, 360 insertions(+), 439 deletions(-) delete mode 100644 tfjs-core/src/ops/non_max_suppression_padded_async_test.ts delete mode 100644 tfjs-core/src/ops/non_max_suppression_padded_test.ts delete mode 100644 tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts delete mode 100644 tfjs-core/src/ops/non_max_suppression_with_score_test.ts diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts index c1f4e2cf6e6..f04fc38f1ed 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts +++ b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -import {KernelConfig, KernelFunc, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, TensorInfo} from '@tensorflow/tfjs-core'; +import {KernelConfig, KernelFunc, NonMaxSuppressionV4, NonMaxSuppressionV4Attrs, NonMaxSuppressionV4Inputs, TensorInfo} from '@tensorflow/tfjs-core'; import {BackendWasm} from '../backend_wasm'; @@ -40,7 +40,7 @@ function setup(backend: BackendWasm): void { ]); } -function kernelFunc(args: { +function nonMaxSuppressionV4(args: { backend: BackendWasm, inputs: NonMaxSuppressionV4Inputs, attrs: NonMaxSuppressionV4Attrs @@ -72,8 +72,8 @@ function kernelFunc(args: { } export const nonMaxSuppressionV4Config: KernelConfig = { - kernelName: 'NonMaxSuppressionV4', + kernelName: NonMaxSuppressionV4, backendName: 'wasm', setupFunc: setup, - kernelFunc: kernelFunc as {} as KernelFunc, + kernelFunc: nonMaxSuppressionV4 as {} as KernelFunc, }; diff --git a/tfjs-core/src/ops/non_max_suppression_async_test.ts b/tfjs-core/src/ops/non_max_suppression_async_test.ts index 85c00ba9a5c..40ffe9d2473 100644 --- a/tfjs-core/src/ops/non_max_suppression_async_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_async_test.ts @@ -16,37 +16,99 @@ */ import * as tf from '../index'; import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysEqual} from '../test_util'; +import {expectArraysClose, expectArraysEqual} from '../test_util'; describeWithFlags('nonMaxSuppressionAsync', ALL_ENVS, () => { - it('select from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - - const before = tf.memory().numTensors; - const indices = await tf.image.nonMaxSuppressionAsync( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - const after = tf.memory().numTensors; - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - expect(after).toEqual(before + 1); + describe('NonMaxSuppressionAsync basic', () => { + it('select from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = await tf.image.nonMaxSuppressionAsync( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); + }); + + it('accepts a tensor-like object', async () => { + const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; + const scores = [1, 2]; + const indices = await tf.image.nonMaxSuppressionAsync(boxes, scores, 10); + expect(indices.shape).toEqual([2]); + expect(indices.dtype).toEqual('int32'); + expectArraysEqual(await indices.data(), [1, 0]); + }); }); - it('accepts a tensor-like object', async () => { - const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; - const scores = [1, 2]; - const indices = await tf.image.nonMaxSuppressionAsync(boxes, scores, 10); - expect(indices.shape).toEqual([2]); - expect(indices.dtype).toEqual('int32'); - expectArraysEqual(await indices.data(), [1, 0]); + describe('NonMaxSuppressionWithScoreAsync', () => { + it('select from three clusters with SoftNMS', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 6; + const iouThreshold = 1.0; + const scoreThreshold = 0; + const softNmsSigma = 0.5; + + const numTensorsBefore = tf.memory().numTensors; + + const {selectedIndices, selectedScores} = + await tf.image.nonMaxSuppressionWithScoreAsync( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + softNmsSigma); + + const numTensorsAfter = tf.memory().numTensors; + + expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); + + expectArraysClose( + await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); + + // The number of tensors should increase by the number of tensors + // returned (i.e. selectedIndices and selectedScores). + expect(numTensorsAfter).toEqual(numTensorsBefore + 2); + }); + }); + + describe('NonMaxSuppressionPaddedAsync', () => { + it('select from three clusters with pad five.', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 5; + const iouThreshold = 0.5; + const scoreThreshold = 0.0; + + const before = tf.memory().numTensors; + + const {selectedIndices, validOutputs} = + await tf.image.nonMaxSuppressionPaddedAsync( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); + + const after = tf.memory().numTensors; + + expectArraysEqual(await selectedIndices.data(), [3, 0, 5, 0, 0]); + expectArraysEqual(await validOutputs.data(), 3); + + // The number of tensors should increase by the number of tensors + // returned (i.e. selectedIndices and selectedScores). + expect(after).toEqual(before + 2); + }); }); }); diff --git a/tfjs-core/src/ops/non_max_suppression_padded_async_test.ts b/tfjs-core/src/ops/non_max_suppression_padded_async_test.ts deleted file mode 100644 index 03e9680dce4..00000000000 --- a/tfjs-core/src/ops/non_max_suppression_padded_async_test.ts +++ /dev/null @@ -1,49 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysEqual} from '../test_util'; - -describeWithFlags('nonMaxSuppressionPaddedAsync', ALL_ENVS, () => { - it('select from three clusters with pad five.', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 5; - const iouThreshold = 0.5; - const scoreThreshold = 0.0; - - const before = tf.memory().numTensors; - - const {selectedIndices, validOutputs} = - await tf.image.nonMaxSuppressionPaddedAsync( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); - - const after = tf.memory().numTensors; - - expectArraysEqual(await selectedIndices.data(), [3, 0, 5, 0, 0]); - expectArraysEqual(await validOutputs.data(), 3); - - // The number of tensors should increase by the number of tensors - // returned (i.e. selectedIndices and selectedScores). - expect(after).toEqual(before + 2); - }); -}); diff --git a/tfjs-core/src/ops/non_max_suppression_padded_test.ts b/tfjs-core/src/ops/non_max_suppression_padded_test.ts deleted file mode 100644 index 5fdfb17d7ff..00000000000 --- a/tfjs-core/src/ops/non_max_suppression_padded_test.ts +++ /dev/null @@ -1,83 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysEqual} from '../test_util'; - -describeWithFlags('nonMaxSuppressionPadded', ALL_ENVS, () => { - it('select from three clusters with pad five.', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 5; - const iouThreshold = 0.5; - const scoreThreshold = 0.0; - - const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); - - expectArraysEqual(await selectedIndices.data(), [3, 0, 5, 0, 0]); - expectArraysEqual(await validOutputs.data(), 3); - }); - - it('select from three clusters with pad five and score threshold.', - async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 6; - const iouThreshold = 0.5; - const scoreThreshold = 0.4; - - const before = tf.memory().numTensors; - const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); - const after = tf.memory().numTensors; - - expectArraysEqual(await selectedIndices.data(), [3, 0, 0, 0, 0, 0]); - expectArraysEqual(await validOutputs.data(), 2); - expect(after).toEqual(before + 2); - }); - - it('select from three clusters with no padding when pad option is false.', - async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 5; - const iouThreshold = 0.5; - const scoreThreshold = 0.0; - - const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, false); - - expectArraysEqual(await selectedIndices.data(), [3, 0, 5]); - expectArraysEqual(await validOutputs.data(), 3); - }); -}); diff --git a/tfjs-core/src/ops/non_max_suppression_test.ts b/tfjs-core/src/ops/non_max_suppression_test.ts index 6988318937e..4645938290c 100644 --- a/tfjs-core/src/ops/non_max_suppression_test.ts +++ b/tfjs-core/src/ops/non_max_suppression_test.ts @@ -16,189 +16,284 @@ */ import * as tf from '../index'; import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysEqual} from '../test_util'; +import {expectArraysClose, expectArraysEqual} from '../test_util'; describeWithFlags('nonMaxSuppression', ALL_ENVS, () => { - it('select from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - - const before = tf.memory().numTensors; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - const after = tf.memory().numTensors; - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - expect(after).toEqual(before + 1); - }); + describe('NonMaxSuppression Basic', () => { + it('select from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - it('select from three clusters flipped coordinates', async () => { - const boxes = tf.tensor2d( - [ - 1, 1, 0, 0, 0, 0.1, 1, 1.1, 0, .9, 1, -0.1, - 0, 10, 1, 11, 1, 10.1, 0, 11.1, 1, 101, 0, 100 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - }); + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); + }); - it('select at most two boxes from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 2; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([2]); - expectArraysEqual(await indices.data(), [3, 0]); - }); + it('select from three clusters flipped coordinates', async () => { + const boxes = tf.tensor2d( + [ + 1, 1, 0, 0, 0, 0.1, 1, 1.1, 0, .9, 1, -0.1, + 0, 10, 1, 11, 1, 10.1, 0, 11.1, 1, 101, 0, 100 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - it('select at most thirty boxes from three clusters', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 30; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([3]); - expectArraysEqual(await indices.data(), [3, 0, 5]); - }); + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); + }); - it('select single box', async () => { - const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); - const scores = tf.tensor1d([0.9]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([1]); - expectArraysEqual(await indices.data(), [0]); - }); + it('select at most two boxes from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 2; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - it('select from ten identical boxes', async () => { - const numBoxes = 10; - const corners = new Array(numBoxes) - .fill(0) - .map(_ => [0, 0, 1, 1]) - .reduce((arr, curr) => arr.concat(curr)); - const boxes = tf.tensor2d(corners, [numBoxes, 4]); - const scores = tf.tensor1d(Array(numBoxes).fill(0.9)); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([1]); - expectArraysEqual(await indices.data(), [0]); - }); + expect(indices.shape).toEqual([2]); + expectArraysEqual(await indices.data(), [3, 0]); + }); - it('inconsistent box and score shapes', () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5]); - const maxOutputSize = 30; - const iouThreshold = 0.5; - const scoreThreshold = 0; - expect( - () => tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) - .toThrowError(/scores has incompatible shape with boxes/); - }); + it('select at most thirty boxes from three clusters', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 30; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - it('invalid iou threshold', () => { - const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); - const scores = tf.tensor1d([0.9]); - const maxOutputSize = 3; - const iouThreshold = 1.2; - const scoreThreshold = 0; - expect( - () => tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) - .toThrowError(/iouThreshold must be in \[0, 1\]/); - }); + expect(indices.shape).toEqual([3]); + expectArraysEqual(await indices.data(), [3, 0, 5]); + }); + + it('select single box', async () => { + const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); + const scores = tf.tensor1d([0.9]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([1]); + expectArraysEqual(await indices.data(), [0]); + }); + + it('select from ten identical boxes', async () => { + const numBoxes = 10; + const corners = new Array(numBoxes) + .fill(0) + .map(_ => [0, 0, 1, 1]) + .reduce((arr, curr) => arr.concat(curr)); + const boxes = tf.tensor2d(corners, [numBoxes, 4]); + const scores = tf.tensor1d(Array(numBoxes).fill(0.9)); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); + + expect(indices.shape).toEqual([1]); + expectArraysEqual(await indices.data(), [0]); + }); + + it('inconsistent box and score shapes', () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5]); + const maxOutputSize = 30; + const iouThreshold = 0.5; + const scoreThreshold = 0; + expect( + () => tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) + .toThrowError(/scores has incompatible shape with boxes/); + }); + + it('invalid iou threshold', () => { + const boxes = tf.tensor2d([0, 0, 1, 1], [1, 4]); + const scores = tf.tensor1d([0.9]); + const maxOutputSize = 3; + const iouThreshold = 1.2; + const scoreThreshold = 0; + expect( + () => tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)) + .toThrowError(/iouThreshold must be in \[0, 1\]/); + }); + + it('empty input', async () => { + const boxes = tf.tensor2d([], [0, 4]); + const scores = tf.tensor1d([]); + const maxOutputSize = 3; + const iouThreshold = 0.5; + const scoreThreshold = 0; + const indices = tf.image.nonMaxSuppression( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - it('empty input', async () => { - const boxes = tf.tensor2d([], [0, 4]); - const scores = tf.tensor1d([]); - const maxOutputSize = 3; - const iouThreshold = 0.5; - const scoreThreshold = 0; - const indices = tf.image.nonMaxSuppression( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold); - - expect(indices.shape).toEqual([0]); - expectArraysEqual(await indices.data(), []); + expect(indices.shape).toEqual([0]); + expectArraysEqual(await indices.data(), []); + }); + + it('accepts a tensor-like object', async () => { + const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; + const scores = [1, 2]; + const indices = tf.image.nonMaxSuppression(boxes, scores, 10); + expect(indices.shape).toEqual([2]); + expect(indices.dtype).toEqual('int32'); + expectArraysEqual(await indices.data(), [1, 0]); + }); + + it('works when inputs are not explicitly initialized on the CPU', + async () => { + // This test ensures that asynchronous backends work with NMS, which + // requires inputs to reside on the CPU. + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const a = tf.tensor1d([0, 1, -2, -4, 4, -4]); + const b = tf.tensor1d([0.15, 0.2, 0.25, 0.5, 0.7, 1.2]); + const scores = a.div(b); + const maxOutputSize = 2; + const iouThreshold = 0.5; + const scoreThreshold = 0; + await scores.data(); + const indices = tf.image.nonMaxSuppression( + boxes, scores as tf.Tensor1D, maxOutputSize, iouThreshold, + scoreThreshold); + + expect(indices.shape).toEqual([2]); + expectArraysEqual(await indices.data(), [4, 1]); + }); }); - it('accepts a tensor-like object', async () => { - const boxes = [[0, 0, 1, 1], [0, 1, 1, 2]]; - const scores = [1, 2]; - const indices = tf.image.nonMaxSuppression(boxes, scores, 10); - expect(indices.shape).toEqual([2]); - expect(indices.dtype).toEqual('int32'); - expectArraysEqual(await indices.data(), [1, 0]); + describe('NonMaxSuppressionWithScore', () => { + it('select from three clusters with SoftNMS', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 6; + const iouThreshold = 1.0; + const scoreThreshold = 0; + const softNmsSigma = 0.5; + + const {selectedIndices, selectedScores} = + tf.image.nonMaxSuppressionWithScore( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + softNmsSigma); + + expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); + + expectArraysClose( + await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); + }); }); - it('works when inputs are not explicitly initialized on the CPU', - async () => { - // This test ensures that asynchronous backends work with NMS, which - // requires inputs to reside on the CPU. - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const a = tf.tensor1d([0, 1, -2, -4, 4, -4]); - const b = tf.tensor1d([0.15, 0.2, 0.25, 0.5, 0.7, 1.2]); - const scores = a.div(b); - const maxOutputSize = 2; - const iouThreshold = 0.5; - const scoreThreshold = 0; - await scores.data(); - const indices = tf.image.nonMaxSuppression( - boxes, scores as tf.Tensor1D, maxOutputSize, iouThreshold, - scoreThreshold); - - expect(indices.shape).toEqual([2]); - expectArraysEqual(await indices.data(), [4, 1]); - }); + describe('NonMaxSuppressionPadded', () => { + it('select from three clusters with pad five.', async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 5; + const iouThreshold = 0.5; + const scoreThreshold = 0; + + const before = tf.memory().numTensors; + const {selectedIndices, validOutputs} = tf.image.nonMaxSuppressionPadded( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, true); + const after = tf.memory().numTensors; + + expectArraysEqual(await selectedIndices.data(), [3, 0, 5, 0, 0]); + expectArraysEqual(await validOutputs.data(), 3); + expect(after).toEqual(before + 2); + }); + + it('select from three clusters with pad five and score threshold.', + async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 6; + const iouThreshold = 0.5; + const scoreThreshold = 0.4; + + const before = tf.memory().numTensors; + const {selectedIndices, validOutputs} = + tf.image.nonMaxSuppressionPadded( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + true); + const after = tf.memory().numTensors; + + expectArraysEqual(await selectedIndices.data(), [3, 0, 0, 0, 0, 0]); + expectArraysEqual(await validOutputs.data(), 2); + expect(after).toEqual(before + 2); + }); + + it('select from three clusters with no padding when pad option is false.', + async () => { + const boxes = tf.tensor2d( + [ + 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, + 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 + ], + [6, 4]); + const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); + const maxOutputSize = 5; + const iouThreshold = 0.5; + const scoreThreshold = 0.0; + + const {selectedIndices, validOutputs} = + tf.image.nonMaxSuppressionPadded( + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + false); + + expectArraysEqual(await selectedIndices.data(), [3, 0, 5]); + expectArraysEqual(await validOutputs.data(), 3); + }); + }); }); diff --git a/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts b/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts deleted file mode 100644 index ab0cfe665d6..00000000000 --- a/tfjs-core/src/ops/non_max_suppression_with_score_async_test.ts +++ /dev/null @@ -1,53 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose, expectArraysEqual} from '../test_util'; - -describeWithFlags('nonMaxSuppressionWithScoreAsync', ALL_ENVS, () => { - it('select from three clusters with SoftNMS', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 6; - const iouThreshold = 1.0; - const scoreThreshold = 0; - const softNmsSigma = 0.5; - - const before = tf.memory().numTensors; - - const {selectedIndices, selectedScores} = - await tf.image.nonMaxSuppressionWithScoreAsync( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, - softNmsSigma); - - const after = tf.memory().numTensors; - - expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); - - expectArraysClose( - await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); - - // The number of tensors should increase by the number of tensors - // returned (i.e. selectedIndices and selectedScores). - expect(after).toEqual(before + 2); - }); -}); diff --git a/tfjs-core/src/ops/non_max_suppression_with_score_test.ts b/tfjs-core/src/ops/non_max_suppression_with_score_test.ts deleted file mode 100644 index a1b4372e93e..00000000000 --- a/tfjs-core/src/ops/non_max_suppression_with_score_test.ts +++ /dev/null @@ -1,49 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose, expectArraysEqual} from '../test_util'; - -describeWithFlags('nonMaxSuppressionWithScore', ALL_ENVS, () => { - it('select from three clusters with SoftNMS.', async () => { - const boxes = tf.tensor2d( - [ - 0, 0, 1, 1, 0, 0.1, 1, 1.1, 0, -0.1, 1, 0.9, - 0, 10, 1, 11, 0, 10.1, 1, 11.1, 0, 100, 1, 101 - ], - [6, 4]); - const scores = tf.tensor1d([0.9, 0.75, 0.6, 0.95, 0.5, 0.3]); - const maxOutputSize = 6; - const iouThreshold = 1.0; - const scoreThreshold = 0; - const softNmsSigma = 0.5; - - const before = tf.memory().numTensors; - const {selectedIndices, selectedScores} = - tf.image.nonMaxSuppressionWithScore( - boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, - softNmsSigma); - const after = tf.memory().numTensors; - - expectArraysEqual(await selectedIndices.data(), [3, 0, 1, 5, 4, 2]); - - expectArraysClose( - await selectedScores.data(), [0.95, 0.9, 0.384, 0.3, 0.256, 0.197]); - - expect(after).toEqual(before + 2); - }); -}); diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index fc267e37e7c..3ac3d392d38 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -151,8 +151,6 @@ import './ops/non_max_suppression_async_test'; import './ops/non_max_suppression_padded_async_test'; import './ops/non_max_suppression_padded_test'; import './ops/non_max_suppression_test'; -import './ops/non_max_suppression_with_score_async_test'; -import './ops/non_max_suppression_with_score_test'; import './ops/norm_test'; import './ops/not_equal_test'; import './ops/one_hot_test'; From beec196a305f3678595c25dd2894382f3b3f6d8e Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 14 Jul 2020 17:23:34 -0700 Subject: [PATCH 6/7] Fix test. --- tfjs-core/src/tests.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index 3ac3d392d38..0e08d5965d2 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -148,8 +148,6 @@ import './ops/multi_rnn_cell_test'; import './ops/multinomial_test'; import './ops/neg_test'; import './ops/non_max_suppression_async_test'; -import './ops/non_max_suppression_padded_async_test'; -import './ops/non_max_suppression_padded_test'; import './ops/non_max_suppression_test'; import './ops/norm_test'; import './ops/not_equal_test'; From e30be8ce76ed561fd4fed44e798455c6e2d7e011 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 16 Jul 2020 17:37:52 -0700 Subject: [PATCH 7/7] Clean up. --- tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts | 2 +- tfjs-core/src/backends/non_max_suppression_impl.ts | 12 +++--------- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts index f04fc38f1ed..6d5b47647b5 100644 --- a/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts +++ b/tfjs-backend-wasm/src/kernels/NonMaxSuppressionV4.ts @@ -28,7 +28,7 @@ let wasmFunc: ( function setup(backend: BackendWasm): void { wasmFunc = backend.wasm.cwrap( - 'NonMaxSuppressionV4', + NonMaxSuppressionV4, 'number', // Result* [ 'number', // boxesId diff --git a/tfjs-core/src/backends/non_max_suppression_impl.ts b/tfjs-core/src/backends/non_max_suppression_impl.ts index 3ccf8429c1e..82a34581f9f 100644 --- a/tfjs-core/src/backends/non_max_suppression_impl.ts +++ b/tfjs-core/src/backends/non_max_suppression_impl.ts @@ -47,15 +47,9 @@ export function nonMaxSuppressionV4Impl( iouThreshold: number, scoreThreshold: number, padToMaxOutputSize: boolean): NamedTensorMap { return nonMaxSuppressionImpl_( - boxes, - scores, - maxOutputSize, - iouThreshold, - scoreThreshold, - 0, /* softNmsSigma */ - false, /* returnScoresTensor */ - padToMaxOutputSize, /* padToMaxOutputSize */ - true, + boxes, scores, maxOutputSize, iouThreshold, scoreThreshold, + 0 /* softNmsSigma */, false /* returnScoresTensor */, + padToMaxOutputSize /* padToMaxOutputSize */, true /* returnValidOutputs */); }