| | |
| | import torch.nn as nn |
| | from torch.autograd import Function |
| |
|
| | from ..utils import ext_loader |
| |
|
| | ext_module = ext_loader.load_ext( |
| | '_ext', ['roi_align_rotated_forward', 'roi_align_rotated_backward']) |
| |
|
| |
|
| | class RoIAlignRotatedFunction(Function): |
| |
|
| | @staticmethod |
| | def symbolic(g, features, rois, out_size, spatial_scale, sample_num, |
| | aligned, clockwise): |
| | if isinstance(out_size, int): |
| | out_h = out_size |
| | out_w = out_size |
| | elif isinstance(out_size, tuple): |
| | assert len(out_size) == 2 |
| | assert isinstance(out_size[0], int) |
| | assert isinstance(out_size[1], int) |
| | out_h, out_w = out_size |
| | else: |
| | raise TypeError( |
| | '"out_size" must be an integer or tuple of integers') |
| | return g.op( |
| | 'mmcv::MMCVRoIAlignRotated', |
| | features, |
| | rois, |
| | output_height_i=out_h, |
| | output_width_i=out_h, |
| | spatial_scale_f=spatial_scale, |
| | sampling_ratio_i=sample_num, |
| | aligned_i=aligned, |
| | clockwise_i=clockwise) |
| |
|
| | @staticmethod |
| | def forward(ctx, |
| | features, |
| | rois, |
| | out_size, |
| | spatial_scale, |
| | sample_num=0, |
| | aligned=True, |
| | clockwise=False): |
| | if isinstance(out_size, int): |
| | out_h = out_size |
| | out_w = out_size |
| | elif isinstance(out_size, tuple): |
| | assert len(out_size) == 2 |
| | assert isinstance(out_size[0], int) |
| | assert isinstance(out_size[1], int) |
| | out_h, out_w = out_size |
| | else: |
| | raise TypeError( |
| | '"out_size" must be an integer or tuple of integers') |
| | ctx.spatial_scale = spatial_scale |
| | ctx.sample_num = sample_num |
| | ctx.aligned = aligned |
| | ctx.clockwise = clockwise |
| | ctx.save_for_backward(rois) |
| | ctx.feature_size = features.size() |
| |
|
| | batch_size, num_channels, data_height, data_width = features.size() |
| | num_rois = rois.size(0) |
| |
|
| | output = features.new_zeros(num_rois, num_channels, out_h, out_w) |
| | ext_module.roi_align_rotated_forward( |
| | features, |
| | rois, |
| | output, |
| | pooled_height=out_h, |
| | pooled_width=out_w, |
| | spatial_scale=spatial_scale, |
| | sample_num=sample_num, |
| | aligned=aligned, |
| | clockwise=clockwise) |
| | return output |
| |
|
| | @staticmethod |
| | def backward(ctx, grad_output): |
| | feature_size = ctx.feature_size |
| | spatial_scale = ctx.spatial_scale |
| | aligned = ctx.aligned |
| | clockwise = ctx.clockwise |
| | sample_num = ctx.sample_num |
| | rois = ctx.saved_tensors[0] |
| | assert feature_size is not None |
| | batch_size, num_channels, data_height, data_width = feature_size |
| |
|
| | out_w = grad_output.size(3) |
| | out_h = grad_output.size(2) |
| |
|
| | grad_input = grad_rois = None |
| |
|
| | if ctx.needs_input_grad[0]: |
| | grad_input = rois.new_zeros(batch_size, num_channels, data_height, |
| | data_width) |
| | ext_module.roi_align_rotated_backward( |
| | grad_output.contiguous(), |
| | rois, |
| | grad_input, |
| | pooled_height=out_h, |
| | pooled_width=out_w, |
| | spatial_scale=spatial_scale, |
| | sample_num=sample_num, |
| | aligned=aligned, |
| | clockwise=clockwise) |
| | return grad_input, grad_rois, None, None, None, None, None |
| |
|
| |
|
| | roi_align_rotated = RoIAlignRotatedFunction.apply |
| |
|
| |
|
| | class RoIAlignRotated(nn.Module): |
| | """RoI align pooling layer for rotated proposals. |
| | |
| | It accepts a feature map of shape (N, C, H, W) and rois with shape |
| | (n, 6) with each roi decoded as (batch_index, center_x, center_y, |
| | w, h, angle). The angle is in radian. |
| | |
| | Args: |
| | out_size (tuple): h, w |
| | spatial_scale (float): scale the input boxes by this number |
| | sample_num (int): number of inputs samples to take for each |
| | output sample. 0 to take samples densely for current models. |
| | aligned (bool): if False, use the legacy implementation in |
| | MMDetection. If True, align the results more perfectly. |
| | Default: True. |
| | clockwise (bool): If True, the angle in each proposal follows a |
| | clockwise fashion in image space, otherwise, the angle is |
| | counterclockwise. Default: False. |
| | |
| | Note: |
| | The implementation of RoIAlign when aligned=True is modified from |
| | https://github.com/facebookresearch/detectron2/ |
| | |
| | The meaning of aligned=True: |
| | |
| | Given a continuous coordinate c, its two neighboring pixel |
| | indices (in our pixel model) are computed by floor(c - 0.5) and |
| | ceil(c - 0.5). For example, c=1.3 has pixel neighbors with discrete |
| | indices [0] and [1] (which are sampled from the underlying signal |
| | at continuous coordinates 0.5 and 1.5). But the original roi_align |
| | (aligned=False) does not subtract the 0.5 when computing |
| | neighboring pixel indices and therefore it uses pixels with a |
| | slightly incorrect alignment (relative to our pixel model) when |
| | performing bilinear interpolation. |
| | |
| | With `aligned=True`, |
| | we first appropriately scale the ROI and then shift it by -0.5 |
| | prior to calling roi_align. This produces the correct neighbors; |
| | |
| | The difference does not make a difference to the model's |
| | performance if ROIAlign is used together with conv layers. |
| | """ |
| |
|
| | def __init__(self, |
| | out_size, |
| | spatial_scale, |
| | sample_num=0, |
| | aligned=True, |
| | clockwise=False): |
| | super(RoIAlignRotated, self).__init__() |
| |
|
| | self.out_size = out_size |
| | self.spatial_scale = float(spatial_scale) |
| | self.sample_num = int(sample_num) |
| | self.aligned = aligned |
| | self.clockwise = clockwise |
| |
|
| | def forward(self, features, rois): |
| | return RoIAlignRotatedFunction.apply(features, rois, self.out_size, |
| | self.spatial_scale, |
| | self.sample_num, self.aligned, |
| | self.clockwise) |
| |
|
