Spaces:
Build error
Build error
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from typing import Tuple, Union | |
| import torch | |
| import torch.nn as nn | |
| from torch.autograd import Function | |
| from torch.autograd.function import once_differentiable | |
| from torch.nn.modules.utils import _pair | |
| from ..utils import ext_loader | |
| ext_module = ext_loader.load_ext( | |
| '_ext', ['bezier_align_forward', 'bezier_align_backward']) | |
| class BezierAlignFunction(Function): | |
| def forward(ctx, | |
| input: torch.Tensor, | |
| beziers: torch.Tensor, | |
| output_size: Union[int, Tuple[int, int]], | |
| spatial_scale: Union[int, float] = 1.0, | |
| sampling_ratio: int = 0, | |
| aligned: bool = True) -> torch.Tensor: | |
| ctx.output_size = _pair(output_size) | |
| ctx.spatial_scale = spatial_scale | |
| ctx.input_shape = input.size() | |
| ctx.sampling_ratio = sampling_ratio | |
| ctx.aligned = aligned | |
| assert beziers.size(1) == 17 | |
| output_shape = (beziers.size(0), input.size(1), ctx.output_size[0], | |
| ctx.output_size[1]) | |
| output = input.new_zeros(output_shape) | |
| ext_module.bezier_align_forward( | |
| input, | |
| beziers, | |
| output, | |
| aligned_height=ctx.output_size[0], | |
| aligned_width=ctx.output_size[1], | |
| spatial_scale=ctx.spatial_scale, | |
| sampling_ratio=ctx.sampling_ratio, | |
| aligned=ctx.aligned) | |
| ctx.save_for_backward(beziers) | |
| return output | |
| def backward(ctx, grad_output: torch.Tensor): | |
| beziers = ctx.saved_tensors[0] | |
| grad_input = grad_output.new_zeros(ctx.input_shape) | |
| grad_output = grad_output.contiguous() | |
| ext_module.bezier_align_backward( | |
| grad_output, | |
| beziers, | |
| grad_input, | |
| aligned_height=ctx.output_size[0], | |
| aligned_width=ctx.output_size[1], | |
| spatial_scale=ctx.spatial_scale, | |
| sampling_ratio=ctx.sampling_ratio, | |
| aligned=ctx.aligned) | |
| return grad_input, None, None, None, None, None | |
| bezier_align = BezierAlignFunction.apply | |
| class BezierAlign(nn.Module): | |
| """Bezier align pooling layer. | |
| Args: | |
| output_size (tuple): h, w | |
| spatial_scale (float): scale the input boxes by this number | |
| sampling_ratio (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. | |
| Note: | |
| The implementation of BezierAlign is modified from | |
| https://github.com/aim-uofa/AdelaiDet | |
| 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, | |
| output_size: Tuple, | |
| spatial_scale: Union[int, float], | |
| sampling_ratio: int, | |
| aligned: bool = True, | |
| ) -> None: | |
| super().__init__() | |
| self.output_size = _pair(output_size) | |
| self.spatial_scale = float(spatial_scale) | |
| self.sampling_ratio = int(sampling_ratio) | |
| self.aligned = aligned | |
| def forward(self, input: torch.Tensor, | |
| beziers: torch.Tensor) -> torch.Tensor: | |
| """BezierAlign forward. | |
| Args: | |
| inputs (Tensor): input features. | |
| beziers (Tensor): beziers for align. | |
| """ | |
| return bezier_align(input, beziers, self.output_size, | |
| self.spatial_scale, self.sampling_ratio, | |
| self.aligned) | |
| def __repr__(self): | |
| s = self.__class__.__name__ | |
| s += f'(output_size={self.output_size}, ' | |
| s += f'spatial_scale={self.spatial_scale})' | |
| s += f'sampling_ratio={self.sampling_ratio})' | |
| s += f'aligned={self.aligned})' | |
| return s | |