1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17"""Faster R-CNN and Mask R-CNN operations."""
18from . import _make
19
20
21def roi_align(data, rois, pooled_size, spatial_scale, sample_ratio=-1, layout="NCHW"):
22    """ROI align operator.
23
24    Parameters
25    ----------
26    data : relay.Expr
27        4-D tensor with shape [batch, channel, height, width]
28
29    rois : relay.Expr
30        2-D tensor with shape [num_roi, 5]. The last dimension should be in format of
31        [batch_index, w_start, h_start, w_end, h_end]
32
33    pooled_size : list/tuple of two ints
34        output size
35
36    spatial_scale : float
37        Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal
38        of total stride in convolutional layers, which should be in range (0.0, 1.0]
39
40    sample_ratio : int
41        Optional sampling ratio of ROI align, using adaptive size by default.
42
43    Returns
44    -------
45    output : relay.Expr
46        4-D tensor with shape [num_roi, channel, pooled_size, pooled_size]
47    """
48    return _make.roi_align(data, rois, pooled_size, spatial_scale, sample_ratio, layout)
49
50
51def roi_pool(data, rois, pooled_size, spatial_scale, layout="NCHW"):
52    """ROI pool operator.
53
54    Parameters
55    ----------
56    data : relay.Expr
57        4-D tensor with shape [batch, channel, height, width]
58
59    rois : relay.Expr
60        2-D tensor with shape [num_roi, 5]. The last dimension should be in format of
61        [batch_index, w_start, h_start, w_end, h_end]
62
63    pooled_size : list/tuple of two ints
64        output size
65
66    spatial_scale : float
67        Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal
68        of total stride in convolutional layers, which should be in range (0.0, 1.0]
69
70    Returns
71    -------
72    output : relay.Expr
73        4-D tensor with shape [num_roi, channel, pooled_size, pooled_size]
74    """
75    return _make.roi_pool(data, rois, pooled_size, spatial_scale, layout)
76
77
78def proposal(
79    cls_prob,
80    bbox_pred,
81    im_info,
82    scales,
83    ratios,
84    feature_stride,
85    threshold,
86    rpn_pre_nms_top_n,
87    rpn_post_nms_top_n,
88    rpn_min_size,
89    iou_loss,
90):
91    """Proposal operator.
92
93    Parameters
94    ----------
95    cls_prob : relay.Expr
96        4-D tensor with shape [batch, 2 * num_anchors, height, width].
97
98    bbox_pred : relay.Expr
99        4-D tensor with shape [batch, 4 * num_anchors, height, width].
100
101    im_info : relay.Expr
102        2-D tensor with shape [batch, 3]. The last dimension should be in format of
103        [im_height, im_width, im_scale]
104
105    scales : list/tuple of float
106        Scales of anchor windows.
107
108    ratios : list/tuple of float
109        Ratios of anchor windows.
110
111    feature_stride : int
112        The size of the receptive field each unit in the convolution layer of the rpn, for example
113        the product of all stride's prior to this layer.
114
115    threshold : float
116        Non-maximum suppression threshold.
117
118    rpn_pre_nms_top_n : int
119        Number of top scoring boxes to apply NMS. -1 to use all boxes.
120
121    rpn_post_nms_top_n : int
122        Number of top scoring boxes to keep after applying NMS to RPN proposals.
123
124    rpn_min_size : int
125        Minimum height or width in proposal.
126
127    iou_loss : bool
128        Usage of IoU loss.
129
130    Returns
131    -------
132    output : relay.Expr
133        2-D tensor with shape [batch * rpn_post_nms_top_n, 5]. The last dimension is in format of
134        [batch_index, w_start, h_start, w_end, h_end].
135    """
136    return _make.proposal(
137        cls_prob,
138        bbox_pred,
139        im_info,
140        scales,
141        ratios,
142        feature_stride,
143        threshold,
144        rpn_pre_nms_top_n,
145        rpn_post_nms_top_n,
146        rpn_min_size,
147        iou_loss,
148    )
149