Mojo function

non_max_suppression

non_max_suppression[dtype: DType](boxes: TileTensor[dtype, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type, element_size=boxes.element_size], scores: TileTensor[dtype, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type, element_size=scores.element_size], output: TileTensor[DType.int64, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], 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 (TileTensor[dtype, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type, element_size=boxes.element_size]): Rank-3 tensor of bounding boxes with shape (batch, num_boxes, 4). Each box is [y1, x1, y2, x2].
scores (TileTensor[dtype, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type, element_size=scores.element_size]): Rank-3 tensor of scores with shape (batch, num_classes, num_boxes).
output (TileTensor[DType.int64, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size]): 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: def(Int64, Int64, Int64) capturing -> None](boxes: TileTensor[dtype, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type, element_size=boxes.element_size], scores: TileTensor[dtype, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type, element_size=scores.element_size], 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.