TensorFlow.js tf.image.nonMaxSuppression()函数介绍

tf.image.nonMaxSuppression()函数是TensorFlow.js中的一个图像处理函数，用于对输入的边界框进行非极大值抑制（Non-Maximum Suppression，简称NMS），以选择具有高置信度的框。

函数签名

tf.image.nonMaxSuppression(
  boxes: tf.Tensor2D | number[][],
  scores: tf.Tensor1D | number[],
  maxOutputSize: number,
  iouThreshold?: number,
  scoreThreshold?: number
): tf.Tensor1D

函数参数

• boxes: tf.Tensor2D | number[][]，形状为 [numBoxes, 4] 的浮点数张量，表示 $[y_{min}, x_{min}, y_{max}, x_{max}]$ 形式的一个或多个边界框；
• scores: tf.Tensor1D | number[]，形状为 (numBoxes,) 的浮点数张量，表示每个边界框的置信度；
• maxOutputSize: number，输出张量的最大长度；
• iouThreshold: number，可选参数，默认值为 0.5，表示区分两个边界框时使用的 IoU 阈值；
• scoreThreshold: number，可选参数，默认值为 -Infinity，表示用于筛选具有高置信度的框的阈值。

函数返回值

本函数的返回值为形状为 (numBoxes,) 的张量，包含非极大值抑制（NMS）的结果。

使用示例

const boxes = tf.tensor2d([
  [0.1, 0.6, 0.4, 0.9],
  [0.3, 0.2, 0.7, 0.8],
  [0.4, 0.5, 0.8, 1.0]
]);

const scores = tf.tensor1d([0.9, 0.75, 0.8]);

const maxOutputSize = 2;

const iouThreshold = 0.3;

const scoreThreshold = 0.5;

const selectedIndices = tf.image.nonMaxSuppression(
    boxes, scores, maxOutputSize, iouThreshold, scoreThreshold);

console.log(selectedIndices.arraySync());

在上面的示例中，我们使用了 tf.image.nonMaxSuppression() 函数对 boxes 的三个边界框进行了非极大值抑制。参数 maxOutputSize 指定了输出张量的最大长度为 2，参数 iouThreshold 指定了 IoU 阈值为 0.3，参数 scoreThreshold 指定了用于筛选具有高置信度的框的阈值为 0.5。函数输出的是非极大值抑制的结果，即选出的两个具有高置信度的框的索引：[0, 2]。

参考链接

更多信息请参考TensorFlow.js API 文档。