import math import torch import torch.nn.functional as F from torch import nn from .inference import make_fcos_postprocessor from .loss import make_fcos_loss_evaluator from maskrcnn_benchmark.layers import Scale class FCOSHead(torch.nn.Module): def __init__(self, cfg, in_channels): """ Arguments: in_channels (int): number of channels of the input feature """ super(FCOSHead, self).__init__() # TODO: Implement the sigmoid version first. num_classes = cfg.MODEL.FCOS.NUM_CLASSES - 1 cls_tower = [] bbox_tower = [] for i in range(cfg.MODEL.FCOS.NUM_CONVS): cls_tower.append( nn.Conv2d( in_channels, in_channels, kernel_size=3, stride=1, padding=1 ) ) if cfg.MODEL.FCOS.USE_GN: cls_tower.append(nn.GroupNorm(32, in_channels)) cls_tower.append(nn.ReLU()) bbox_tower.append( nn.Conv2d( in_channels, in_channels, kernel_size=3, stride=1, padding=1 ) ) if cfg.MODEL.FCOS.USE_GN: bbox_tower.append(nn.GroupNorm(32, in_channels)) bbox_tower.append(nn.ReLU()) self.add_module('cls_tower', nn.Sequential(*cls_tower)) self.add_module('bbox_tower', nn.Sequential(*bbox_tower)) self.cls_logits = nn.Conv2d( in_channels, num_classes, kernel_size=3, stride=1, padding=1 ) self.bbox_pred = nn.Conv2d( in_channels, 4, kernel_size=3, stride=1, padding=1 ) self.centerness = nn.Conv2d( in_channels, 1, kernel_size=3, stride=1, padding=1 ) # initialization for modules in [self.cls_tower, self.bbox_tower, self.cls_logits, self.bbox_pred, self.centerness]: for l in modules.modules(): if isinstance(l, nn.Conv2d): torch.nn.init.normal_(l.weight, std=0.01) torch.nn.init.constant_(l.bias, 0) # initialize the bias for focal loss prior_prob = cfg.MODEL.FCOS.PRIOR_PROB bias_value = -math.log((1 - prior_prob) / prior_prob) torch.nn.init.constant_(self.cls_logits.bias, bias_value) if cfg.MODEL.FCOS.USE_STRIDE_SCALE_INIT: strides = cfg.MODEL.FCOS.FPN_STRIDES self.scales = nn.ModuleList([Scale(init_value=stride) for stride in strides]) else: self.scales = nn.ModuleList([Scale(init_value=1.0) for _ in range(5)]) def forward(self, x): logits = [] bbox_reg = [] centerness = [] for l, feature in enumerate(x): cls_tower = self.cls_tower(feature) logits.append(self.cls_logits(cls_tower)) centerness.append(self.centerness(cls_tower)) bbox_reg.append(torch.exp(self.scales[l]( self.bbox_pred(self.bbox_tower(feature)) ))) return logits, bbox_reg, centerness class FCOSModule(torch.nn.Module): """ Module for FCOS computation. Takes feature maps from the backbone and FCOS outputs and losses. Only Test on FPN now. """ def __init__(self, cfg, in_channels): super(FCOSModule, self).__init__() head = FCOSHead(cfg, in_channels) box_selector_test = make_fcos_postprocessor(cfg) loss_evaluator = make_fcos_loss_evaluator(cfg) self.head = head self.box_selector_test = box_selector_test self.loss_evaluator = loss_evaluator self.fpn_strides = cfg.MODEL.FCOS.FPN_STRIDES self.vis_labels = cfg.MODEL.FCOS.DEBUG.VIS_LABELS def forward(self, images, features, targets=None): """ Arguments: images (ImageList): images for which we want to compute the predictions features (list[Tensor]): features computed from the images that are used for computing the predictions. Each tensor in the list correspond to different feature levels targets (list[BoxList): ground-truth boxes present in the image (optional) Returns: boxes (list[BoxList]): the predicted boxes from the RPN, one BoxList per image. losses (dict[Tensor]): the losses for the model during training. During testing, it is an empty dict. """ box_cls, box_regression, centerness = self.head(features) locations = self.compute_locations(features) if self.vis_labels: show_image(images, targets) if self.training: return self._forward_train( locations, box_cls, box_regression, centerness, targets ) else: return self._forward_test( locations, box_cls, box_regression, centerness, images.image_sizes ) def _forward_train(self, locations, box_cls, box_regression, centerness, targets): loss_box_cls, loss_box_reg, loss_centerness = self.loss_evaluator( locations, box_cls, box_regression, centerness, targets ) losses = { "loss_cls": loss_box_cls, "loss_reg": loss_box_reg, "loss_centerness": loss_centerness } return None, losses def _forward_test(self, locations, box_cls, box_regression, centerness, image_sizes): boxes = self.box_selector_test( locations, box_cls, box_regression, centerness, image_sizes ) return boxes, {} def compute_locations(self, features): """ iter each fpn level """ locations = [] for level, feature in enumerate(features): h, w = feature.size()[-2:] locations_per_level = self.compute_locations_per_level( h, w, self.fpn_strides[level], feature.device ) locations.append(locations_per_level) return locations def compute_locations_per_level(self, h, w, stride, device): shifts_x = torch.arange( 0, w * stride, step=stride, dtype=torch.float32, device=device ) shifts_y = torch.arange( 0, h * stride, step=stride, dtype=torch.float32, device=device ) shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x) shift_x = shift_x.reshape(-1) shift_y = shift_y.reshape(-1) locations = torch.stack((shift_x, shift_y), dim=1) + stride // 2 return locations def build_fcos(cfg, in_channels): return FCOSModule(cfg, in_channels) def show_image(images, targets): import numpy as np import matplotlib.pyplot as plt img = images.tensors[0].permute((1, 2, 0)).cpu().numpy() + np.array([102.9801, 115.9465, 122.7717]) plt.imshow(img/255) for x1, y1, x2, y2 in targets[0].bbox.cpu().numpy().tolist(): w = x2 - x1 + 1 h = y2 - y1 + 1 rect = plt.Rectangle((x1, y1), w, h, fill=False, color=(1, 0, 0), linewidth=2) plt.axes().add_patch(rect) plt.show()