Mojo function

non_max_suppression

non_max_suppression[dtype: DType](boxes: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], scores: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], output: LayoutTensor[DType.int64, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], max_output_boxes_per_class: Int, iou_threshold: Float32, score_threshold: Float32)

Perform Non-Maximum Suppression (NMS) on bounding boxes.

This is a buffer semantic overload that writes results directly to an output tensor. NMS iteratively selects boxes with highest scores while suppressing nearby boxes with high overlap (IoU).

Parameters:

dtype (DType): The data type for box coordinates and scores.

Args:

boxes (LayoutTensor): Rank-3 tensor of bounding boxes with shape (batch, num_boxes, 4). Each box is [y1, x1, y2, x2].
scores (LayoutTensor): Rank-3 tensor of scores with shape (batch, num_classes, num_boxes).
output (LayoutTensor): Rank-2 output tensor to store selected boxes as (N, 3) where each row is [batch_idx, class_idx, box_idx].
max_output_boxes_per_class (Int): Maximum number of boxes to select per class.
iou_threshold (Float32): IoU threshold for suppression. Boxes with IoU > threshold are suppressed.
score_threshold (Float32): Minimum score threshold. Boxes with score < threshold are filtered out.

non_max_suppression[dtype: DType, func: fn(Int64, Int64, Int64) capturing -> None](boxes: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], scores: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], max_output_boxes_per_class: Int, iou_threshold: Float32, score_threshold: Float32)

Implements the NonMaxSuppression operator from the ONNX spec https://github.com/onnx/onnx/blob/main/docs/Operators.md#nonmaxsuppression.