tfm.vision.nms.sorted_non_max_suppression_padded

A wrapper that handles non-maximum suppression.

tfm.vision.nms.sorted_non_max_suppression_padded(
    scores, boxes, max_output_size, iou_threshold
)

  • The boxes are sorted by scores unless the box is a dot (all coordinates are zero).
  • Boxes with higher scores can be used to suppress boxes with lower scores.

The overal design of the algorithm is to handle boxes tile-by-tile:

boxes = boxes.pad_to_multiply_of(tile_size) num_tiles = len(boxes) // tile_size output_boxes = [] for i in range(num_tiles): box_tile = boxes[itile_size : (i+1)tile_size] for j in range(i - 1): suppressing_tile = boxes[jtile_size : (j+1)tile_size] iou = bbox_overlap(box_tile, suppressing_tile) # if the box is suppressed in iou, clear it to a dot box_tile *= _update_boxes(iou) # Iteratively handle the diagnal tile. iou = _box_overlap(box_tile, box_tile) iou_changed = True while iou_changed: # boxes that are not suppressed by anything else suppressing_boxes = _get_suppressing_boxes(iou) # boxes that are suppressed by suppressing_boxes suppressed_boxes = _get_suppressed_boxes(iou, suppressing_boxes) # clear iou to 0 for boxes that are suppressed, as they cannot be used # to suppress other boxes any more new_iou = _clear_iou(iou, suppressed_boxes) iou_changed = (new_iou != iou) iou = new_iou # remaining boxes that can still suppress others, are selected boxes. output_boxes.append(_get_suppressing_boxes(iou)) if len(output_boxes) >= max_output_size: break

scores a tensor with a shape of [batch_size, anchors].
boxes a tensor with a shape of [batch_size, anchors, 4].
max_output_size a scalar integer Tensor representing the maximum number of boxes to be selected by non max suppression.
iou_threshold a float representing the threshold for deciding whether boxes overlap too much with respect to IOU.

nms_scores a tensor with a shape of [batch_size, anchors]. It has same dtype as input scores.
nms_proposals a tensor with a shape of [batch_size, anchors, 4]. It has same dtype as input boxes.