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_shape_func
def non_max_suppression_shape_func[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], max_output_boxes_per_class: Int, iou_threshold: Float32, score_threshold: Float32) -> IndexList[Int(2)]
Compute the output shape for NMS without allocating the output buffer.
This function performs a dry-run of NMS to determine how many boxes will be selected, allowing proper output buffer allocation. Can be removed once the graph compiler supports value semantic kernels that allocate their own output.
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). - 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). - max_output_boxes_per_class (
Int): Maximum number of boxes to select per class. - iou_threshold (
Float32): IoU threshold for suppression. - score_threshold (
Float32): Minimum score threshold.
Returns:
IndexList[Int(2)]: A 2-element IndexList specifying the output shape (num_selected_boxes, 3).