nms.py 3.58 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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.
Yao Wang committed
17 18 19
"""Non-maximum suppression operations."""
from __future__ import absolute_import as _abs
from . import _make
20
from ...expr import TupleWrapper
Yao Wang committed
21

22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
def get_valid_counts(data,
                     score_threshold):
    """Get valid count of bounding boxes given a score threshold.
    Also moves valid boxes to the top of input data.

    Parameters
    ----------
    data : relay.Expr
        Input data. 3-D tensor with shape [batch_size, num_anchors, 6].

    score_threshold : optional, float
        Lower limit of score for valid bounding boxes.

    Returns
    -------
    valid_count : relay.Expr
        1-D tensor for valid number of boxes.

    out_tensor : relay.Expr
        Rearranged data tensor.
    """
    return TupleWrapper(_make.get_valid_counts(data, score_threshold), 2)


def non_max_suppression(data,
                        valid_count,
                        max_output_size=-1,
                        iou_threshold=0.5,
                        force_suppress=False,
                        top_k=-1,
52 53
                        coord_start=2,
                        score_index=1,
54 55 56
                        id_index=0,
                        return_indices=True,
                        invalid_to_bottom=False):
Yao Wang committed
57 58 59 60 61 62 63 64 65 66 67 68
    """Non-maximum suppression operator for object detection.

    Parameters
    ----------
    data : relay.Expr
        3-D tensor with shape [batch_size, num_anchors, 6].
        The last dimension should be in format of
        [class_id, score, box_left, box_top, box_right, box_bottom].

    valid_count : relay.Expr
        1-D tensor for valid number of boxes.

69 70 71 72 73
    max_output_size : int, optional
        Max number of output valid boxes for each instance.
        By default all valid boxes are returned.

    iou_threshold : float, optional
Yao Wang committed
74 75 76 77 78
        Non-maximum suppression threshold.

    force_suppress : bool, optional
        Suppress all detections regardless of class_id.

79
    top_k : int, optional
Yao Wang committed
80 81
        Keep maximum top k detections before nms, -1 for no limit.

82 83 84 85 86 87
    coord_start : int, optional
        The starting index of the consecutive 4 coordinates.

    score_index : int, optional
        Index of the scores/confidence of boxes.

88 89 90 91 92 93 94 95 96
    id_index : int, optional
        index of the class categories, -1 to disable.

    return_indices : bool, optional
        Whether to return box indices in input data.

    invalid_to_bottom : bool, optional
        Whether to move all valid bounding boxes to the top.

Yao Wang committed
97 98 99 100 101
    Returns
    -------
    out : relay.Expr
        3-D tensor with shape [batch_size, num_anchors, 6].
    """
102 103
    return _make.non_max_suppression(data, valid_count, max_output_size,
                                     iou_threshold, force_suppress, top_k,
104 105
                                     coord_start, score_index, id_index,
                                     return_indices, invalid_to_bottom)