# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Non-maximum suppression operators.""" from . import _ffi_api def all_class_non_max_suppression( boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, output_format="onnx", ): """Non-maximum suppression operator for object detection, corresponding to ONNX NonMaxSuppression and TensorFlow combined_non_max_suppression. NMS is performed for each class separately. Parameters ---------- boxes : relax.Expr 3-D tensor with shape (batch_size, num_boxes, 4) scores: relax.Expr 3-D tensor with shape (batch_size, num_classes, num_boxes) max_output_boxes_per_class : relax.Expr The maxinum number of output selected boxes per class iou_threshold : relax.Expr IoU test threshold score_threshold : relax.Expr Score threshold to filter out low score boxes early output_format : str, optional "onnx" or "tensorflow", see below. Returns ------- out : relax.Expr If `output_format` is "onnx", the output is two tensors. The first is `indices` of size `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor `num_total_detection` of shape `(1,)` representing the total number of selected boxes. The three values in `indices` encode batch, class, and box indices. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come first, in descending of scores, followed by boxes from batch 0, class 1 etc. The output uses dynamic_strided_slice to trim to only valid detections, so the first tensor has shape (num_total_detection, 3) containing only valid rows. If `output_format` is "tensorflow", the output is three tensors, the first is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size `(batch_size,)` representing the total number of selected boxes per batch. The two values in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at batch b, only the first `num_total_detection[b]` entries are valid. The second axis of `indices` and `scores` are sorted within each class by box scores, but not across classes. So the box indices and scores for the class 0 come first in a sorted order, followed by the class 1 etc. """ return _ffi_api.all_class_non_max_suppression( boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, output_format ) def get_valid_counts(data, score_threshold=0, id_index=0, score_index=1): """Get valid count of bounding boxes given a score threshold. Also moves valid boxes to the top of input data. Parameters ---------- data : relax.Expr 3-D tensor with shape [batch_size, num_anchors, elem_length]. score_threshold : float, optional Lower limit of score for valid bounding boxes. id_index : int, optional Index of the class categories. Set to ``-1`` to disable the class-id check. score_index : int, optional Index of the scores/confidence of boxes. Returns ------- out : relax.Expr A tuple ``(valid_count, out_tensor, out_indices)`` where ``valid_count`` has shape ``[batch_size]``, ``out_tensor`` has shape ``[batch_size, num_anchors, elem_length]``, and ``out_indices`` has shape ``[batch_size, num_anchors]``. """ return _ffi_api.get_valid_counts(data, score_threshold, id_index, score_index) def non_max_suppression( data, valid_count, indices, max_output_size=-1, iou_threshold=0.5, force_suppress=False, top_k=-1, coord_start=2, score_index=1, id_index=0, return_indices=True, invalid_to_bottom=False, soft_nms_sigma=0.0, score_threshold=0.0, ): """Non-maximum suppression operator for object detection. Parameters ---------- data : relax.Expr 3-D tensor with shape [batch_size, num_anchors, elem_length]. valid_count : relax.Expr 1-D tensor for valid number of boxes. indices : relax.Expr 2-D tensor with shape [batch_size, num_anchors]. max_output_size : int, optional Max number of output valid boxes, -1 for no limit. iou_threshold : float, optional Non-maximum suppression IoU threshold. force_suppress : bool, optional Whether to suppress all detections regardless of class_id. When ``id_index`` is ``-1``, all valid boxes are treated as belonging to the same class, so this flag has the same effect as ``True``. top_k : int, optional Keep maximum top k detections before nms, -1 for no limit. coord_start : int, optional Start index of the consecutive 4 coordinates. score_index : int, optional Index of the scores/confidence of boxes. id_index : int, optional Index of the class categories. Set to ``-1`` to suppress boxes across all classes. return_indices : bool, optional Whether to return box indices in input data. invalid_to_bottom : bool, optional Whether to move valid bounding boxes to the top of the returned tensor. This option only affects the ``return_indices=False`` path. soft_nms_sigma : float, optional Sigma for soft-NMS Gaussian penalty. When ``0.0`` (default), standard hard NMS is used. Positive values decay overlapping box scores instead of suppressing them outright. score_threshold : float, optional Post-decay minimum score for a box to remain eligible during soft-NMS. Only used when ``soft_nms_sigma > 0``. This is distinct from ``get_valid_counts.score_threshold``, which filters boxes before NMS. Defaults to ``0.0``. Returns ------- out : relax.Expr The return tuple shape depends on ``soft_nms_sigma``. If ``return_indices`` is ``True`` and ``soft_nms_sigma`` is ``0.0``, returns a 2-tuple ``(box_indices, valid_box_count)`` with shapes ``[batch_size, num_anchors]`` and ``[batch_size, 1]``. If ``return_indices`` is ``True`` and ``soft_nms_sigma > 0``, returns a 3-tuple ``(out_data, box_indices, valid_box_count)`` where decayed ``out_data`` is prepended and has the same shape as the input data. Otherwise returns the modified data tensor. """ return _ffi_api.non_max_suppression( data, valid_count, indices, max_output_size, iou_threshold, force_suppress, top_k, coord_start, score_index, id_index, return_indices, invalid_to_bottom, soft_nms_sigma, score_threshold, )