detectron2/projects/TensorMask/tensormask/arch.py

# Copyright (c) Facebook, Inc. and its affiliates.
import copy
import math
from typing import List
import torch
import torch.nn.functional as F
from fvcore.nn import sigmoid_focal_loss_star_jit, smooth_l1_loss
from torch import nn

from detectron2.layers import ShapeSpec, batched_nms, cat, paste_masks_in_image
from detectron2.modeling.anchor_generator import DefaultAnchorGenerator
from detectron2.modeling.backbone import build_backbone
from detectron2.modeling.box_regression import Box2BoxTransform
from detectron2.modeling.meta_arch.build import META_ARCH_REGISTRY
from detectron2.modeling.meta_arch.retinanet import permute_to_N_HWA_K
from detectron2.structures import Boxes, ImageList, Instances

from tensormask.layers import SwapAlign2Nat

__all__ = ["TensorMask"]


def permute_all_cls_and_box_to_N_HWA_K_and_concat(pred_logits, pred_anchor_deltas, num_classes=80):
    """
    Rearrange the tensor layout from the network output, i.e.:
    list[Tensor]: #lvl tensors of shape (N, A x K, Hi, Wi)
    to per-image predictions, i.e.:
    Tensor: of shape (N x sum(Hi x Wi x A), K)
    """
    # for each feature level, permute the outputs to make them be in the
    # same format as the labels.
    pred_logits_flattened = [permute_to_N_HWA_K(x, num_classes) for x in pred_logits]
    pred_anchor_deltas_flattened = [permute_to_N_HWA_K(x, 4) for x in pred_anchor_deltas]
    # concatenate on the first dimension (representing the feature levels), to
    # take into account the way the labels were generated (with all feature maps
    # being concatenated as well)
    pred_logits = cat(pred_logits_flattened, dim=1).view(-1, num_classes)
    pred_anchor_deltas = cat(pred_anchor_deltas_flattened, dim=1).view(-1, 4)
    return pred_logits, pred_anchor_deltas


def _assignment_rule(
    gt_boxes,
    anchor_boxes,
    unit_lengths,
    min_anchor_size,
    scale_thresh=2.0,
    spatial_thresh=1.0,
    uniqueness_on=True,
):
    """
    Given two lists of boxes of N ground truth boxes and M anchor boxes,
    compute the assignment between the two, following the assignment rules in
    https://arxiv.org/abs/1903.12174.
    The box order must be (xmin, ymin, xmax, ymax), so please make sure to convert
    to BoxMode.XYXY_ABS before calling this function.

    Args:
        gt_boxes, anchor_boxes (Boxes): two Boxes. Contains N & M boxes/anchors, respectively.
        unit_lengths (Tensor): Contains the unit lengths of M anchor boxes.
        min_anchor_size (float): Minimum size of the anchor, in pixels
        scale_thresh (float): The `scale` threshold: the maximum size of the anchor
                              should not be greater than scale_thresh x max(h, w) of
                              the ground truth box.
        spatial_thresh (float): The `spatial` threshold: the l2 distance between the
                              center of the anchor and the ground truth box should not
                              be greater than spatial_thresh x u where u is the unit length.

    Returns:
        matches (Tensor[int64]): a vector of length M, where matches[i] is a matched
                ground-truth index in [0, N)
        match_labels (Tensor[int8]): a vector of length M, where pred_labels[i] indicates
            whether a prediction is a true or false positive or ignored
    """
    gt_boxes, anchor_boxes = gt_boxes.tensor, anchor_boxes.tensor
    N = gt_boxes.shape[0]
    M = anchor_boxes.shape[0]
    if N == 0 or M == 0:
        return (
            gt_boxes.new_full((N,), 0, dtype=torch.int64),
            gt_boxes.new_full((N,), -1, dtype=torch.int8),
        )

    # Containment rule
    lt = torch.min(gt_boxes[:, None, :2], anchor_boxes[:, :2])  # [N,M,2]
    rb = torch.max(gt_boxes[:, None, 2:], anchor_boxes[:, 2:])  # [N,M,2]
    union = cat([lt, rb], dim=2)  # [N,M,4]

    dummy_gt_boxes = torch.zeros_like(gt_boxes)
    anchor = dummy_gt_boxes[:, None, :] + anchor_boxes[:, :]  # [N,M,4]

    contain_matrix = torch.all(union == anchor, dim=2)  # [N,M]

    # Centrality rule, scale
    gt_size_lower = torch.max(gt_boxes[:, 2:] - gt_boxes[:, :2], dim=1)[0]  # [N]
    gt_size_upper = gt_size_lower * scale_thresh  # [N]
    # Fall back for small objects
    gt_size_upper[gt_size_upper < min_anchor_size] = min_anchor_size
    # Due to sampling of locations, the anchor sizes are deducted with sampling strides
    anchor_size = (
        torch.max(anchor_boxes[:, 2:] - anchor_boxes[:, :2], dim=1)[0] - unit_lengths
    )  # [M]

    size_diff_upper = gt_size_upper[:, None] - anchor_size  # [N,M]
    scale_matrix = size_diff_upper >= 0  # [N,M]

    # Centrality rule, spatial
    gt_center = (gt_boxes[:, 2:] + gt_boxes[:, :2]) / 2  # [N,2]
    anchor_center = (anchor_boxes[:, 2:] + anchor_boxes[:, :2]) / 2  # [M,2]
    offset_center = gt_center[:, None, :] - anchor_center[:, :]  # [N,M,2]
    offset_center /= unit_lengths[:, None]  # [N,M,2]
    spatial_square = spatial_thresh * spatial_thresh
    spatial_matrix = torch.sum(offset_center * offset_center, dim=2) <= spatial_square

    assign_matrix = (contain_matrix & scale_matrix & spatial_matrix).int()

    # assign_matrix is N (gt) x M (predicted)
    # Max over gt elements (dim 0) to find best gt candidate for each prediction
    matched_vals, matches = assign_matrix.max(dim=0)
    match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)

    match_labels[matched_vals == 0] = 0
    match_labels[matched_vals == 1] = 1

    # find all the elements that match to ground truths multiple times
    not_unique_idxs = assign_matrix.sum(dim=0) > 1
    if uniqueness_on:
        match_labels[not_unique_idxs] = 0
    else:
        match_labels[not_unique_idxs] = -1

    return matches, match_labels


# TODO make the paste_mask function in d2 core support mask list
def _paste_mask_lists_in_image(masks, boxes, image_shape, threshold=0.5):
    """
    Paste a list of masks that are of various resolutions (e.g., 28 x 28) into an image.
    The location, height, and width for pasting each mask is determined by their
    corresponding bounding boxes in boxes.

    Args:
        masks (list(Tensor)): A list of Tensor of shape (1, Hmask_i, Wmask_i).
                            Values are in [0, 1]. The list length, Bimg, is the
                            number of detected object instances in the image.
        boxes (Boxes): A Boxes of length Bimg. boxes.tensor[i] and masks[i] correspond
                            to the same object instance.
        image_shape (tuple): height, width
        threshold (float): A threshold in [0, 1] for converting the (soft) masks to
            binary masks.

    Returns:
        img_masks (Tensor): A tensor of shape (Bimg, Himage, Wimage), where Bimg is the
        number of detected object instances and Himage, Wimage are the image width
        and height. img_masks[i] is a binary mask for object instance i.
    """
    if len(masks) == 0:
        return torch.empty((0, 1) + image_shape, dtype=torch.uint8)

    # Loop over masks groups. Each group has the same mask prediction size.
    img_masks = []
    ind_masks = []
    mask_sizes = torch.tensor([m.shape[-1] for m in masks])
    unique_sizes = torch.unique(mask_sizes)
    for msize in unique_sizes.tolist():
        cur_ind = torch.where(mask_sizes == msize)[0]
        ind_masks.append(cur_ind)

        cur_masks = cat([masks[i] for i in cur_ind])
        cur_boxes = boxes[cur_ind]
        img_masks.append(paste_masks_in_image(cur_masks, cur_boxes, image_shape, threshold))

    img_masks = cat(img_masks)
    ind_masks = cat(ind_masks)

    img_masks_out = torch.empty_like(img_masks)
    img_masks_out[ind_masks, :, :] = img_masks

    return img_masks_out


def _postprocess(results, result_mask_info, output_height, output_width, mask_threshold=0.5):
    """
    Post-process the output boxes for TensorMask.
    The input images are often resized when entering an object detector.
    As a result, we often need the outputs of the detector in a different
    resolution from its inputs.

    This function will postprocess the raw outputs of TensorMask
    to produce outputs according to the desired output resolution.

    Args:
        results (Instances): the raw outputs from the detector.
            `results.image_size` contains the input image resolution the detector sees.
            This object might be modified in-place. Note that it does not contain the field
            `pred_masks`, which is provided by another input `result_masks`.
        result_mask_info (list[Tensor], Boxes): a pair of two items for mask related results.
                The first item is a list of #detection tensors, each is the predicted masks.
                The second item is the anchors corresponding to the predicted masks.
        output_height, output_width: the desired output resolution.

    Returns:
        Instances: the postprocessed output from the model, based on the output resolution
    """
    scale_x, scale_y = (output_width / results.image_size[1], output_height / results.image_size[0])
    results = Instances((output_height, output_width), **results.get_fields())

    output_boxes = results.pred_boxes
    output_boxes.tensor[:, 0::2] *= scale_x
    output_boxes.tensor[:, 1::2] *= scale_y
    output_boxes.clip(results.image_size)

    inds_nonempty = output_boxes.nonempty()
    results = results[inds_nonempty]
    result_masks, result_anchors = result_mask_info
    if result_masks:
        result_anchors.tensor[:, 0::2] *= scale_x
        result_anchors.tensor[:, 1::2] *= scale_y
        result_masks = [x for (i, x) in zip(inds_nonempty.tolist(), result_masks) if i]
        results.pred_masks = _paste_mask_lists_in_image(
            result_masks,
            result_anchors[inds_nonempty],
            results.image_size,
            threshold=mask_threshold,
        )
    return results


class TensorMaskAnchorGenerator(DefaultAnchorGenerator):
    """
    For a set of image sizes and feature maps, computes a set of anchors for TensorMask.
    It also computes the unit lengths and indexes for each anchor box.
    """

    def grid_anchors_with_unit_lengths_and_indexes(self, grid_sizes):
        anchors = []
        unit_lengths = []
        indexes = []
        for lvl, (size, stride, base_anchors) in enumerate(
            zip(grid_sizes, self.strides, self.cell_anchors)
        ):
            grid_height, grid_width = size
            device = base_anchors.device
            shifts_x = torch.arange(
                0, grid_width * stride, step=stride, dtype=torch.float32, device=device
            )
            shifts_y = torch.arange(
                0, grid_height * stride, step=stride, dtype=torch.float32, device=device
            )
            shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
            shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=2)
            # Stack anchors in shapes of (HWA, 4)
            cur_anchor = (shifts[:, :, None, :] + base_anchors.view(1, 1, -1, 4)).view(-1, 4)
            anchors.append(cur_anchor)
            unit_lengths.append(
                torch.full((cur_anchor.shape[0],), stride, dtype=torch.float32, device=device)
            )
            # create mask indexes using mesh grid
            shifts_l = torch.full((1,), lvl, dtype=torch.int64, device=device)
            shifts_i = torch.zeros((1,), dtype=torch.int64, device=device)
            shifts_h = torch.arange(0, grid_height, dtype=torch.int64, device=device)
            shifts_w = torch.arange(0, grid_width, dtype=torch.int64, device=device)
            shifts_a = torch.arange(0, base_anchors.shape[0], dtype=torch.int64, device=device)
            grids = torch.meshgrid(shifts_l, shifts_i, shifts_h, shifts_w, shifts_a)

            indexes.append(torch.stack(grids, dim=5).view(-1, 5))

        return anchors, unit_lengths, indexes

    def forward(self, features):
        """
        Returns:
            list[list[Boxes]]: a list of #image elements. Each is a list of #feature level Boxes.
                The Boxes contains anchors of this image on the specific feature level.
            list[list[Tensor]]: a list of #image elements. Each is a list of #feature level tensors.
                The tensor contains strides, or unit lengths for the anchors.
            list[list[Tensor]]: a list of #image elements. Each is a list of #feature level tensors.
                The Tensor contains indexes for the anchors, with the last dimension meaning
                (L, N, H, W, A), where L is level, I is image (not set yet), H is height,
                W is width, and A is anchor.
        """
        num_images = len(features[0])
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_list, lengths_list, indexes_list = self.grid_anchors_with_unit_lengths_and_indexes(
            grid_sizes
        )

        # Convert anchors from Tensor to Boxes
        anchors_per_im = [Boxes(x) for x in anchors_list]

        # TODO it can be simplified to not return duplicated information for
        # each image, just like detectron2's own AnchorGenerator
        anchors = [copy.deepcopy(anchors_per_im) for _ in range(num_images)]
        unit_lengths = [copy.deepcopy(lengths_list) for _ in range(num_images)]
        indexes = [copy.deepcopy(indexes_list) for _ in range(num_images)]

        return anchors, unit_lengths, indexes


@META_ARCH_REGISTRY.register()
class TensorMask(nn.Module):
    """
    TensorMask model. Creates FPN backbone, anchors and a head for classification
    and box regression. Calculates and applies proper losses to class, box, and
    masks.
    """

    def __init__(self, cfg):
        super().__init__()

        # fmt: off
        self.num_classes              = cfg.MODEL.TENSOR_MASK.NUM_CLASSES
        self.in_features              = cfg.MODEL.TENSOR_MASK.IN_FEATURES
        self.anchor_sizes             = cfg.MODEL.ANCHOR_GENERATOR.SIZES
        self.num_levels               = len(cfg.MODEL.ANCHOR_GENERATOR.SIZES)
        # Loss parameters:
        self.focal_loss_alpha         = cfg.MODEL.TENSOR_MASK.FOCAL_LOSS_ALPHA
        self.focal_loss_gamma         = cfg.MODEL.TENSOR_MASK.FOCAL_LOSS_GAMMA
        # Inference parameters:
        self.score_threshold          = cfg.MODEL.TENSOR_MASK.SCORE_THRESH_TEST
        self.topk_candidates          = cfg.MODEL.TENSOR_MASK.TOPK_CANDIDATES_TEST
        self.nms_threshold            = cfg.MODEL.TENSOR_MASK.NMS_THRESH_TEST
        self.detections_im            = cfg.TEST.DETECTIONS_PER_IMAGE
        # Mask parameters:
        self.mask_on                  = cfg.MODEL.MASK_ON
        self.mask_loss_weight         = cfg.MODEL.TENSOR_MASK.MASK_LOSS_WEIGHT
        self.mask_pos_weight          = torch.tensor(cfg.MODEL.TENSOR_MASK.POSITIVE_WEIGHT,
                                                     dtype=torch.float32)
        self.bipyramid_on             = cfg.MODEL.TENSOR_MASK.BIPYRAMID_ON
        # fmt: on

        # build the backbone
        self.backbone = build_backbone(cfg)

        backbone_shape = self.backbone.output_shape()
        feature_shapes = [backbone_shape[f] for f in self.in_features]
        feature_strides = [x.stride for x in feature_shapes]
        # build anchors
        self.anchor_generator = TensorMaskAnchorGenerator(cfg, feature_shapes)
        self.num_anchors = self.anchor_generator.num_cell_anchors[0]
        anchors_min_level = cfg.MODEL.ANCHOR_GENERATOR.SIZES[0]
        self.mask_sizes = [size // feature_strides[0] for size in anchors_min_level]
        self.min_anchor_size = min(anchors_min_level) - feature_strides[0]

        # head of the TensorMask
        self.head = TensorMaskHead(
            cfg, self.num_levels, self.num_anchors, self.mask_sizes, feature_shapes
        )
        # box transform
        self.box2box_transform = Box2BoxTransform(weights=cfg.MODEL.TENSOR_MASK.BBOX_REG_WEIGHTS)
        self.register_buffer("pixel_mean", torch.tensor(cfg.MODEL.PIXEL_MEAN).view(-1, 1, 1), False)
        self.register_buffer("pixel_std", torch.tensor(cfg.MODEL.PIXEL_STD).view(-1, 1, 1), False)

    @property
    def device(self):
        return self.pixel_mean.device

    def forward(self, batched_inputs):
        """
        Args:
            batched_inputs: a list, batched outputs of :class:`DetectionTransform` .
                Each item in the list contains the inputs for one image.
            For now, each item in the list is a dict that contains:
                image: Tensor, image in (C, H, W) format.
                instances: Instances
                Other information that's included in the original dicts, such as:
                    "height", "width" (int): the output resolution of the model, used in inference.
                        See :meth:`postprocess` for details.
         Returns:
            losses (dict[str: Tensor]): mapping from a named loss to a tensor
                storing the loss. Used during training only.
        """
        images = self.preprocess_image(batched_inputs)
        if "instances" in batched_inputs[0]:
            gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
        else:
            gt_instances = None

        features = self.backbone(images.tensor)
        features = [features[f] for f in self.in_features]
        # apply the TensorMask head
        pred_logits, pred_deltas, pred_masks = self.head(features)
        # generate anchors based on features, is it image specific?
        anchors, unit_lengths, indexes = self.anchor_generator(features)

        if self.training:
            # get ground truths for class labels and box targets, it will label each anchor
            gt_class_info, gt_delta_info, gt_mask_info, num_fg = self.get_ground_truth(
                anchors, unit_lengths, indexes, gt_instances
            )
            # compute the loss
            return self.losses(
                gt_class_info,
                gt_delta_info,
                gt_mask_info,
                num_fg,
                pred_logits,
                pred_deltas,
                pred_masks,
            )
        else:
            # do inference to get the output
            results = self.inference(pred_logits, pred_deltas, pred_masks, anchors, indexes, images)
            processed_results = []
            for results_im, input_im, image_size in zip(
                results, batched_inputs, images.image_sizes
            ):
                height = input_im.get("height", image_size[0])
                width = input_im.get("width", image_size[1])
                # this is to do post-processing with the image size
                result_box, result_mask = results_im
                r = _postprocess(result_box, result_mask, height, width)
                processed_results.append({"instances": r})
            return processed_results

    def losses(
        self,
        gt_class_info,
        gt_delta_info,
        gt_mask_info,
        num_fg,
        pred_logits,
        pred_deltas,
        pred_masks,
    ):
        """
        Args:
            For `gt_class_info`, `gt_delta_info`, `gt_mask_info` and `num_fg` parameters, see
                :meth:`TensorMask.get_ground_truth`.
            For `pred_logits`, `pred_deltas` and `pred_masks`, see
                :meth:`TensorMaskHead.forward`.

        Returns:
            losses (dict[str: Tensor]): mapping from a named loss to a scalar tensor
                storing the loss. Used during training only. The potential dict keys are:
                "loss_cls", "loss_box_reg" and "loss_mask".
        """
        gt_classes_target, gt_valid_inds = gt_class_info
        gt_deltas, gt_fg_inds = gt_delta_info
        gt_masks, gt_mask_inds = gt_mask_info
        loss_normalizer = torch.tensor(max(1, num_fg), dtype=torch.float32, device=self.device)

        # classification and regression
        pred_logits, pred_deltas = permute_all_cls_and_box_to_N_HWA_K_and_concat(
            pred_logits, pred_deltas, self.num_classes
        )
        loss_cls = (
            sigmoid_focal_loss_star_jit(
                pred_logits[gt_valid_inds],
                gt_classes_target[gt_valid_inds],
                alpha=self.focal_loss_alpha,
                gamma=self.focal_loss_gamma,
                reduction="sum",
            )
            / loss_normalizer
        )

        if num_fg == 0:
            loss_box_reg = pred_deltas.sum() * 0
        else:
            loss_box_reg = (
                smooth_l1_loss(pred_deltas[gt_fg_inds], gt_deltas, beta=0.0, reduction="sum")
                / loss_normalizer
            )
        losses = {"loss_cls": loss_cls, "loss_box_reg": loss_box_reg}

        # mask prediction
        if self.mask_on:
            loss_mask = 0
            for lvl in range(self.num_levels):
                cur_level_factor = 2 ** lvl if self.bipyramid_on else 1
                for anc in range(self.num_anchors):
                    cur_gt_mask_inds = gt_mask_inds[lvl][anc]
                    if cur_gt_mask_inds is None:
                        loss_mask += pred_masks[lvl][anc][0, 0, 0, 0] * 0
                    else:
                        cur_mask_size = self.mask_sizes[anc] * cur_level_factor
                        # TODO maybe there are numerical issues when mask sizes are large
                        cur_size_divider = torch.tensor(
                            self.mask_loss_weight / (cur_mask_size ** 2),
                            dtype=torch.float32,
                            device=self.device,
                        )

                        cur_pred_masks = pred_masks[lvl][anc][
                            cur_gt_mask_inds[:, 0],  # N
                            :,  # V x U
                            cur_gt_mask_inds[:, 1],  # H
                            cur_gt_mask_inds[:, 2],  # W
                        ]

                        loss_mask += F.binary_cross_entropy_with_logits(
                            cur_pred_masks.view(-1, cur_mask_size, cur_mask_size),  # V, U
                            gt_masks[lvl][anc].to(dtype=torch.float32),
                            reduction="sum",
                            weight=cur_size_divider,
                            pos_weight=self.mask_pos_weight,
                        )
            losses["loss_mask"] = loss_mask / loss_normalizer
        return losses

    @torch.no_grad()
    def get_ground_truth(self, anchors, unit_lengths, indexes, targets):
        """
        Args:
            anchors (list[list[Boxes]]): a list of N=#image elements. Each is a
                list of #feature level Boxes. The Boxes contains anchors of
                this image on the specific feature level.
            unit_lengths (list[list[Tensor]]): a list of N=#image elements. Each is a
                list of #feature level Tensor. The tensor contains unit lengths for anchors of
                this image on the specific feature level.
            indexes (list[list[Tensor]]): a list of N=#image elements. Each is a
                list of #feature level Tensor. The tensor contains the 5D index of
                each anchor, the second dimension means (L, N, H, W, A), where L
                is level, I is image, H is height, W is width, and A is anchor.
            targets (list[Instances]): a list of N `Instances`s. The i-th
                `Instances` contains the ground-truth per-instance annotations
                for the i-th input image.  Specify `targets` during training only.

        Returns:
            gt_class_info (Tensor, Tensor): A pair of two tensors for classification.
                The first one is an integer tensor of shape (R, #classes) storing ground-truth
                labels for each anchor. R is the total number of anchors in the batch.
                The second one is an integer tensor of shape (R,), to indicate which
                anchors are valid for loss computation, which anchors are not.
            gt_delta_info (Tensor, Tensor): A pair of two tensors for boxes.
                The first one, of shape (F, 4). F=#foreground anchors.
                The last dimension represents ground-truth box2box transform
                targets (dx, dy, dw, dh) that map each anchor to its matched ground-truth box.
                Only foreground anchors have values in this tensor. Could be `None` if F=0.
                The second one, of shape (R,), is an integer tensor indicating which anchors
                are foreground ones used for box regression. Could be `None` if F=0.
            gt_mask_info (list[list[Tensor]], list[list[Tensor]]): A pair of two lists for masks.
                The first one is a list of P=#feature level elements. Each is a
                list of A=#anchor tensors. Each tensor contains the ground truth
                masks of the same size and for the same feature level. Could be `None`.
                The second one is a list of P=#feature level elements. Each is a
                list of A=#anchor tensors. Each tensor contains the location of the ground truth
                masks of the same size and for the same feature level. The second dimension means
                (N, H, W), where N is image, H is height, and W is width. Could be `None`.
            num_fg (int): F=#foreground anchors, used later for loss normalization.
        """
        gt_classes = []
        gt_deltas = []
        gt_masks = [[[] for _ in range(self.num_anchors)] for _ in range(self.num_levels)]
        gt_mask_inds = [[[] for _ in range(self.num_anchors)] for _ in range(self.num_levels)]

        anchors = [Boxes.cat(anchors_i) for anchors_i in anchors]
        unit_lengths = [cat(unit_lengths_i) for unit_lengths_i in unit_lengths]
        indexes = [cat(indexes_i) for indexes_i in indexes]

        num_fg = 0
        for i, (anchors_im, unit_lengths_im, indexes_im, targets_im) in enumerate(
            zip(anchors, unit_lengths, indexes, targets)
        ):
            # Initialize all
            gt_classes_i = torch.full_like(
                unit_lengths_im, self.num_classes, dtype=torch.int64, device=self.device
            )
            # Ground truth classes
            has_gt = len(targets_im) > 0
            if has_gt:
                # Compute the pairwise matrix
                gt_matched_inds, anchor_labels = _assignment_rule(
                    targets_im.gt_boxes, anchors_im, unit_lengths_im, self.min_anchor_size
                )
                # Find the foreground instances
                fg_inds = anchor_labels == 1
                fg_anchors = anchors_im[fg_inds]
                num_fg += len(fg_anchors)
                # Find the ground truths for foreground instances
                gt_fg_matched_inds = gt_matched_inds[fg_inds]
                # Assign labels for foreground instances
                gt_classes_i[fg_inds] = targets_im.gt_classes[gt_fg_matched_inds]
                # Anchors with label -1 are ignored, others are left as negative
                gt_classes_i[anchor_labels == -1] = -1

                # Boxes
                # Ground truth box regression, only for foregrounds
                matched_gt_boxes = targets_im[gt_fg_matched_inds].gt_boxes
                # Compute box regression offsets for foregrounds only
                gt_deltas_i = self.box2box_transform.get_deltas(
                    fg_anchors.tensor, matched_gt_boxes.tensor
                )
                gt_deltas.append(gt_deltas_i)

                # Masks
                if self.mask_on:
                    # Compute masks for each level and each anchor
                    matched_indexes = indexes_im[fg_inds, :]
                    for lvl in range(self.num_levels):
                        ids_lvl = matched_indexes[:, 0] == lvl
                        if torch.any(ids_lvl):
                            cur_level_factor = 2 ** lvl if self.bipyramid_on else 1
                            for anc in range(self.num_anchors):
                                ids_lvl_anchor = ids_lvl & (matched_indexes[:, 4] == anc)
                                if torch.any(ids_lvl_anchor):
                                    gt_masks[lvl][anc].append(
                                        targets_im[
                                            gt_fg_matched_inds[ids_lvl_anchor]
                                        ].gt_masks.crop_and_resize(
                                            fg_anchors[ids_lvl_anchor].tensor,
                                            self.mask_sizes[anc] * cur_level_factor,
                                        )
                                    )
                                    # Select (N, H, W) dimensions
                                    gt_mask_inds_lvl_anc = matched_indexes[ids_lvl_anchor, 1:4]
                                    # Set the image index to the current image
                                    gt_mask_inds_lvl_anc[:, 0] = i
                                    gt_mask_inds[lvl][anc].append(gt_mask_inds_lvl_anc)
            gt_classes.append(gt_classes_i)

        # Classes and boxes
        gt_classes = cat(gt_classes)
        gt_valid_inds = gt_classes >= 0
        gt_fg_inds = gt_valid_inds & (gt_classes < self.num_classes)
        gt_classes_target = torch.zeros(
            (gt_classes.shape[0], self.num_classes), dtype=torch.float32, device=self.device
        )
        gt_classes_target[gt_fg_inds, gt_classes[gt_fg_inds]] = 1
        gt_deltas = cat(gt_deltas) if gt_deltas else None

        # Masks
        gt_masks = [[cat(mla) if mla else None for mla in ml] for ml in gt_masks]
        gt_mask_inds = [[cat(ila) if ila else None for ila in il] for il in gt_mask_inds]
        return (
            (gt_classes_target, gt_valid_inds),
            (gt_deltas, gt_fg_inds),
            (gt_masks, gt_mask_inds),
            num_fg,
        )

    def inference(self, pred_logits, pred_deltas, pred_masks, anchors, indexes, images):
        """
        Arguments:
            pred_logits, pred_deltas, pred_masks: Same as the output of:
                meth:`TensorMaskHead.forward`
            anchors, indexes: Same as the input of meth:`TensorMask.get_ground_truth`
            images (ImageList): the input images

        Returns:
            results (List[Instances]): a list of #images elements.
        """
        assert len(anchors) == len(images)
        results = []

        pred_logits = [permute_to_N_HWA_K(x, self.num_classes) for x in pred_logits]
        pred_deltas = [permute_to_N_HWA_K(x, 4) for x in pred_deltas]

        pred_logits = cat(pred_logits, dim=1)
        pred_deltas = cat(pred_deltas, dim=1)

        for img_idx, (anchors_im, indexes_im) in enumerate(zip(anchors, indexes)):
            # Get the size of the current image
            image_size = images.image_sizes[img_idx]

            logits_im = pred_logits[img_idx]
            deltas_im = pred_deltas[img_idx]

            if self.mask_on:
                masks_im = [[mla[img_idx] for mla in ml] for ml in pred_masks]
            else:
                masks_im = [None] * self.num_levels
            results_im = self.inference_single_image(
                logits_im,
                deltas_im,
                masks_im,
                Boxes.cat(anchors_im),
                cat(indexes_im),
                tuple(image_size),
            )
            results.append(results_im)
        return results

    def inference_single_image(
        self, pred_logits, pred_deltas, pred_masks, anchors, indexes, image_size
    ):
        """
        Single-image inference. Return bounding-box detection results by thresholding
        on scores and applying non-maximum suppression (NMS).

        Arguments:
            pred_logits (list[Tensor]): list of #feature levels. Each entry contains
                tensor of size (AxHxW, K)
            pred_deltas (list[Tensor]): Same shape as 'pred_logits' except that K becomes 4.
            pred_masks (list[list[Tensor]]): List of #feature levels, each is a list of #anchors.
                Each entry contains tensor of size (M_i*M_i, H, W). `None` if mask_on=False.
            anchors (list[Boxes]): list of #feature levels. Each entry contains
                a Boxes object, which contains all the anchors for that
                image in that feature level.
            image_size (tuple(H, W)): a tuple of the image height and width.

        Returns:
            Same as `inference`, but for only one image.
        """
        pred_logits = pred_logits.flatten().sigmoid_()
        # We get top locations across all levels to accelerate the inference speed,
        # which does not seem to affect the accuracy.
        # First select values above the threshold
        logits_top_idxs = torch.where(pred_logits > self.score_threshold)[0]
        # Then get the top values
        num_topk = min(self.topk_candidates, logits_top_idxs.shape[0])
        pred_prob, topk_idxs = pred_logits[logits_top_idxs].sort(descending=True)
        # Keep top k scoring values
        pred_prob = pred_prob[:num_topk]
        # Keep top k values
        top_idxs = logits_top_idxs[topk_idxs[:num_topk]]

        # class index
        cls_idxs = top_idxs % self.num_classes
        # HWA index
        top_idxs //= self.num_classes
        # predict boxes
        pred_boxes = self.box2box_transform.apply_deltas(
            pred_deltas[top_idxs], anchors[top_idxs].tensor
        )
        # apply nms
        keep = batched_nms(pred_boxes, pred_prob, cls_idxs, self.nms_threshold)
        # pick the top ones
        keep = keep[: self.detections_im]

        results = Instances(image_size)
        results.pred_boxes = Boxes(pred_boxes[keep])
        results.scores = pred_prob[keep]
        results.pred_classes = cls_idxs[keep]

        # deal with masks
        result_masks, result_anchors = [], None
        if self.mask_on:
            # index and anchors, useful for masks
            top_indexes = indexes[top_idxs]
            top_anchors = anchors[top_idxs]
            result_indexes = top_indexes[keep]
            result_anchors = top_anchors[keep]
            # Get masks and do sigmoid
            for lvl, _, h, w, anc in result_indexes.tolist():
                cur_size = self.mask_sizes[anc] * (2 ** lvl if self.bipyramid_on else 1)
                result_masks.append(
                    torch.sigmoid(pred_masks[lvl][anc][:, h, w].view(1, cur_size, cur_size))
                )

        return results, (result_masks, result_anchors)

    def preprocess_image(self, batched_inputs):
        """
        Normalize, pad and batch the input images.
        """
        images = [x["image"].to(self.device) for x in batched_inputs]
        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
        images = ImageList.from_tensors(images, self.backbone.size_divisibility)
        return images


class TensorMaskHead(nn.Module):
    def __init__(self, cfg, num_levels, num_anchors, mask_sizes, input_shape: List[ShapeSpec]):
        """
        TensorMask head.
        """
        super().__init__()
        # fmt: off
        self.in_features        = cfg.MODEL.TENSOR_MASK.IN_FEATURES
        in_channels             = input_shape[0].channels
        num_classes             = cfg.MODEL.TENSOR_MASK.NUM_CLASSES
        cls_channels            = cfg.MODEL.TENSOR_MASK.CLS_CHANNELS
        num_convs               = cfg.MODEL.TENSOR_MASK.NUM_CONVS
        # box parameters
        bbox_channels           = cfg.MODEL.TENSOR_MASK.BBOX_CHANNELS
        # mask parameters
        self.mask_on            = cfg.MODEL.MASK_ON
        self.mask_sizes         = mask_sizes
        mask_channels           = cfg.MODEL.TENSOR_MASK.MASK_CHANNELS
        self.align_on           = cfg.MODEL.TENSOR_MASK.ALIGNED_ON
        self.bipyramid_on       = cfg.MODEL.TENSOR_MASK.BIPYRAMID_ON
        # fmt: on

        # class subnet
        cls_subnet = []
        cur_channels = in_channels
        for _ in range(num_convs):
            cls_subnet.append(
                nn.Conv2d(cur_channels, cls_channels, kernel_size=3, stride=1, padding=1)
            )
            cur_channels = cls_channels
            cls_subnet.append(nn.ReLU())

        self.cls_subnet = nn.Sequential(*cls_subnet)
        self.cls_score = nn.Conv2d(
            cur_channels, num_anchors * num_classes, kernel_size=3, stride=1, padding=1
        )
        modules_list = [self.cls_subnet, self.cls_score]

        # box subnet
        bbox_subnet = []
        cur_channels = in_channels
        for _ in range(num_convs):
            bbox_subnet.append(
                nn.Conv2d(cur_channels, bbox_channels, kernel_size=3, stride=1, padding=1)
            )
            cur_channels = bbox_channels
            bbox_subnet.append(nn.ReLU())

        self.bbox_subnet = nn.Sequential(*bbox_subnet)
        self.bbox_pred = nn.Conv2d(
            cur_channels, num_anchors * 4, kernel_size=3, stride=1, padding=1
        )
        modules_list.extend([self.bbox_subnet, self.bbox_pred])

        # mask subnet
        if self.mask_on:
            mask_subnet = []
            cur_channels = in_channels
            for _ in range(num_convs):
                mask_subnet.append(
                    nn.Conv2d(cur_channels, mask_channels, kernel_size=3, stride=1, padding=1)
                )
                cur_channels = mask_channels
                mask_subnet.append(nn.ReLU())

            self.mask_subnet = nn.Sequential(*mask_subnet)
            modules_list.append(self.mask_subnet)
            for mask_size in self.mask_sizes:
                cur_mask_module = "mask_pred_%02d" % mask_size
                self.add_module(
                    cur_mask_module,
                    nn.Conv2d(
                        cur_channels, mask_size * mask_size, kernel_size=1, stride=1, padding=0
                    ),
                )
                modules_list.append(getattr(self, cur_mask_module))
            if self.align_on:
                if self.bipyramid_on:
                    for lvl in range(num_levels):
                        cur_mask_module = "align2nat_%02d" % lvl
                        lambda_val = 2 ** lvl
                        setattr(self, cur_mask_module, SwapAlign2Nat(lambda_val))
                    # Also the fusing layer, stay at the same channel size
                    mask_fuse = [
                        nn.Conv2d(cur_channels, cur_channels, kernel_size=3, stride=1, padding=1),
                        nn.ReLU(),
                    ]
                    self.mask_fuse = nn.Sequential(*mask_fuse)
                    modules_list.append(self.mask_fuse)
                else:
                    self.align2nat = SwapAlign2Nat(1)

        # Initialization
        for modules in modules_list:
            for layer in modules.modules():
                if isinstance(layer, nn.Conv2d):
                    torch.nn.init.normal_(layer.weight, mean=0, std=0.01)
                    torch.nn.init.constant_(layer.bias, 0)

        # Use prior in model initialization to improve stability
        bias_value = -(math.log((1 - 0.01) / 0.01))
        torch.nn.init.constant_(self.cls_score.bias, bias_value)

    def forward(self, features):
        """
        Arguments:
            features (list[Tensor]): FPN feature map tensors in high to low resolution.
                Each tensor in the list correspond to different feature levels.

        Returns:
            pred_logits (list[Tensor]): #lvl tensors, each has shape (N, AxK, Hi, Wi).
                The tensor predicts the classification probability
                at each spatial position for each of the A anchors and K object
                classes.
            pred_deltas (list[Tensor]): #lvl tensors, each has shape (N, Ax4, Hi, Wi).
                The tensor predicts 4-vector (dx,dy,dw,dh) box
                regression values for every anchor. These values are the
                relative offset between the anchor and the ground truth box.
            pred_masks (list(list[Tensor])): #lvl list of tensors, each is a list of
                A tensors of shape (N, M_{i,a}, Hi, Wi).
                The tensor predicts a dense set of M_ixM_i masks at every location.
        """
        pred_logits = [self.cls_score(self.cls_subnet(x)) for x in features]
        pred_deltas = [self.bbox_pred(self.bbox_subnet(x)) for x in features]

        pred_masks = None
        if self.mask_on:
            mask_feats = [self.mask_subnet(x) for x in features]

            if self.bipyramid_on:
                mask_feat_high_res = mask_feats[0]
                H, W = mask_feat_high_res.shape[-2:]
                mask_feats_up = []
                for lvl, mask_feat in enumerate(mask_feats):
                    lambda_val = 2.0 ** lvl
                    mask_feat_up = mask_feat
                    if lvl > 0:
                        mask_feat_up = F.interpolate(
                            mask_feat, scale_factor=lambda_val, mode="bilinear", align_corners=False
                        )
                    mask_feats_up.append(
                        self.mask_fuse(mask_feat_up[:, :, :H, :W] + mask_feat_high_res)
                    )
                mask_feats = mask_feats_up

            pred_masks = []
            for lvl, mask_feat in enumerate(mask_feats):
                cur_masks = []
                for mask_size in self.mask_sizes:
                    cur_mask_module = getattr(self, "mask_pred_%02d" % mask_size)
                    cur_mask = cur_mask_module(mask_feat)
                    if self.align_on:
                        if self.bipyramid_on:
                            cur_mask_module = getattr(self, "align2nat_%02d" % lvl)
                            cur_mask = cur_mask_module(cur_mask)
                        else:
                            cur_mask = self.align2nat(cur_mask)
                    cur_masks.append(cur_mask)
                pred_masks.append(cur_masks)
        return pred_logits, pred_deltas, pred_masks