def roi_pool(input, boxes, output_size, spatial_scale=1.0): """ Performs Region of Interest (RoI) Pool operator described in Fast R-CNN Arguments: input (Tensor[N, C, H, W]): input tensor boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) format where the regions will be taken from. If a single Tensor is passed, then the first column should contain the batch index. If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i in a batch output_size (int or Tuple[int, int]): the size of the output after the cropping is performed, as (height, width) spatial_scale (float): a scaling factor that maps the input coordinates to the box coordinates. Default: 1.0 Returns: output (Tensor[K, C, output_size[0], output_size[1]]) """ rois = boxes if not isinstance(rois, torch.Tensor): rois = convert_boxes_to_roi_format(rois) # TODO: Change this to support backwards, which we # do not currently support when JIT tracing. if torch._C._get_tracing_state(): _lazy_import() output, _ = torch.ops.torchvision.roi_pool(input, rois, spatial_scale, output_size[0], output_size[1]) return output return _RoIPoolFunction.apply(input, rois, output_size, spatial_scale)
def roi_align(input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1): """ Performs Region of Interest (RoI) Align operator described in Mask R-CNN Arguments: input (Tensor[N, C, H, W]): input tensor boxes (Tensor[K, 5] or List[Tensor[L, 4]]): the box coordinates in (x1, y1, x2, y2) format where the regions will be taken from. If a single Tensor is passed, then the first column should contain the batch index. If a list of Tensors is passed, then each Tensor will correspond to the boxes for an element i in a batch output_size (int or Tuple[int, int]): the size of the output after the cropping is performed, as (height, width) spatial_scale (float): a scaling factor that maps the input coordinates to the box coordinates. Default: 1.0 sampling_ratio (int): number of sampling points in the interpolation grid used to compute the output value of each pooled output bin. If > 0, then exactly sampling_ratio x sampling_ratio grid points are used. If <= 0, then an adaptive number of grid points are used (computed as ceil(roi_width / pooled_w), and likewise for height). Default: -1 Returns: output (Tensor[K, C, output_size[0], output_size[1]]) """ rois = boxes if not isinstance(rois, torch.Tensor): rois = convert_boxes_to_roi_format(rois) _lazy_import() return torch.ops.torchvision.roi_align(input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio)
def nms(boxes, scores, iou_threshold): """ Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union (IoU). NMS iteratively removes lower scoring boxes which have an IoU greater than iou_threshold with another (higher scoring) box. Parameters ---------- boxes : Tensor[N, 4]) boxes to perform NMS on. They are expected to be in (x1, y1, x2, y2) format scores : Tensor[N] scores for each one of the boxes iou_threshold : float discards all overlapping boxes with IoU < iou_threshold Returns ------- keep : Tensor int64 tensor with the indices of the elements that have been kept by NMS, sorted in decreasing order of scores """ _lazy_import() return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
def forward(ctx, input, rois, output_size, spatial_scale): ctx.output_size = _pair(output_size) ctx.spatial_scale = spatial_scale ctx.input_shape = input.size() _C = _lazy_import() output, argmax = _C.roi_pool_forward(input, rois, spatial_scale, output_size[0], output_size[1]) ctx.save_for_backward(rois, argmax) return output
def backward(ctx, grad_output): rois, argmax = ctx.saved_tensors output_size = ctx.output_size spatial_scale = ctx.spatial_scale bs, ch, h, w = ctx.input_shape _C = _lazy_import() grad_input = _C.roi_pool_backward(grad_output, rois, argmax, spatial_scale, output_size[0], output_size[1], bs, ch, h, w) return grad_input, None, None, None
def backward(ctx, grad_output): rois, = ctx.saved_tensors output_size = ctx.output_size spatial_scale = ctx.spatial_scale sampling_ratio = ctx.sampling_ratio bs, ch, h, w = ctx.input_shape _C = _lazy_import() grad_input = _C.roi_align_backward(grad_output, rois, spatial_scale, output_size[0], output_size[1], bs, ch, h, w, sampling_ratio) return grad_input, None, None, None, None
def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio): ctx.save_for_backward(roi) ctx.output_size = _pair(output_size) ctx.spatial_scale = spatial_scale ctx.sampling_ratio = sampling_ratio ctx.input_shape = input.size() _C = _lazy_import() output = _C.roi_align_forward(input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio) return output
def nms(boxes, scores, iou_threshold): """ Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union (IoU). NMS iteratively removes lower scoring boxes which have an IoU greater than iou_threshold with another (higher scoring) box. Arguments: boxes (Tensor[N, 4]): boxes to perform NMS on scores (Tensor[N]): scores for each one of the boxes iou_threshold (float): discards all overlapping boxes with IoU < iou_threshold Returns: keep (Tensor): int64 tensor with the indices of the elements that have been kept by NMS, sorted in decreasing order of scores """ _C = _lazy_import() return _C.nms(boxes, scores, iou_threshold)