[jacinto-ai/pytorch-jacinto-ai-devkit.git] / modules / pytorch_jacinto_ai / vision / ops / roi_pool.py
1 import torch
2 from torch import nn
4 from torch.autograd import Function
5 from torch.autograd.function import once_differentiable
7 from torch.nn.modules.utils import _pair
9 from ..extension import _lazy_import
10 from ._utils import convert_boxes_to_roi_format
13 class _RoIPoolFunction(Function):
14 @staticmethod
15 def forward(ctx, input, rois, output_size, spatial_scale):
16 ctx.output_size = _pair(output_size)
17 ctx.spatial_scale = spatial_scale
18 ctx.input_shape = input.size()
19 _C = _lazy_import()
20 output, argmax = _C.roi_pool_forward(
21 input, rois, spatial_scale,
22 output_size[0], output_size[1])
23 ctx.save_for_backward(rois, argmax)
24 return output
26 @staticmethod
27 @once_differentiable
28 def backward(ctx, grad_output):
29 rois, argmax = ctx.saved_tensors
30 output_size = ctx.output_size
31 spatial_scale = ctx.spatial_scale
32 bs, ch, h, w = ctx.input_shape
33 _C = _lazy_import()
34 grad_input = _C.roi_pool_backward(
35 grad_output, rois, argmax, spatial_scale,
36 output_size[0], output_size[1], bs, ch, h, w)
37 return grad_input, None, None, None
40 def roi_pool(input, boxes, output_size, spatial_scale=1.0):
41 """
42 Performs Region of Interest (RoI) Pool operator described in Fast R-CNN
44 Arguments:
45 input (Tensor[N, C, H, W]): input tensor
46 boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2)
47 format where the regions will be taken from. If a single Tensor is passed,
48 then the first column should contain the batch index. If a list of Tensors
49 is passed, then each Tensor will correspond to the boxes for an element i
50 in a batch
51 output_size (int or Tuple[int, int]): the size of the output after the cropping
52 is performed, as (height, width)
53 spatial_scale (float): a scaling factor that maps the input coordinates to
54 the box coordinates. Default: 1.0
56 Returns:
57 output (Tensor[K, C, output_size[0], output_size[1]])
58 """
59 rois = boxes
60 if not isinstance(rois, torch.Tensor):
61 rois = convert_boxes_to_roi_format(rois)
62 return _RoIPoolFunction.apply(input, rois, output_size, spatial_scale)
65 class RoIPool(nn.Module):
66 """
67 See roi_pool
68 """
69 def __init__(self, output_size, spatial_scale):
70 super(RoIPool, self).__init__()
71 self.output_size = output_size
72 self.spatial_scale = spatial_scale
74 def forward(self, input, rois):
75 return roi_pool(input, rois, self.output_size, self.spatial_scale)
77 def __repr__(self):
78 tmpstr = self.__class__.__name__ + '('
79 tmpstr += 'output_size=' + str(self.output_size)
80 tmpstr += ', spatial_scale=' + str(self.spatial_scale)
81 tmpstr += ')'
82 return tmpstr