Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import logging | |
| import math | |
| from typing import List, Tuple | |
| import torch | |
| from fvcore.nn import sigmoid_focal_loss_jit | |
| from torch import Tensor, nn | |
| from torch.nn import functional as F | |
| from detectron2.config import configurable | |
| from detectron2.layers import CycleBatchNormList, ShapeSpec, batched_nms, cat, get_norm | |
| from detectron2.structures import Boxes, ImageList, Instances, pairwise_iou | |
| from detectron2.utils.events import get_event_storage | |
| from ..anchor_generator import build_anchor_generator | |
| from ..backbone import Backbone, build_backbone | |
| from ..box_regression import Box2BoxTransform, _dense_box_regression_loss | |
| from ..matcher import Matcher | |
| from .build import META_ARCH_REGISTRY | |
| from .dense_detector import DenseDetector, permute_to_N_HWA_K # noqa | |
| __all__ = ["RetinaNet"] | |
| logger = logging.getLogger(__name__) | |
| class RetinaNet(DenseDetector): | |
| """ | |
| Implement RetinaNet in :paper:`RetinaNet`. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| backbone: Backbone, | |
| head: nn.Module, | |
| head_in_features, | |
| anchor_generator, | |
| box2box_transform, | |
| anchor_matcher, | |
| num_classes, | |
| focal_loss_alpha=0.25, | |
| focal_loss_gamma=2.0, | |
| smooth_l1_beta=0.0, | |
| box_reg_loss_type="smooth_l1", | |
| test_score_thresh=0.05, | |
| test_topk_candidates=1000, | |
| test_nms_thresh=0.5, | |
| max_detections_per_image=100, | |
| pixel_mean, | |
| pixel_std, | |
| vis_period=0, | |
| input_format="BGR", | |
| ): | |
| """ | |
| NOTE: this interface is experimental. | |
| Args: | |
| backbone: a backbone module, must follow detectron2's backbone interface | |
| head (nn.Module): a module that predicts logits and regression deltas | |
| for each level from a list of per-level features | |
| head_in_features (Tuple[str]): Names of the input feature maps to be used in head | |
| anchor_generator (nn.Module): a module that creates anchors from a | |
| list of features. Usually an instance of :class:`AnchorGenerator` | |
| box2box_transform (Box2BoxTransform): defines the transform from anchors boxes to | |
| instance boxes | |
| anchor_matcher (Matcher): label the anchors by matching them with ground truth. | |
| num_classes (int): number of classes. Used to label background proposals. | |
| # Loss parameters: | |
| focal_loss_alpha (float): focal_loss_alpha | |
| focal_loss_gamma (float): focal_loss_gamma | |
| smooth_l1_beta (float): smooth_l1_beta | |
| box_reg_loss_type (str): Options are "smooth_l1", "giou", "diou", "ciou" | |
| # Inference parameters: | |
| test_score_thresh (float): Inference cls score threshold, only anchors with | |
| score > INFERENCE_TH are considered for inference (to improve speed) | |
| test_topk_candidates (int): Select topk candidates before NMS | |
| test_nms_thresh (float): Overlap threshold used for non-maximum suppression | |
| (suppress boxes with IoU >= this threshold) | |
| max_detections_per_image (int): | |
| Maximum number of detections to return per image during inference | |
| (100 is based on the limit established for the COCO dataset). | |
| pixel_mean, pixel_std: see :class:`DenseDetector`. | |
| """ | |
| super().__init__( | |
| backbone, head, head_in_features, pixel_mean=pixel_mean, pixel_std=pixel_std | |
| ) | |
| self.num_classes = num_classes | |
| # Anchors | |
| self.anchor_generator = anchor_generator | |
| self.box2box_transform = box2box_transform | |
| self.anchor_matcher = anchor_matcher | |
| # Loss parameters: | |
| self.focal_loss_alpha = focal_loss_alpha | |
| self.focal_loss_gamma = focal_loss_gamma | |
| self.smooth_l1_beta = smooth_l1_beta | |
| self.box_reg_loss_type = box_reg_loss_type | |
| # Inference parameters: | |
| self.test_score_thresh = test_score_thresh | |
| self.test_topk_candidates = test_topk_candidates | |
| self.test_nms_thresh = test_nms_thresh | |
| self.max_detections_per_image = max_detections_per_image | |
| # Vis parameters | |
| self.vis_period = vis_period | |
| self.input_format = input_format | |
| def from_config(cls, cfg): | |
| backbone = build_backbone(cfg) | |
| backbone_shape = backbone.output_shape() | |
| feature_shapes = [backbone_shape[f] for f in cfg.MODEL.RETINANET.IN_FEATURES] | |
| head = RetinaNetHead(cfg, feature_shapes) | |
| anchor_generator = build_anchor_generator(cfg, feature_shapes) | |
| return { | |
| "backbone": backbone, | |
| "head": head, | |
| "anchor_generator": anchor_generator, | |
| "box2box_transform": Box2BoxTransform(weights=cfg.MODEL.RETINANET.BBOX_REG_WEIGHTS), | |
| "anchor_matcher": Matcher( | |
| cfg.MODEL.RETINANET.IOU_THRESHOLDS, | |
| cfg.MODEL.RETINANET.IOU_LABELS, | |
| allow_low_quality_matches=True, | |
| ), | |
| "pixel_mean": cfg.MODEL.PIXEL_MEAN, | |
| "pixel_std": cfg.MODEL.PIXEL_STD, | |
| "num_classes": cfg.MODEL.RETINANET.NUM_CLASSES, | |
| "head_in_features": cfg.MODEL.RETINANET.IN_FEATURES, | |
| # Loss parameters: | |
| "focal_loss_alpha": cfg.MODEL.RETINANET.FOCAL_LOSS_ALPHA, | |
| "focal_loss_gamma": cfg.MODEL.RETINANET.FOCAL_LOSS_GAMMA, | |
| "smooth_l1_beta": cfg.MODEL.RETINANET.SMOOTH_L1_LOSS_BETA, | |
| "box_reg_loss_type": cfg.MODEL.RETINANET.BBOX_REG_LOSS_TYPE, | |
| # Inference parameters: | |
| "test_score_thresh": cfg.MODEL.RETINANET.SCORE_THRESH_TEST, | |
| "test_topk_candidates": cfg.MODEL.RETINANET.TOPK_CANDIDATES_TEST, | |
| "test_nms_thresh": cfg.MODEL.RETINANET.NMS_THRESH_TEST, | |
| "max_detections_per_image": cfg.TEST.DETECTIONS_PER_IMAGE, | |
| # Vis parameters | |
| "vis_period": cfg.VIS_PERIOD, | |
| "input_format": cfg.INPUT.FORMAT, | |
| } | |
| def forward_training(self, images, features, predictions, gt_instances): | |
| # Transpose the Hi*Wi*A dimension to the middle: | |
| pred_logits, pred_anchor_deltas = self._transpose_dense_predictions( | |
| predictions, [self.num_classes, 4] | |
| ) | |
| anchors = self.anchor_generator(features) | |
| gt_labels, gt_boxes = self.label_anchors(anchors, gt_instances) | |
| return self.losses(anchors, pred_logits, gt_labels, pred_anchor_deltas, gt_boxes) | |
| def losses(self, anchors, pred_logits, gt_labels, pred_anchor_deltas, gt_boxes): | |
| """ | |
| Args: | |
| anchors (list[Boxes]): a list of #feature level Boxes | |
| gt_labels, gt_boxes: see output of :meth:`RetinaNet.label_anchors`. | |
| Their shapes are (N, R) and (N, R, 4), respectively, where R is | |
| the total number of anchors across levels, i.e. sum(Hi x Wi x Ai) | |
| pred_logits, pred_anchor_deltas: both are list[Tensor]. Each element in the | |
| list corresponds to one level and has shape (N, Hi * Wi * Ai, K or 4). | |
| Where K is the number of classes used in `pred_logits`. | |
| Returns: | |
| dict[str, Tensor]: | |
| mapping from a named loss to a scalar tensor storing the loss. | |
| Used during training only. The dict keys are: "loss_cls" and "loss_box_reg" | |
| """ | |
| num_images = len(gt_labels) | |
| gt_labels = torch.stack(gt_labels) # (N, R) | |
| valid_mask = gt_labels >= 0 | |
| pos_mask = (gt_labels >= 0) & (gt_labels != self.num_classes) | |
| num_pos_anchors = pos_mask.sum().item() | |
| get_event_storage().put_scalar("num_pos_anchors", num_pos_anchors / num_images) | |
| normalizer = self._ema_update("loss_normalizer", max(num_pos_anchors, 1), 100) | |
| # classification and regression loss | |
| gt_labels_target = F.one_hot(gt_labels[valid_mask], num_classes=self.num_classes + 1)[ | |
| :, :-1 | |
| ] # no loss for the last (background) class | |
| loss_cls = sigmoid_focal_loss_jit( | |
| cat(pred_logits, dim=1)[valid_mask], | |
| gt_labels_target.to(pred_logits[0].dtype), | |
| alpha=self.focal_loss_alpha, | |
| gamma=self.focal_loss_gamma, | |
| reduction="sum", | |
| ) | |
| loss_box_reg = _dense_box_regression_loss( | |
| anchors, | |
| self.box2box_transform, | |
| pred_anchor_deltas, | |
| gt_boxes, | |
| pos_mask, | |
| box_reg_loss_type=self.box_reg_loss_type, | |
| smooth_l1_beta=self.smooth_l1_beta, | |
| ) | |
| return { | |
| "loss_cls": loss_cls / normalizer, | |
| "loss_box_reg": loss_box_reg / normalizer, | |
| } | |
| def label_anchors(self, anchors, gt_instances): | |
| """ | |
| Args: | |
| anchors (list[Boxes]): A list of #feature level Boxes. | |
| The Boxes contains anchors of this image on the specific feature level. | |
| gt_instances (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. | |
| Returns: | |
| list[Tensor]: List of #img tensors. i-th element is a vector of labels whose length is | |
| the total number of anchors across all feature maps (sum(Hi * Wi * A)). | |
| Label values are in {-1, 0, ..., K}, with -1 means ignore, and K means background. | |
| list[Tensor]: i-th element is a Rx4 tensor, where R is the total number of anchors | |
| across feature maps. The values are the matched gt boxes for each anchor. | |
| Values are undefined for those anchors not labeled as foreground. | |
| """ | |
| anchors = Boxes.cat(anchors) # Rx4 | |
| gt_labels = [] | |
| matched_gt_boxes = [] | |
| for gt_per_image in gt_instances: | |
| match_quality_matrix = pairwise_iou(gt_per_image.gt_boxes, anchors) | |
| matched_idxs, anchor_labels = self.anchor_matcher(match_quality_matrix) | |
| del match_quality_matrix | |
| if len(gt_per_image) > 0: | |
| matched_gt_boxes_i = gt_per_image.gt_boxes.tensor[matched_idxs] | |
| gt_labels_i = gt_per_image.gt_classes[matched_idxs] | |
| # Anchors with label 0 are treated as background. | |
| gt_labels_i[anchor_labels == 0] = self.num_classes | |
| # Anchors with label -1 are ignored. | |
| gt_labels_i[anchor_labels == -1] = -1 | |
| else: | |
| matched_gt_boxes_i = torch.zeros_like(anchors.tensor) | |
| gt_labels_i = torch.zeros_like(matched_idxs) + self.num_classes | |
| gt_labels.append(gt_labels_i) | |
| matched_gt_boxes.append(matched_gt_boxes_i) | |
| return gt_labels, matched_gt_boxes | |
| def forward_inference( | |
| self, images: ImageList, features: List[Tensor], predictions: List[List[Tensor]] | |
| ): | |
| pred_logits, pred_anchor_deltas = self._transpose_dense_predictions( | |
| predictions, [self.num_classes, 4] | |
| ) | |
| anchors = self.anchor_generator(features) | |
| results: List[Instances] = [] | |
| for img_idx, image_size in enumerate(images.image_sizes): | |
| scores_per_image = [x[img_idx].sigmoid_() for x in pred_logits] | |
| deltas_per_image = [x[img_idx] for x in pred_anchor_deltas] | |
| results_per_image = self.inference_single_image( | |
| anchors, scores_per_image, deltas_per_image, image_size | |
| ) | |
| results.append(results_per_image) | |
| return results | |
| def inference_single_image( | |
| self, | |
| anchors: List[Boxes], | |
| box_cls: List[Tensor], | |
| box_delta: List[Tensor], | |
| image_size: Tuple[int, int], | |
| ): | |
| """ | |
| Single-image inference. Return bounding-box detection results by thresholding | |
| on scores and applying non-maximum suppression (NMS). | |
| Arguments: | |
| anchors (list[Boxes]): list of #feature levels. Each entry contains | |
| a Boxes object, which contains all the anchors in that feature level. | |
| box_cls (list[Tensor]): list of #feature levels. Each entry contains | |
| tensor of size (H x W x A, K) | |
| box_delta (list[Tensor]): Same shape as 'box_cls' except that K becomes 4. | |
| image_size (tuple(H, W)): a tuple of the image height and width. | |
| Returns: | |
| Same as `inference`, but for only one image. | |
| """ | |
| pred = self._decode_multi_level_predictions( | |
| anchors, | |
| box_cls, | |
| box_delta, | |
| self.test_score_thresh, | |
| self.test_topk_candidates, | |
| image_size, | |
| ) | |
| keep = batched_nms( # per-class NMS | |
| pred.pred_boxes.tensor, pred.scores, pred.pred_classes, self.test_nms_thresh | |
| ) | |
| return pred[keep[: self.max_detections_per_image]] | |
| class RetinaNetHead(nn.Module): | |
| """ | |
| The head used in RetinaNet for object classification and box regression. | |
| It has two subnets for the two tasks, with a common structure but separate parameters. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| input_shape: List[ShapeSpec], | |
| num_classes, | |
| num_anchors, | |
| conv_dims: List[int], | |
| norm="", | |
| prior_prob=0.01, | |
| ): | |
| """ | |
| NOTE: this interface is experimental. | |
| Args: | |
| input_shape (List[ShapeSpec]): input shape | |
| num_classes (int): number of classes. Used to label background proposals. | |
| num_anchors (int): number of generated anchors | |
| conv_dims (List[int]): dimensions for each convolution layer | |
| norm (str or callable): | |
| Normalization for conv layers except for the two output layers. | |
| See :func:`detectron2.layers.get_norm` for supported types. | |
| prior_prob (float): Prior weight for computing bias | |
| """ | |
| super().__init__() | |
| self._num_features = len(input_shape) | |
| if norm == "BN" or norm == "SyncBN": | |
| logger.info( | |
| f"Using domain-specific {norm} in RetinaNetHead with len={self._num_features}." | |
| ) | |
| bn_class = nn.BatchNorm2d if norm == "BN" else nn.SyncBatchNorm | |
| def norm(c): | |
| return CycleBatchNormList( | |
| length=self._num_features, bn_class=bn_class, num_features=c | |
| ) | |
| else: | |
| norm_name = str(type(get_norm(norm, 32))) | |
| if "BN" in norm_name: | |
| logger.warning( | |
| f"Shared BatchNorm (type={norm_name}) may not work well in RetinaNetHead." | |
| ) | |
| cls_subnet = [] | |
| bbox_subnet = [] | |
| for in_channels, out_channels in zip( | |
| [input_shape[0].channels] + list(conv_dims), conv_dims | |
| ): | |
| cls_subnet.append( | |
| nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) | |
| ) | |
| if norm: | |
| cls_subnet.append(get_norm(norm, out_channels)) | |
| cls_subnet.append(nn.ReLU()) | |
| bbox_subnet.append( | |
| nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) | |
| ) | |
| if norm: | |
| bbox_subnet.append(get_norm(norm, out_channels)) | |
| bbox_subnet.append(nn.ReLU()) | |
| self.cls_subnet = nn.Sequential(*cls_subnet) | |
| self.bbox_subnet = nn.Sequential(*bbox_subnet) | |
| self.cls_score = nn.Conv2d( | |
| conv_dims[-1], num_anchors * num_classes, kernel_size=3, stride=1, padding=1 | |
| ) | |
| self.bbox_pred = nn.Conv2d( | |
| conv_dims[-1], num_anchors * 4, kernel_size=3, stride=1, padding=1 | |
| ) | |
| # Initialization | |
| for modules in [self.cls_subnet, self.bbox_subnet, self.cls_score, self.bbox_pred]: | |
| 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 - prior_prob) / prior_prob)) | |
| torch.nn.init.constant_(self.cls_score.bias, bias_value) | |
| def from_config(cls, cfg, input_shape: List[ShapeSpec]): | |
| num_anchors = build_anchor_generator(cfg, input_shape).num_cell_anchors | |
| assert ( | |
| len(set(num_anchors)) == 1 | |
| ), "Using different number of anchors between levels is not currently supported!" | |
| num_anchors = num_anchors[0] | |
| return { | |
| "input_shape": input_shape, | |
| "num_classes": cfg.MODEL.RETINANET.NUM_CLASSES, | |
| "conv_dims": [input_shape[0].channels] * cfg.MODEL.RETINANET.NUM_CONVS, | |
| "prior_prob": cfg.MODEL.RETINANET.PRIOR_PROB, | |
| "norm": cfg.MODEL.RETINANET.NORM, | |
| "num_anchors": num_anchors, | |
| } | |
| def forward(self, features: List[Tensor]): | |
| """ | |
| Arguments: | |
| features (list[Tensor]): FPN feature map tensors in high to low resolution. | |
| Each tensor in the list correspond to different feature levels. | |
| Returns: | |
| 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. | |
| bbox_reg (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. | |
| """ | |
| assert len(features) == self._num_features | |
| logits = [] | |
| bbox_reg = [] | |
| for feature in features: | |
| logits.append(self.cls_score(self.cls_subnet(feature))) | |
| bbox_reg.append(self.bbox_pred(self.bbox_subnet(feature))) | |
| return logits, bbox_reg | |