ai.onnx.NonMaxSuppression

ai.onnx · standard ONNX operator · ONNX opset ≥ 11

Description

Filters boxes whose intersection over union (IoU) with a higher-scoring selected box exceeds iou_threshold. Processing is independent per batch and class, and returns logical int64 [batch_index, class_index, box_index] triples with lossless uint32 WebGPU storage. The int64 selection limit uses a semantics-preserving saturating projection to uint32. Because the result length is data-dependent, callers provide its exact shape. The current backend implements float32 boxes, scores, and thresholds.

See the ONNX NonMaxSuppression spec for the reference semantics.

Inputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`boxes`	`boxes`	`T`	same as logical dtype	`3`	—	Bounding box coordinates with shape `[num_batches, spatial_dimension, 4]`; box format is controlled by `center_point_box`.	required
`scores`	`scores`	`T`	same as logical dtype	`3`	—	Per-class confidence scores with shape `[num_batches, num_classes, spatial_dimension]`.	required
`max_output_boxes_per_class`	`max_output_boxes_per_class`	`M`	`uint32`	`0`	—	Optional logical int64 scalar limiting boxes selected per batch and class. When omitted, the ONNX default is zero and the output is empty. WebGPU maps non-positive values to zero and values above uint32 range to `UINT32_MAX`; this preserves results because no group can select more boxes than its finite input.	optional
`iou_threshold`	`iou_threshold`	`T`	same as logical dtype	`0`	—	Optional scalar IoU threshold in `[0, 1]`; boxes whose IoU is strictly greater are suppressed. Defaults to zero.	optional
`score_threshold`	`score_threshold`	`T`	same as logical dtype	`0`	—	Optional scalar score threshold. When present, only boxes whose score is strictly greater than the threshold are considered; when omitted, scores are not filtered.	optional

Outputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`selected_indices`	`selected_indices`	`I`	`uint32`	`2`	—	Logical int64 selected-box indices with shape `[num_selected_indices, 3]`, each row containing `[batch_index, class_index, box_index]`. WebGPU stores the bounded indices as uint32.	required

Attributes

Default values (overridable per request):

Attribute	Default	Description
`center_point_box`	`0`	Box coordinate format: 0 for corner format `[y1, x1, y2, x2]`, 1 for center format `[x_center, y_center, width, height]`.

Type constraints

Variable	Allowed dtypes
`T`	`float32`
`M`	`int64`
`I`	`int64`

Files

metadata.json — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
nms-bitmask.wgsl.jinja
nms-groups-compact.wgsl.jinja
nms-groups-select.wgsl.jinja
non-max-suppression.wgsl.jinja

Use with `@huggingface/kernels`

The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.

The explicit outputs entries provide shape and logical dtype metadata for the results listed below:

selected_indices

Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model version.

Replace each *Data placeholder with a typed array containing the corresponding input data.

import { getKernel } from "@huggingface/kernels";

const kernel = await getKernel("webgpu-kernels/ai.onnx.NonMaxSuppression", { version: 1 });
// Explicit destinations request optional results or supply metadata that cannot be inferred.
const { selected_indices } = await kernel({
  boxes: { data: boxesData, shape: [1, 1, 4] },
  scores: { data: scoresData, shape: [1, 1, 1] },
  max_output_boxes_per_class: { data: max_output_boxes_per_classData, shape: [] },
  iou_threshold: { data: iou_thresholdData, shape: [] },
  score_threshold: { data: score_thresholdData, shape: [] },
}, {
  outputs: { selected_indices: { shape: [1, 3], dtype: "int64" } },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.