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
)

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.

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.