mmdet/models/dense_heads/centernet_update_head.py

# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List, Optional, Sequence, Tuple

import torch
import torch.nn as nn
from mmcv.cnn import Scale
from mmengine.structures import InstanceData
from torch import Tensor

from mmdet.registry import MODELS
from mmdet.structures.bbox import bbox2distance
from mmdet.utils import (ConfigType, InstanceList, OptConfigType,
                         OptInstanceList, reduce_mean)
from ..utils import multi_apply
from .anchor_free_head import AnchorFreeHead

INF = 1000000000
RangeType = Sequence[Tuple[int, int]]


def _transpose(tensor_list: List[Tensor],
               num_point_list: list) -> List[Tensor]:
    """This function is used to transpose image first tensors to level first
    ones."""
    for img_idx in range(len(tensor_list)):
        tensor_list[img_idx] = torch.split(
            tensor_list[img_idx], num_point_list, dim=0)

    tensors_level_first = []
    for targets_per_level in zip(*tensor_list):
        tensors_level_first.append(torch.cat(targets_per_level, dim=0))
    return tensors_level_first


@MODELS.register_module()
class CenterNetUpdateHead(AnchorFreeHead):
    """CenterNetUpdateHead is an improved version of CenterNet in CenterNet2.
    Paper link `<https://arxiv.org/abs/2103.07461>`_.

    Args:
        num_classes (int): Number of categories excluding the background
            category.
        in_channels (int): Number of channel in the input feature map.
        regress_ranges (Sequence[Tuple[int, int]]): Regress range of multiple
            level points.
        hm_min_radius (int): Heatmap target minimum radius of cls branch.
            Defaults to 4.
        hm_min_overlap (float): Heatmap target minimum overlap of cls branch.
            Defaults to 0.8.
        more_pos_thresh (float): The filtering threshold when the cls branch
            adds more positive samples. Defaults to 0.2.
        more_pos_topk (int): The maximum number of additional positive samples
            added to each gt. Defaults to 9.
        soft_weight_on_reg (bool): Whether to use the soft target of the
            cls branch as the soft weight of the bbox branch.
            Defaults to False.
        loss_cls (:obj:`ConfigDict` or dict): Config of cls loss. Defaults to
            dict(type='GaussianFocalLoss', loss_weight=1.0)
        loss_bbox (:obj:`ConfigDict` or dict): Config of bbox loss. Defaults to
             dict(type='GIoULoss', loss_weight=2.0).
        norm_cfg (:obj:`ConfigDict` or dict, optional): dictionary to construct
            and config norm layer.  Defaults to
            ``norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)``.
        train_cfg (:obj:`ConfigDict` or dict, optional): Training config.
            Unused in CenterNet. Reserved for compatibility with
            SingleStageDetector.
        test_cfg (:obj:`ConfigDict` or dict, optional): Testing config
            of CenterNet.
    """

    def __init__(self,
                 num_classes: int,
                 in_channels: int,
                 regress_ranges: RangeType = ((0, 80), (64, 160), (128, 320),
                                              (256, 640), (512, INF)),
                 hm_min_radius: int = 4,
                 hm_min_overlap: float = 0.8,
                 more_pos_thresh: float = 0.2,
                 more_pos_topk: int = 9,
                 soft_weight_on_reg: bool = False,
                 loss_cls: ConfigType = dict(
                     type='GaussianFocalLoss',
                     pos_weight=0.25,
                     neg_weight=0.75,
                     loss_weight=1.0),
                 loss_bbox: ConfigType = dict(
                     type='GIoULoss', loss_weight=2.0),
                 norm_cfg: OptConfigType = dict(
                     type='GN', num_groups=32, requires_grad=True),
                 train_cfg: OptConfigType = None,
                 test_cfg: OptConfigType = None,
                 **kwargs) -> None:
        super().__init__(
            num_classes=num_classes,
            in_channels=in_channels,
            loss_cls=loss_cls,
            loss_bbox=loss_bbox,
            norm_cfg=norm_cfg,
            train_cfg=train_cfg,
            test_cfg=test_cfg,
            **kwargs)
        self.soft_weight_on_reg = soft_weight_on_reg
        self.hm_min_radius = hm_min_radius
        self.more_pos_thresh = more_pos_thresh
        self.more_pos_topk = more_pos_topk
        self.delta = (1 - hm_min_overlap) / (1 + hm_min_overlap)
        self.sigmoid_clamp = 0.0001

        # GaussianFocalLoss must be sigmoid mode
        self.use_sigmoid_cls = True
        self.cls_out_channels = num_classes

        self.regress_ranges = regress_ranges
        self.scales = nn.ModuleList([Scale(1.0) for _ in self.strides])

    def _init_predictor(self) -> None:
        """Initialize predictor layers of the head."""
        self.conv_cls = nn.Conv2d(
            self.feat_channels, self.num_classes, 3, padding=1)
        self.conv_reg = nn.Conv2d(self.feat_channels, 4, 3, padding=1)

    def forward(self, x: Tuple[Tensor]) -> Tuple[List[Tensor], List[Tensor]]:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of each level outputs.

            - cls_scores (list[Tensor]): Box scores for each scale level, \
            each is a 4D-tensor, the channel number is num_classes.
            - bbox_preds (list[Tensor]): Box energies / deltas for each \
            scale level, each is a 4D-tensor, the channel number is 4.
        """
        return multi_apply(self.forward_single, x, self.scales, self.strides)

    def forward_single(self, x: Tensor, scale: Scale,
                       stride: int) -> Tuple[Tensor, Tensor]:
        """Forward features of a single scale level.

        Args:
            x (Tensor): FPN feature maps of the specified stride.
            scale (:obj:`mmcv.cnn.Scale`): Learnable scale module to resize
                the bbox prediction.
            stride (int): The corresponding stride for feature maps.

        Returns:
            tuple: scores for each class, bbox predictions of
            input feature maps.
        """
        cls_score, bbox_pred, _, _ = super().forward_single(x)
        # scale the bbox_pred of different level
        # float to avoid overflow when enabling FP16
        bbox_pred = scale(bbox_pred).float()
        # bbox_pred needed for gradient computation has been modified
        # by F.relu(bbox_pred) when run with PyTorch 1.10. So replace
        # F.relu(bbox_pred) with bbox_pred.clamp(min=0)
        bbox_pred = bbox_pred.clamp(min=0)
        if not self.training:
            bbox_pred *= stride
        return cls_score, bbox_pred

    def loss_by_feat(
        self,
        cls_scores: List[Tensor],
        bbox_preds: List[Tensor],
        batch_gt_instances: InstanceList,
        batch_img_metas: List[dict],
        batch_gt_instances_ignore: OptInstanceList = None
    ) -> Dict[str, Tensor]:
        """Calculate the loss based on the features extracted by the detection
        head.

        Args:
            cls_scores (list[Tensor]): Box scores for each scale level,
                each is a 4D-tensor, the channel number is num_classes.
            bbox_preds (list[Tensor]): Box energies / deltas for each scale
                level, each is a 4D-tensor, the channel number is 4.
            batch_gt_instances (list[:obj:`InstanceData`]): Batch of
                gt_instance.  It usually includes ``bboxes`` and ``labels``
                attributes.
            batch_img_metas (list[dict]): Meta information of each image, e.g.,
                image size, scaling factor, etc.
            batch_gt_instances_ignore (list[:obj:`InstanceData`], optional):
                Batch of gt_instances_ignore. It includes ``bboxes`` attribute
                data that is ignored during training and testing.
                Defaults to None.

        Returns:
            dict[str, Tensor]: A dictionary of loss components.
        """
        num_imgs = cls_scores[0].size(0)
        assert len(cls_scores) == len(bbox_preds)
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
        all_level_points = self.prior_generator.grid_priors(
            featmap_sizes,
            dtype=bbox_preds[0].dtype,
            device=bbox_preds[0].device)

        # 1 flatten outputs
        flatten_cls_scores = [
            cls_score.permute(0, 2, 3, 1).reshape(-1, self.cls_out_channels)
            for cls_score in cls_scores
        ]
        flatten_bbox_preds = [
            bbox_pred.permute(0, 2, 3, 1).reshape(-1, 4)
            for bbox_pred in bbox_preds
        ]
        flatten_cls_scores = torch.cat(flatten_cls_scores)
        flatten_bbox_preds = torch.cat(flatten_bbox_preds)

        # repeat points to align with bbox_preds
        flatten_points = torch.cat(
            [points.repeat(num_imgs, 1) for points in all_level_points])

        assert (torch.isfinite(flatten_bbox_preds).all().item())

        # 2 calc reg and cls branch targets
        cls_targets, bbox_targets = self.get_targets(all_level_points,
                                                     batch_gt_instances)

        # 3 add more pos index for cls branch
        featmap_sizes = flatten_points.new_tensor(featmap_sizes)
        pos_inds, cls_labels = self.add_cls_pos_inds(flatten_points,
                                                     flatten_bbox_preds,
                                                     featmap_sizes,
                                                     batch_gt_instances)

        # 4 calc cls loss
        if pos_inds is None:
            # num_gts=0
            num_pos_cls = bbox_preds[0].new_tensor(0, dtype=torch.float)
        else:
            num_pos_cls = bbox_preds[0].new_tensor(
                len(pos_inds), dtype=torch.float)
        num_pos_cls = max(reduce_mean(num_pos_cls), 1.0)
        flatten_cls_scores = flatten_cls_scores.sigmoid().clamp(
            min=self.sigmoid_clamp, max=1 - self.sigmoid_clamp)
        cls_loss = self.loss_cls(
            flatten_cls_scores,
            cls_targets,
            pos_inds=pos_inds,
            pos_labels=cls_labels,
            avg_factor=num_pos_cls)

        # 5 calc reg loss
        pos_bbox_inds = torch.nonzero(
            bbox_targets.max(dim=1)[0] >= 0).squeeze(1)
        pos_bbox_preds = flatten_bbox_preds[pos_bbox_inds]
        pos_bbox_targets = bbox_targets[pos_bbox_inds]

        bbox_weight_map = cls_targets.max(dim=1)[0]
        bbox_weight_map = bbox_weight_map[pos_bbox_inds]
        bbox_weight_map = bbox_weight_map if self.soft_weight_on_reg \
            else torch.ones_like(bbox_weight_map)
        num_pos_bbox = max(reduce_mean(bbox_weight_map.sum()), 1.0)

        if len(pos_bbox_inds) > 0:
            pos_points = flatten_points[pos_bbox_inds]
            pos_decoded_bbox_preds = self.bbox_coder.decode(
                pos_points, pos_bbox_preds)
            pos_decoded_target_preds = self.bbox_coder.decode(
                pos_points, pos_bbox_targets)
            bbox_loss = self.loss_bbox(
                pos_decoded_bbox_preds,
                pos_decoded_target_preds,
                weight=bbox_weight_map,
                avg_factor=num_pos_bbox)
        else:
            bbox_loss = flatten_bbox_preds.sum() * 0

        return dict(loss_cls=cls_loss, loss_bbox=bbox_loss)

    def get_targets(
        self,
        points: List[Tensor],
        batch_gt_instances: InstanceList,
    ) -> Tuple[Tensor, Tensor]:
        """Compute classification and bbox targets for points in multiple
        images.

        Args:
            points (list[Tensor]): Points of each fpn level, each has shape
                (num_points, 2).
            batch_gt_instances (list[:obj:`InstanceData`]): Batch of
                gt_instance.  It usually includes ``bboxes`` and ``labels``
                attributes.

        Returns:
            tuple: Targets of each level.

            - concat_lvl_labels (Tensor): Labels of all level and batch.
            - concat_lvl_bbox_targets (Tensor): BBox targets of all \
            level and batch.
        """
        assert len(points) == len(self.regress_ranges)

        num_levels = len(points)
        # the number of points per img, per lvl
        num_points = [center.size(0) for center in points]

        # expand regress ranges to align with points
        expanded_regress_ranges = [
            points[i].new_tensor(self.regress_ranges[i])[None].expand_as(
                points[i]) for i in range(num_levels)
        ]
        # concat all levels points and regress ranges
        concat_regress_ranges = torch.cat(expanded_regress_ranges, dim=0)
        concat_points = torch.cat(points, dim=0)
        concat_strides = torch.cat([
            concat_points.new_ones(num_points[i]) * self.strides[i]
            for i in range(num_levels)
        ])

        # get labels and bbox_targets of each image
        cls_targets_list, bbox_targets_list = multi_apply(
            self._get_targets_single,
            batch_gt_instances,
            points=concat_points,
            regress_ranges=concat_regress_ranges,
            strides=concat_strides)

        bbox_targets_list = _transpose(bbox_targets_list, num_points)
        cls_targets_list = _transpose(cls_targets_list, num_points)
        concat_lvl_bbox_targets = torch.cat(bbox_targets_list, 0)
        concat_lvl_cls_targets = torch.cat(cls_targets_list, dim=0)
        return concat_lvl_cls_targets, concat_lvl_bbox_targets

    def _get_targets_single(self, gt_instances: InstanceData, points: Tensor,
                            regress_ranges: Tensor,
                            strides: Tensor) -> Tuple[Tensor, Tensor]:
        """Compute classification and bbox targets for a single image."""
        num_points = points.size(0)
        num_gts = len(gt_instances)
        gt_bboxes = gt_instances.bboxes
        gt_labels = gt_instances.labels

        if num_gts == 0:
            return gt_labels.new_full((num_points,
                                       self.num_classes),
                                      self.num_classes), \
                   gt_bboxes.new_full((num_points, 4), -1)

        # Calculate the regression tblr target corresponding to all points
        points = points[:, None].expand(num_points, num_gts, 2)
        gt_bboxes = gt_bboxes[None].expand(num_points, num_gts, 4)
        strides = strides[:, None, None].expand(num_points, num_gts, 2)

        bbox_target = bbox2distance(points, gt_bboxes)  # M x N x 4

        # condition1: inside a gt bbox
        inside_gt_bbox_mask = bbox_target.min(dim=2)[0] > 0  # M x N

        # condition2: Calculate the nearest points from
        # the upper, lower, left and right ranges from
        # the center of the gt bbox
        centers = ((gt_bboxes[..., [0, 1]] + gt_bboxes[..., [2, 3]]) / 2)
        centers_discret = ((centers / strides).int() * strides).float() + \
            strides / 2

        centers_discret_dist = points - centers_discret
        dist_x = centers_discret_dist[..., 0].abs()
        dist_y = centers_discret_dist[..., 1].abs()
        inside_gt_center3x3_mask = (dist_x <= strides[..., 0]) & \
                                   (dist_y <= strides[..., 0])

        # condition3: limit the regression range for each location
        bbox_target_wh = bbox_target[..., :2] + bbox_target[..., 2:]
        crit = (bbox_target_wh**2).sum(dim=2)**0.5 / 2
        inside_fpn_level_mask = (crit >= regress_ranges[:, [0]]) & \
                                (crit <= regress_ranges[:, [1]])
        bbox_target_mask = inside_gt_bbox_mask & \
            inside_gt_center3x3_mask & \
            inside_fpn_level_mask

        # Calculate the distance weight map
        gt_center_peak_mask = ((centers_discret_dist**2).sum(dim=2) == 0)
        weighted_dist = ((points - centers)**2).sum(dim=2)  # M x N
        weighted_dist[gt_center_peak_mask] = 0

        areas = (gt_bboxes[..., 2] - gt_bboxes[..., 0]) * (
            gt_bboxes[..., 3] - gt_bboxes[..., 1])
        radius = self.delta**2 * 2 * areas
        radius = torch.clamp(radius, min=self.hm_min_radius**2)
        weighted_dist = weighted_dist / radius

        # Calculate bbox_target
        bbox_weighted_dist = weighted_dist.clone()
        bbox_weighted_dist[bbox_target_mask == 0] = INF * 1.0
        min_dist, min_inds = bbox_weighted_dist.min(dim=1)
        bbox_target = bbox_target[range(len(bbox_target)),
                                  min_inds]  # M x N x 4 --> M x 4
        bbox_target[min_dist == INF] = -INF

        # Convert to feature map scale
        bbox_target /= strides[:, 0, :].repeat(1, 2)

        # Calculate cls_target
        cls_target = self._create_heatmaps_from_dist(weighted_dist, gt_labels)

        return cls_target, bbox_target

    @torch.no_grad()
    def add_cls_pos_inds(
        self, flatten_points: Tensor, flatten_bbox_preds: Tensor,
        featmap_sizes: Tensor, batch_gt_instances: InstanceList
    ) -> Tuple[Optional[Tensor], Optional[Tensor]]:
        """Provide additional adaptive positive samples to the classification
        branch.

        Args:
            flatten_points (Tensor): The point after flatten, including
                batch image and all levels. The shape is (N, 2).
            flatten_bbox_preds (Tensor): The bbox predicts after flatten,
                including batch image and all levels. The shape is (N, 4).
            featmap_sizes (Tensor): Feature map size of all layers.
                The shape is (5, 2).
            batch_gt_instances (list[:obj:`InstanceData`]): Batch of
                gt_instance.  It usually includes ``bboxes`` and ``labels``
                attributes.

        Returns:
           tuple:

           - pos_inds (Tensor): Adaptively selected positive sample index.
           - cls_labels (Tensor): Corresponding positive class label.
        """
        outputs = self._get_center3x3_region_index_targets(
            batch_gt_instances, featmap_sizes)
        cls_labels, fpn_level_masks, center3x3_inds, \
            center3x3_bbox_targets, center3x3_masks = outputs

        num_gts, total_level, K = cls_labels.shape[0], len(
            self.strides), center3x3_masks.shape[-1]

        if num_gts == 0:
            return None, None

        # The out-of-bounds index is forcibly set to 0
        # to prevent loss calculation errors
        center3x3_inds[center3x3_masks == 0] = 0
        reg_pred_center3x3 = flatten_bbox_preds[center3x3_inds]
        center3x3_points = flatten_points[center3x3_inds].view(-1, 2)

        center3x3_bbox_targets_expand = center3x3_bbox_targets.view(
            -1, 4).clamp(min=0)

        pos_decoded_bbox_preds = self.bbox_coder.decode(
            center3x3_points, reg_pred_center3x3.view(-1, 4))
        pos_decoded_target_preds = self.bbox_coder.decode(
            center3x3_points, center3x3_bbox_targets_expand)
        center3x3_bbox_loss = self.loss_bbox(
            pos_decoded_bbox_preds,
            pos_decoded_target_preds,
            None,
            reduction_override='none').view(num_gts, total_level,
                                            K) / self.loss_bbox.loss_weight

        # Invalid index Loss set to infinity
        center3x3_bbox_loss[center3x3_masks == 0] = INF

        # 4 is the center point of the sampled 9 points, the center point
        # of gt bbox after discretization.
        # The center point of gt bbox after discretization
        # must be a positive sample, so we force its loss to be set to 0.
        center3x3_bbox_loss.view(-1, K)[fpn_level_masks.view(-1), 4] = 0
        center3x3_bbox_loss = center3x3_bbox_loss.view(num_gts, -1)

        loss_thr = torch.kthvalue(
            center3x3_bbox_loss, self.more_pos_topk, dim=1)[0]

        loss_thr[loss_thr > self.more_pos_thresh] = self.more_pos_thresh
        new_pos = center3x3_bbox_loss < loss_thr.view(num_gts, 1)
        pos_inds = center3x3_inds.view(num_gts, -1)[new_pos]
        cls_labels = cls_labels.view(num_gts,
                                     1).expand(num_gts,
                                               total_level * K)[new_pos]
        return pos_inds, cls_labels

    def _create_heatmaps_from_dist(self, weighted_dist: Tensor,
                                   cls_labels: Tensor) -> Tensor:
        """Generate heatmaps of classification branch based on weighted
        distance map."""
        heatmaps = weighted_dist.new_zeros(
            (weighted_dist.shape[0], self.num_classes))
        for c in range(self.num_classes):
            inds = (cls_labels == c)  # N
            if inds.int().sum() == 0:
                continue
            heatmaps[:, c] = torch.exp(-weighted_dist[:, inds].min(dim=1)[0])
            zeros = heatmaps[:, c] < 1e-4
            heatmaps[zeros, c] = 0
        return heatmaps

    def _get_center3x3_region_index_targets(self,
                                            bacth_gt_instances: InstanceList,
                                            shapes_per_level: Tensor) -> tuple:
        """Get the center (and the 3x3 region near center) locations and target
        of each objects."""
        cls_labels = []
        inside_fpn_level_masks = []
        center3x3_inds = []
        center3x3_masks = []
        center3x3_bbox_targets = []

        total_levels = len(self.strides)
        batch = len(bacth_gt_instances)

        shapes_per_level = shapes_per_level.long()
        area_per_level = (shapes_per_level[:, 0] * shapes_per_level[:, 1])

        # Select a total of 9 positions of 3x3 in the center of the gt bbox
        # as candidate positive samples
        K = 9
        dx = shapes_per_level.new_tensor([-1, 0, 1, -1, 0, 1, -1, 0,
                                          1]).view(1, 1, K)
        dy = shapes_per_level.new_tensor([-1, -1, -1, 0, 0, 0, 1, 1,
                                          1]).view(1, 1, K)

        regress_ranges = shapes_per_level.new_tensor(self.regress_ranges).view(
            len(self.regress_ranges), 2)  # L x 2
        strides = shapes_per_level.new_tensor(self.strides)

        start_coord_pre_level = []
        _start = 0
        for level in range(total_levels):
            start_coord_pre_level.append(_start)
            _start = _start + batch * area_per_level[level]
        start_coord_pre_level = shapes_per_level.new_tensor(
            start_coord_pre_level).view(1, total_levels, 1)
        area_per_level = area_per_level.view(1, total_levels, 1)

        for im_i in range(batch):
            gt_instance = bacth_gt_instances[im_i]
            gt_bboxes = gt_instance.bboxes
            gt_labels = gt_instance.labels
            num_gts = gt_bboxes.shape[0]
            if num_gts == 0:
                continue

            cls_labels.append(gt_labels)

            gt_bboxes = gt_bboxes[:, None].expand(num_gts, total_levels, 4)
            expanded_strides = strides[None, :,
                                       None].expand(num_gts, total_levels, 2)
            expanded_regress_ranges = regress_ranges[None].expand(
                num_gts, total_levels, 2)
            expanded_shapes_per_level = shapes_per_level[None].expand(
                num_gts, total_levels, 2)

            # calc reg_target
            centers = ((gt_bboxes[..., [0, 1]] + gt_bboxes[..., [2, 3]]) / 2)
            centers_inds = (centers / expanded_strides).long()
            centers_discret = centers_inds * expanded_strides \
                + expanded_strides // 2

            bbox_target = bbox2distance(centers_discret,
                                        gt_bboxes)  # M x N x 4

            # calc inside_fpn_level_mask
            bbox_target_wh = bbox_target[..., :2] + bbox_target[..., 2:]
            crit = (bbox_target_wh**2).sum(dim=2)**0.5 / 2
            inside_fpn_level_mask = \
                (crit >= expanded_regress_ranges[..., 0]) & \
                (crit <= expanded_regress_ranges[..., 1])

            inside_gt_bbox_mask = bbox_target.min(dim=2)[0] >= 0
            inside_fpn_level_mask = inside_gt_bbox_mask & inside_fpn_level_mask
            inside_fpn_level_masks.append(inside_fpn_level_mask)

            # calc center3x3_ind and mask
            expand_ws = expanded_shapes_per_level[..., 1:2].expand(
                num_gts, total_levels, K)
            expand_hs = expanded_shapes_per_level[..., 0:1].expand(
                num_gts, total_levels, K)
            centers_inds_x = centers_inds[..., 0:1]
            centers_inds_y = centers_inds[..., 1:2]

            center3x3_idx = start_coord_pre_level + \
                im_i * area_per_level + \
                (centers_inds_y + dy) * expand_ws + \
                (centers_inds_x + dx)
            center3x3_mask = \
                ((centers_inds_y + dy) < expand_hs) & \
                ((centers_inds_y + dy) >= 0) & \
                ((centers_inds_x + dx) < expand_ws) & \
                ((centers_inds_x + dx) >= 0)

            # recalc center3x3 region reg target
            bbox_target = bbox_target / expanded_strides.repeat(1, 1, 2)
            center3x3_bbox_target = bbox_target[..., None, :].expand(
                num_gts, total_levels, K, 4).clone()
            center3x3_bbox_target[..., 0] += dx
            center3x3_bbox_target[..., 1] += dy
            center3x3_bbox_target[..., 2] -= dx
            center3x3_bbox_target[..., 3] -= dy
            # update center3x3_mask
            center3x3_mask = center3x3_mask & (
                center3x3_bbox_target.min(dim=3)[0] >= 0)  # n x L x K

            center3x3_inds.append(center3x3_idx)
            center3x3_masks.append(center3x3_mask)
            center3x3_bbox_targets.append(center3x3_bbox_target)

        if len(inside_fpn_level_masks) > 0:
            cls_labels = torch.cat(cls_labels, dim=0)
            inside_fpn_level_masks = torch.cat(inside_fpn_level_masks, dim=0)
            center3x3_inds = torch.cat(center3x3_inds, dim=0).long()
            center3x3_bbox_targets = torch.cat(center3x3_bbox_targets, dim=0)
            center3x3_masks = torch.cat(center3x3_masks, dim=0)
        else:
            cls_labels = shapes_per_level.new_zeros(0).long()
            inside_fpn_level_masks = shapes_per_level.new_zeros(
                (0, total_levels)).bool()
            center3x3_inds = shapes_per_level.new_zeros(
                (0, total_levels, K)).long()
            center3x3_bbox_targets = shapes_per_level.new_zeros(
                (0, total_levels, K, 4)).float()
            center3x3_masks = shapes_per_level.new_zeros(
                (0, total_levels, K)).bool()
        return cls_labels, inside_fpn_level_masks, center3x3_inds, \
            center3x3_bbox_targets, center3x3_masks