# 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