tfm.vision.box_ops.top_k_boxes

Sorts and select top k boxes according to the scores.

tfm.vision.box_ops.top_k_boxes(
    boxes, scores, k
)

Args
boxes	a tensor of shape [batch_size, N, 4] representing the coordinate of the boxes. N is the number of boxes per image.
scores	a tensor of shsape [batch_size, N] representing the socre of the boxes.
k	an integer or a tensor indicating the top k number.

Returns
selected_boxes	a tensor of shape [batch_size, k, 4] representing the selected top k box coordinates.
selected_scores	a tensor of shape [batch_size, k] representing the selected top k box scores.