coreml-community/ControlNet-v1-1-Annotators-cpu

1# Copyright (c) Facebook, Inc. and its affiliates.2import math3from typing import List, Optional4import torch5from torch import nn6from torchvision.ops import RoIPool7 8from annotator.oneformer.detectron2.layers import ROIAlign, ROIAlignRotated, cat, nonzero_tuple, shapes_to_tensor9from annotator.oneformer.detectron2.structures import Boxes10from annotator.oneformer.detectron2.utils.tracing import assert_fx_safe, is_fx_tracing11 12"""13To export ROIPooler to torchscript, in this file, variables that should be annotated with14`Union[List[Boxes], List[RotatedBoxes]]` are only annotated with `List[Boxes]`.15 16TODO: Correct these annotations when torchscript support `Union`.17https://github.com/pytorch/pytorch/issues/4141218"""19 20__all__ = ["ROIPooler"]21 22 23def assign_boxes_to_levels(24    box_lists: List[Boxes],25    min_level: int,26    max_level: int,27    canonical_box_size: int,28    canonical_level: int,29):30    """31    Map each box in `box_lists` to a feature map level index and return the assignment32    vector.33 34    Args:35        box_lists (list[Boxes] | list[RotatedBoxes]): A list of N Boxes or N RotatedBoxes,36            where N is the number of images in the batch.37        min_level (int): Smallest feature map level index. The input is considered index 0,38            the output of stage 1 is index 1, and so.39        max_level (int): Largest feature map level index.40        canonical_box_size (int): A canonical box size in pixels (sqrt(box area)).41        canonical_level (int): The feature map level index on which a canonically-sized box42            should be placed.43 44    Returns:45        A tensor of length M, where M is the total number of boxes aggregated over all46            N batch images. The memory layout corresponds to the concatenation of boxes47            from all images. Each element is the feature map index, as an offset from48            `self.min_level`, for the corresponding box (so value i means the box is at49            `self.min_level + i`).50    """51    box_sizes = torch.sqrt(cat([boxes.area() for boxes in box_lists]))52    # Eqn.(1) in FPN paper53    level_assignments = torch.floor(54        canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8)55    )56    # clamp level to (min, max), in case the box size is too large or too small57    # for the available feature maps58    level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)59    return level_assignments.to(torch.int64) - min_level60 61 62# script the module to avoid hardcoded device type63@torch.jit.script_if_tracing64def _convert_boxes_to_pooler_format(boxes: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor:65    sizes = sizes.to(device=boxes.device)66    indices = torch.repeat_interleave(67        torch.arange(len(sizes), dtype=boxes.dtype, device=boxes.device), sizes68    )69    return cat([indices[:, None], boxes], dim=1)70 71 72def convert_boxes_to_pooler_format(box_lists: List[Boxes]):73    """74    Convert all boxes in `box_lists` to the low-level format used by ROI pooling ops75    (see description under Returns).76 77    Args:78        box_lists (list[Boxes] | list[RotatedBoxes]):79            A list of N Boxes or N RotatedBoxes, where N is the number of images in the batch.80 81    Returns:82        When input is list[Boxes]:83            A tensor of shape (M, 5), where M is the total number of boxes aggregated over all84            N batch images.85            The 5 columns are (batch index, x0, y0, x1, y1), where batch index86            is the index in [0, N) identifying which batch image the box with corners at87            (x0, y0, x1, y1) comes from.88        When input is list[RotatedBoxes]:89            A tensor of shape (M, 6), where M is the total number of boxes aggregated over all90            N batch images.91            The 6 columns are (batch index, x_ctr, y_ctr, width, height, angle_degrees),92            where batch index is the index in [0, N) identifying which batch image the93            rotated box (x_ctr, y_ctr, width, height, angle_degrees) comes from.94    """95    boxes = torch.cat([x.tensor for x in box_lists], dim=0)96    # __len__ returns Tensor in tracing.97    sizes = shapes_to_tensor([x.__len__() for x in box_lists])98    return _convert_boxes_to_pooler_format(boxes, sizes)99 100 101@torch.jit.script_if_tracing102def _create_zeros(103    batch_target: Optional[torch.Tensor],104    channels: int,105    height: int,106    width: int,107    like_tensor: torch.Tensor,108) -> torch.Tensor:109    batches = batch_target.shape[0] if batch_target is not None else 0110    sizes = (batches, channels, height, width)111    return torch.zeros(sizes, dtype=like_tensor.dtype, device=like_tensor.device)112 113 114class ROIPooler(nn.Module):115    """116    Region of interest feature map pooler that supports pooling from one or more117    feature maps.118    """119 120    def __init__(121        self,122        output_size,123        scales,124        sampling_ratio,125        pooler_type,126        canonical_box_size=224,127        canonical_level=4,128    ):129        """130        Args:131            output_size (int, tuple[int] or list[int]): output size of the pooled region,132                e.g., 14 x 14. If tuple or list is given, the length must be 2.133            scales (list[float]): The scale for each low-level pooling op relative to134                the input image. For a feature map with stride s relative to the input135                image, scale is defined as 1/s. The stride must be power of 2.136                When there are multiple scales, they must form a pyramid, i.e. they must be137                a monotically decreasing geometric sequence with a factor of 1/2.138            sampling_ratio (int): The `sampling_ratio` parameter for the ROIAlign op.139            pooler_type (string): Name of the type of pooling operation that should be applied.140                For instance, "ROIPool" or "ROIAlignV2".141            canonical_box_size (int): A canonical box size in pixels (sqrt(box area)). The default142                is heuristically defined as 224 pixels in the FPN paper (based on ImageNet143                pre-training).144            canonical_level (int): The feature map level index from which a canonically-sized box145                should be placed. The default is defined as level 4 (stride=16) in the FPN paper,146                i.e., a box of size 224x224 will be placed on the feature with stride=16.147                The box placement for all boxes will be determined from their sizes w.r.t148                canonical_box_size. For example, a box whose area is 4x that of a canonical box149                should be used to pool features from feature level ``canonical_level+1``.150 151                Note that the actual input feature maps given to this module may not have152                sufficiently many levels for the input boxes. If the boxes are too large or too153                small for the input feature maps, the closest level will be used.154        """155        super().__init__()156 157        if isinstance(output_size, int):158            output_size = (output_size, output_size)159        assert len(output_size) == 2160        assert isinstance(output_size[0], int) and isinstance(output_size[1], int)161        self.output_size = output_size162 163        if pooler_type == "ROIAlign":164            self.level_poolers = nn.ModuleList(165                ROIAlign(166                    output_size, spatial_scale=scale, sampling_ratio=sampling_ratio, aligned=False167                )168                for scale in scales169            )170        elif pooler_type == "ROIAlignV2":171            self.level_poolers = nn.ModuleList(172                ROIAlign(173                    output_size, spatial_scale=scale, sampling_ratio=sampling_ratio, aligned=True174                )175                for scale in scales176            )177        elif pooler_type == "ROIPool":178            self.level_poolers = nn.ModuleList(179                RoIPool(output_size, spatial_scale=scale) for scale in scales180            )181        elif pooler_type == "ROIAlignRotated":182            self.level_poolers = nn.ModuleList(183                ROIAlignRotated(output_size, spatial_scale=scale, sampling_ratio=sampling_ratio)184                for scale in scales185            )186        else:187            raise ValueError("Unknown pooler type: {}".format(pooler_type))188 189        # Map scale (defined as 1 / stride) to its feature map level under the190        # assumption that stride is a power of 2.191        min_level = -(math.log2(scales[0]))192        max_level = -(math.log2(scales[-1]))193        assert math.isclose(min_level, int(min_level)) and math.isclose(194            max_level, int(max_level)195        ), "Featuremap stride is not power of 2!"196        self.min_level = int(min_level)197        self.max_level = int(max_level)198        assert (199            len(scales) == self.max_level - self.min_level + 1200        ), "[ROIPooler] Sizes of input featuremaps do not form a pyramid!"201        assert 0 <= self.min_level and self.min_level <= self.max_level202        self.canonical_level = canonical_level203        assert canonical_box_size > 0204        self.canonical_box_size = canonical_box_size205 206    def forward(self, x: List[torch.Tensor], box_lists: List[Boxes]):207        """208        Args:209            x (list[Tensor]): A list of feature maps of NCHW shape, with scales matching those210                used to construct this module.211            box_lists (list[Boxes] | list[RotatedBoxes]):212                A list of N Boxes or N RotatedBoxes, where N is the number of images in the batch.213                The box coordinates are defined on the original image and214                will be scaled by the `scales` argument of :class:`ROIPooler`.215 216        Returns:217            Tensor:218                A tensor of shape (M, C, output_size, output_size) where M is the total number of219                boxes aggregated over all N batch images and C is the number of channels in `x`.220        """221        num_level_assignments = len(self.level_poolers)222 223        if not is_fx_tracing():224            torch._assert(225                isinstance(x, list) and isinstance(box_lists, list),226                "Arguments to pooler must be lists",227            )228        assert_fx_safe(229            len(x) == num_level_assignments,230            "unequal value, num_level_assignments={}, but x is list of {} Tensors".format(231                num_level_assignments, len(x)232            ),233        )234        assert_fx_safe(235            len(box_lists) == x[0].size(0),236            "unequal value, x[0] batch dim 0 is {}, but box_list has length {}".format(237                x[0].size(0), len(box_lists)238            ),239        )240        if len(box_lists) == 0:241            return _create_zeros(None, x[0].shape[1], *self.output_size, x[0])242 243        pooler_fmt_boxes = convert_boxes_to_pooler_format(box_lists)244 245        if num_level_assignments == 1:246            return self.level_poolers[0](x[0], pooler_fmt_boxes)247 248        level_assignments = assign_boxes_to_levels(249            box_lists, self.min_level, self.max_level, self.canonical_box_size, self.canonical_level250        )251 252        num_channels = x[0].shape[1]253        output_size = self.output_size[0]254 255        output = _create_zeros(pooler_fmt_boxes, num_channels, output_size, output_size, x[0])256 257        for level, pooler in enumerate(self.level_poolers):258            inds = nonzero_tuple(level_assignments == level)[0]259            pooler_fmt_boxes_level = pooler_fmt_boxes[inds]260            # Use index_put_ instead of advance indexing, to avoid pytorch/issues/49852261            output.index_put_((inds,), pooler(x[level], pooler_fmt_boxes_level))262 263        return output264