Spaces:
Build error
Build error
# Copyright (c) Facebook, Inc. and its affiliates. | |
import math | |
from typing import List | |
import torch | |
from torch import nn | |
from torchvision.ops import RoIPool | |
from detectron2.layers import ROIAlign, ROIAlignRotated, cat, nonzero_tuple | |
from detectron2.structures import Boxes | |
""" | |
To export ROIPooler to torchscript, in this file, variables that should be annotated with | |
`Union[List[Boxes], List[RotatedBoxes]]` are only annotated with `List[Boxes]`. | |
TODO: Correct these annotations when torchscript support `Union`. | |
https://github.com/pytorch/pytorch/issues/41412 | |
""" | |
__all__ = ["ROIPooler"] | |
def assign_boxes_to_levels( | |
box_lists: List[Boxes], | |
min_level: int, | |
max_level: int, | |
canonical_box_size: int, | |
canonical_level: int, | |
): | |
""" | |
Map each box in `box_lists` to a feature map level index and return the assignment | |
vector. | |
Args: | |
box_lists (list[Boxes] | list[RotatedBoxes]): A list of N Boxes or N RotatedBoxes, | |
where N is the number of images in the batch. | |
min_level (int): Smallest feature map level index. The input is considered index 0, | |
the output of stage 1 is index 1, and so. | |
max_level (int): Largest feature map level index. | |
canonical_box_size (int): A canonical box size in pixels (sqrt(box area)). | |
canonical_level (int): The feature map level index on which a canonically-sized box | |
should be placed. | |
Returns: | |
A tensor of length M, where M is the total number of boxes aggregated over all | |
N batch images. The memory layout corresponds to the concatenation of boxes | |
from all images. Each element is the feature map index, as an offset from | |
`self.min_level`, for the corresponding box (so value i means the box is at | |
`self.min_level + i`). | |
""" | |
box_sizes = torch.sqrt(cat([boxes.area() for boxes in box_lists])) | |
# Eqn.(1) in FPN paper | |
level_assignments = torch.floor( | |
canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8) | |
) | |
# clamp level to (min, max), in case the box size is too large or too small | |
# for the available feature maps | |
level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level) | |
return level_assignments.to(torch.int64) - min_level | |
def _fmt_box_list(box_tensor, batch_index: int): | |
repeated_index = torch.full_like( | |
box_tensor[:, :1], batch_index, dtype=box_tensor.dtype, device=box_tensor.device | |
) | |
return cat((repeated_index, box_tensor), dim=1) | |
def convert_boxes_to_pooler_format(box_lists: List[Boxes]): | |
""" | |
Convert all boxes in `box_lists` to the low-level format used by ROI pooling ops | |
(see description under Returns). | |
Args: | |
box_lists (list[Boxes] | list[RotatedBoxes]): | |
A list of N Boxes or N RotatedBoxes, where N is the number of images in the batch. | |
Returns: | |
When input is list[Boxes]: | |
A tensor of shape (M, 5), where M is the total number of boxes aggregated over all | |
N batch images. | |
The 5 columns are (batch index, x0, y0, x1, y1), where batch index | |
is the index in [0, N) identifying which batch image the box with corners at | |
(x0, y0, x1, y1) comes from. | |
When input is list[RotatedBoxes]: | |
A tensor of shape (M, 6), where M is the total number of boxes aggregated over all | |
N batch images. | |
The 6 columns are (batch index, x_ctr, y_ctr, width, height, angle_degrees), | |
where batch index is the index in [0, N) identifying which batch image the | |
rotated box (x_ctr, y_ctr, width, height, angle_degrees) comes from. | |
""" | |
pooler_fmt_boxes = cat( | |
[_fmt_box_list(box_list.tensor, i) for i, box_list in enumerate(box_lists)], dim=0 | |
) | |
return pooler_fmt_boxes | |
class ROIPooler(nn.Module): | |
""" | |
Region of interest feature map pooler that supports pooling from one or more | |
feature maps. | |
""" | |
def __init__( | |
self, | |
output_size, | |
scales, | |
sampling_ratio, | |
pooler_type, | |
canonical_box_size=224, | |
canonical_level=4, | |
): | |
""" | |
Args: | |
output_size (int, tuple[int] or list[int]): output size of the pooled region, | |
e.g., 14 x 14. If tuple or list is given, the length must be 2. | |
scales (list[float]): The scale for each low-level pooling op relative to | |
the input image. For a feature map with stride s relative to the input | |
image, scale is defined as 1/s. The stride must be power of 2. | |
When there are multiple scales, they must form a pyramid, i.e. they must be | |
a monotically decreasing geometric sequence with a factor of 1/2. | |
sampling_ratio (int): The `sampling_ratio` parameter for the ROIAlign op. | |
pooler_type (string): Name of the type of pooling operation that should be applied. | |
For instance, "ROIPool" or "ROIAlignV2". | |
canonical_box_size (int): A canonical box size in pixels (sqrt(box area)). The default | |
is heuristically defined as 224 pixels in the FPN paper (based on ImageNet | |
pre-training). | |
canonical_level (int): The feature map level index from which a canonically-sized box | |
should be placed. The default is defined as level 4 (stride=16) in the FPN paper, | |
i.e., a box of size 224x224 will be placed on the feature with stride=16. | |
The box placement for all boxes will be determined from their sizes w.r.t | |
canonical_box_size. For example, a box whose area is 4x that of a canonical box | |
should be used to pool features from feature level ``canonical_level+1``. | |
Note that the actual input feature maps given to this module may not have | |
sufficiently many levels for the input boxes. If the boxes are too large or too | |
small for the input feature maps, the closest level will be used. | |
""" | |
super().__init__() | |
if isinstance(output_size, int): | |
output_size = (output_size, output_size) | |
assert len(output_size) == 2 | |
assert isinstance(output_size[0], int) and isinstance(output_size[1], int) | |
self.output_size = output_size | |
if pooler_type == "ROIAlign": | |
self.level_poolers = nn.ModuleList( | |
ROIAlign( | |
output_size, spatial_scale=scale, sampling_ratio=sampling_ratio, aligned=False | |
) | |
for scale in scales | |
) | |
elif pooler_type == "ROIAlignV2": | |
self.level_poolers = nn.ModuleList( | |
ROIAlign( | |
output_size, spatial_scale=scale, sampling_ratio=sampling_ratio, aligned=True | |
) | |
for scale in scales | |
) | |
elif pooler_type == "ROIPool": | |
self.level_poolers = nn.ModuleList( | |
RoIPool(output_size, spatial_scale=scale) for scale in scales | |
) | |
elif pooler_type == "ROIAlignRotated": | |
self.level_poolers = nn.ModuleList( | |
ROIAlignRotated(output_size, spatial_scale=scale, sampling_ratio=sampling_ratio) | |
for scale in scales | |
) | |
else: | |
raise ValueError("Unknown pooler type: {}".format(pooler_type)) | |
# Map scale (defined as 1 / stride) to its feature map level under the | |
# assumption that stride is a power of 2. | |
min_level = -(math.log2(scales[0])) | |
max_level = -(math.log2(scales[-1])) | |
assert math.isclose(min_level, int(min_level)) and math.isclose( | |
max_level, int(max_level) | |
), "Featuremap stride is not power of 2!" | |
self.min_level = int(min_level) | |
self.max_level = int(max_level) | |
assert ( | |
len(scales) == self.max_level - self.min_level + 1 | |
), "[ROIPooler] Sizes of input featuremaps do not form a pyramid!" | |
assert 0 <= self.min_level and self.min_level <= self.max_level | |
self.canonical_level = canonical_level | |
assert canonical_box_size > 0 | |
self.canonical_box_size = canonical_box_size | |
def forward(self, x: List[torch.Tensor], box_lists: List[Boxes]): | |
""" | |
Args: | |
x (list[Tensor]): A list of feature maps of NCHW shape, with scales matching those | |
used to construct this module. | |
box_lists (list[Boxes] | list[RotatedBoxes]): | |
A list of N Boxes or N RotatedBoxes, where N is the number of images in the batch. | |
The box coordinates are defined on the original image and | |
will be scaled by the `scales` argument of :class:`ROIPooler`. | |
Returns: | |
Tensor: | |
A tensor of shape (M, C, output_size, output_size) where M is the total number of | |
boxes aggregated over all N batch images and C is the number of channels in `x`. | |
""" | |
num_level_assignments = len(self.level_poolers) | |
assert isinstance(x, list) and isinstance( | |
box_lists, list | |
), "Arguments to pooler must be lists" | |
assert ( | |
len(x) == num_level_assignments | |
), "unequal value, num_level_assignments={}, but x is list of {} Tensors".format( | |
num_level_assignments, len(x) | |
) | |
assert len(box_lists) == x[0].size( | |
0 | |
), "unequal value, x[0] batch dim 0 is {}, but box_list has length {}".format( | |
x[0].size(0), len(box_lists) | |
) | |
if len(box_lists) == 0: | |
return torch.zeros( | |
(0, x[0].shape[1]) + self.output_size, device=x[0].device, dtype=x[0].dtype | |
) | |
pooler_fmt_boxes = convert_boxes_to_pooler_format(box_lists) | |
if num_level_assignments == 1: | |
return self.level_poolers[0](x[0], pooler_fmt_boxes) | |
level_assignments = assign_boxes_to_levels( | |
box_lists, self.min_level, self.max_level, self.canonical_box_size, self.canonical_level | |
) | |
num_boxes = pooler_fmt_boxes.size(0) | |
num_channels = x[0].shape[1] | |
if len(self.output_size) == 1: | |
output_size = self.output_size[0] | |
dtype, device = x[0].dtype, x[0].device | |
output = torch.zeros( | |
(num_boxes, num_channels, output_size, output_size), dtype=dtype, device=device | |
) | |
else: | |
output_size = self.output_size[0] | |
output_size1 = self.output_size[1] | |
dtype, device = x[0].dtype, x[0].device | |
output = torch.zeros( | |
(num_boxes, num_channels, output_size, output_size1), dtype=dtype, device=device | |
) | |
for level, pooler in enumerate(self.level_poolers): | |
inds = nonzero_tuple(level_assignments == level)[0] | |
pooler_fmt_boxes_level = pooler_fmt_boxes[inds] | |
# Use index_put_ instead of advance indexing, to avoid pytorch/issues/49852 | |
output.index_put_((inds,), pooler(x[level], pooler_fmt_boxes_level)) | |
return output | |