For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

non_max_suppression

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

def non_max_suppression[dtype: DType, FuncType: def(Int64, Int64, Int64) -> None & ImplicitlyCopyable](boxes: TileTensor[dtype, Storage=boxes.Storage, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type], scores: TileTensor[dtype, Storage=scores.Storage, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type], max_output_boxes_per_class: Int, iou_threshold: Float32, score_threshold: Float32, func: FuncType)

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

Parameters:

• ​dtype (DType): The data type for box coordinates and scores.
• ​FuncType (def(Int64, Int64, Int64) -> None & ImplicitlyCopyable): Type of the func callback invoked for each selected box.

Args:

• ​boxes (TileTensor[dtype, Storage=boxes.Storage, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type]): Rank-3 tensor of bounding boxes with shape (batch, num_boxes, 4). Each box is [y1, x1, y2, x2].
• ​scores (TileTensor[dtype, Storage=scores.Storage, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type]): Rank-3 tensor of detection scores with shape (batch, num_classes, num_boxes).
• ​max_output_boxes_per_class (Int): Maximum number of boxes to select per class.
• ​iou_threshold (Float32): IoU threshold for suppression. Boxes with IoU above this value are suppressed.
• ​score_threshold (Float32): Minimum score for a box to be considered. Boxes below this are filtered out.
• ​func (FuncType): Callback invoked for each selected box with the batch index, class index, and box index.