# Copyright (c) 2019 Western Digital Corporation or its affiliates.
import warnings
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from .dense_test_mixins import BBoxTestMixin
from hikframe.dl.transfer.hikdet.mmdet.core import (build_anchor_generator, build_assigner,
                            build_bbox_coder, build_sampler,
                            images_to_levels, multi_apply, multiclass_nms,
                        show_img, show_bbox)
from ..builder import build_loss
from .base_dense_head import BaseDenseHead


def show_pos_anchor(img_meta, gt_anchors, gt_bboxes, is_show=True):
    # 显示正样本
    assert 'img' in img_meta, print('Collect类中的meta_keys需要新增‘img’，用于可视化调试')
    img = img_meta['img'].data.numpy()
    mean = img_meta['img_norm_cfg']['mean']
    std = img_meta['img_norm_cfg']['std']
    # 默认输入是rgb数据，需要切换为bgr显示
    img = np.transpose(img.copy(), (1, 2, 0))[..., ::-1]
    img = img * std.reshape([1, 1, 3])[..., ::-1] + mean.reshape([1, 1, 3])[..., ::-1]
    img = img.astype(np.uint8)
    img = show_bbox(img, gt_anchors.cpu().numpy(), is_show=False)
    return show_bbox(img, gt_bboxes.cpu().numpy(), color=(255, 255, 255), is_show=is_show)


####### the order is 13*13, 26*26, 52*52
def bbox_to_levels(all_pos_bboxes, all_pos_gt_bboxes, all_anchor_strides):
    pred_bboxes = {}
    gt_bboxes= {}

    for i in range(len(all_anchor_strides)):
        for j in range(all_anchor_strides[i].shape[0]):
            stride_tmp = int(all_anchor_strides[i][j])
            if stride_tmp in pred_bboxes.keys():
                pred_bboxes[stride_tmp] = torch.cat((pred_bboxes[stride_tmp], all_pos_bboxes[i][j].unsqueeze(0)))
                gt_bboxes[stride_tmp] = torch.cat((gt_bboxes[stride_tmp], all_pos_gt_bboxes[i][j].unsqueeze(0)))

    return pred_bboxes, gt_bboxes


def change_dim_pred_map(pred_map):
    pred_map = pred_map.view(pred_map.size()[0], pred_map.size()[1], -1)
    pred_map = pred_map.permute(0,2,1)
    pred_map = pred_map.view(pred_map.size()[0], pred_map.size()[1]*3, -1)
    return pred_map

class BaseYOLOHead(BaseDenseHead, BBoxTestMixin):
    """YOLOV3Head Paper link: https://arxiv.org/abs/1804.02767.

    Args:
        num_classes (int): The number of object classes (w/o background)
        in_channels (List[int]): Number of input channels per scale.
        out_channels (List[int]): The number of output channels per scale
            before the final 1x1 layer. Default: (1024, 512, 256).
        anchor_generator (dict): Config dict for anchor generator
        bbox_coder (dict): Config of bounding box coder.
        featmap_strides (List[int]): The stride of each scale.
            Should be in descending order. Default: (32, 16, 8).
        one_hot_smoother (float): Set a non-zero value to enable label-smooth
            Default: 0.
        conv_cfg (dict): Config dict for convolution layer. Default: None.
        norm_cfg (dict): Dictionary to construct and config norm layer.
            Default: dict(type='BN', requires_grad=True)
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='LeakyReLU', negative_slope=0.1).
        loss_cls (dict): Config of classification loss.
        loss_conf (dict): Config of confidence loss.
        loss_xy (dict): Config of xy coordinate loss.
        loss_wh (dict): Config of wh coordinate loss.
        train_cfg (dict): Training config of YOLOV3 head. Default: None.
        test_cfg (dict): Testing config of YOLOV3 head. Default: None.
    """

    def __init__(self,
                 num_classes,
                 in_channels,
                 out_channels=(1024, 512, 256),
                 anchor_generator=dict(
                     type='YOLOAnchorGenerator',
                     base_sizes=[[(116, 90), (156, 198), (373, 326)],
                                 [(30, 61), (62, 45), (59, 119)],
                                 [(10, 13), (16, 30), (33, 23)]],
                     strides=[32, 16, 8]),
                 bbox_coder=dict(type='YOLOBBoxCoder'),
                 featmap_strides=[32, 16, 8],
                 one_hot_smoother=0.,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN', requires_grad=True),
                 act_cfg=dict(type='LeakyReLU', negative_slope=0.1),
                 loss_cls=dict(
                     type='CrossEntropyLoss',
                     use_sigmoid=True,
                     loss_weight=1.0),
                 loss_conf=dict(
                     type='CrossEntropyLoss',
                     use_sigmoid=True,
                     loss_weight=1.0),
                 loss_xy=dict(
                     type='CrossEntropyLoss',
                     use_sigmoid=True,
                     loss_weight=1.0),
                 loss_wh=dict(type='MSELoss', loss_weight=1.0),
                 loss_bbox=None,
                 train_cfg=None,
                 test_cfg=None):
        super(BaseYOLOHead, self).__init__()
        # Check params
        assert (len(in_channels) == len(featmap_strides))

        self.num_classes = num_classes
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.featmap_strides = featmap_strides
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        self.debug = False
        if self.train_cfg:
            self.assigner = build_assigner(self.train_cfg.assigner)
            if hasattr(self.train_cfg, 'sampler'):
                sampler_cfg = self.train_cfg.sampler
            else:
                sampler_cfg = dict(type='PseudoSampler')  # yolo系列不需随机采样等操作
            self.sampler = build_sampler(sampler_cfg, context=self)
            self.debug = self.train_cfg.debug

        self.one_hot_smoother = one_hot_smoother

        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg

        self.bbox_coder = build_bbox_coder(bbox_coder)
        self.anchor_generator = build_anchor_generator(anchor_generator)

        self.loss_cls = build_loss(loss_cls)
        self.loss_conf = build_loss(loss_conf)
        self.loss_xy = build_loss(loss_xy)
        self.loss_wh = build_loss(loss_wh)

        self.valid_loss_bbox = False
        if loss_bbox is not None:
            self.valid_loss_bbox = True
            self.loss_bbox = build_loss(loss_bbox)

        # usually the numbers of anchors for each level are the same
        # except SSD detectors
        self.num_anchors = self.anchor_generator.num_base_anchors[0]
        assert len(
            self.anchor_generator.num_base_anchors) == len(featmap_strides)
        self._init_layers()

    @property
    def num_levels(self):
        return len(self.featmap_strides)

    @property
    def num_attrib(self):
        """int: number of attributes in pred_map, bboxes (4) +
        objectness (1) + num_classes"""

        return 5 + self.num_classes

    def _init_layers(self):
        raise NotImplementedError

    def init_weights(self):
        pass

    def forward(self, feats):
        """Forward features from the upstream network.

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

        Returns:
            tuple[Tensor]: A tuple of multi-level predication map, each is a
                4D-tensor of shape (batch_size, 5+num_classes, height, width).
        """

        raise NotImplementedError

    def get_bboxes(self, pred_maps, img_metas, cfg=None, rescale=False, with_nms=True):
        """Transform network output for a batch into bbox predictions.

        Args:
            pred_maps (list[Tensor]): Raw predictions for a batch of images.
            img_metas (list[dict]): Meta information of each image, e.g.,
                image size, scaling factor, etc.
            cfg (mmcv.Config | None): Test / postprocessing configuration,
                if None, test_cfg would be used. Default: None.
            rescale (bool): If True, return boxes in original image space.
                Default: False.
            with_nms (bool): If True, do nms before return boxes.
                Default: True.

        Returns:
            list[tuple[Tensor, Tensor]]: Each item in result_list is 2-tuple.
                The first item is an (n, 5) tensor, where the first 4 columns
                are bounding box positions (tl_x, tl_y, br_x, br_y) and the
                5-th column is a score between 0 and 1. The second item is a
                (n,) tensor where each item is the predicted class label of the
                corresponding box.
        """
        result_list = []
        num_levels = len(pred_maps)
        for img_id in range(len(img_metas)):
            pred_maps_list = [
                pred_maps[i][img_id].detach() for i in range(num_levels)
            ]
            scale_factor = img_metas[img_id]['scale_factor']
            if 'pad_param' in img_metas[img_id]:
                pad_param = img_metas[img_id]['pad_param']
            else:
                pad_param = None
            proposals = self._get_bboxes_single(pred_maps_list, scale_factor,
                                                cfg, rescale, pad_param, with_nms)
            result_list.append(proposals)
        return result_list

    def _get_bboxes_single(self,
                           pred_maps_list,
                           scale_factor,
                           cfg,
                           rescale=False,
                           pad_param=None,
                           with_nms=True):
        """Transform outputs for a single batch item into bbox predictions.

        Args:
            pred_maps_list (list[Tensor]): Prediction maps for different scales
                of each single image in the batch.
            scale_factor (ndarray): Scale factor of the image arrange as
                (w_scale, h_scale, w_scale, h_scale).
            cfg (mmcv.Config): Test / postprocessing configuration,
                if None, test_cfg would be used.
            rescale (bool): If True, return boxes in original image space.

        Returns:
            Tensor: Labeled boxes in shape (n, 5), where the first 4 columns
                are bounding box positions (tl_x, tl_y, br_x, br_y) and the
                5-th column is a score between 0 and 1.
        """
        cfg = self.test_cfg if cfg is None else cfg
        assert len(pred_maps_list) == self.num_levels
        multi_lvl_bboxes = []
        multi_lvl_cls_scores = []
        multi_lvl_conf_scores = []
        num_levels = len(pred_maps_list)
        featmap_sizes = [
            pred_maps_list[i].shape[-2:] for i in range(num_levels)
        ]
        multi_lvl_anchors = self.anchor_generator.grid_anchors(
            featmap_sizes, pred_maps_list[0][0].device)
        for i in range(self.num_levels):
            # get some key info for current scale
            pred_map = pred_maps_list[i]
            stride = self.featmap_strides[i]

            # (h, w, num_anchors*num_attrib) -> (h*w*num_anchors, num_attrib)
            pred_map = pred_map.permute(1, 2, 0).reshape(-1, self.num_attrib)

            pred_map[..., :2] = torch.sigmoid(pred_map[..., :2])
            bbox_pred = self.bbox_coder.decode(multi_lvl_anchors[i],
                                               pred_map[..., :4], stride)
            # conf and cls
            conf_pred = torch.sigmoid(pred_map[..., 4]).view(-1)
            cls_pred = torch.sigmoid(pred_map[..., 5:]).view(
                -1, self.num_classes)  # Cls pred one-hot.

            # Filtering out all predictions with conf < conf_thr
            conf_thr = cfg.get('conf_thr', -1)
            conf_inds = conf_pred.ge(conf_thr).nonzero().flatten()
            bbox_pred = bbox_pred[conf_inds, :]
            cls_pred = cls_pred[conf_inds, :]
            conf_pred = conf_pred[conf_inds]

            # Get top-k prediction
            nms_pre = cfg.get('nms_pre', -1)
            if 0 < nms_pre < conf_pred.size(0):
                _, topk_inds = conf_pred.topk(nms_pre)
                bbox_pred = bbox_pred[topk_inds, :]
                cls_pred = cls_pred[topk_inds, :]
                conf_pred = conf_pred[topk_inds]

            # Save the result of current scale
            multi_lvl_bboxes.append(bbox_pred)
            multi_lvl_cls_scores.append(cls_pred)
            multi_lvl_conf_scores.append(conf_pred)

        # Merge the results of different scales together
        multi_lvl_bboxes = torch.cat(multi_lvl_bboxes)
        multi_lvl_cls_scores = torch.cat(multi_lvl_cls_scores)
        multi_lvl_conf_scores = torch.cat(multi_lvl_conf_scores)

        if with_nms and multi_lvl_conf_scores.size(0) == 0:
            return torch.zeros((0, 5)), torch.zeros((0,))

        if rescale:
            # 先去掉pad，如果有的话
            if pad_param is not None:
                multi_lvl_bboxes -= multi_lvl_bboxes.new_tensor(
                    [pad_param[2], pad_param[0], pad_param[2], pad_param[0]])
            multi_lvl_bboxes /= multi_lvl_bboxes.new_tensor(scale_factor)

        # In mmdet 2.x, the class_id for background is num_classes.
        # i.e., the last column.
        padding = multi_lvl_cls_scores.new_zeros(multi_lvl_cls_scores.shape[0],
                                                 1)
        multi_lvl_cls_scores = torch.cat([multi_lvl_cls_scores, padding],
                                         dim=1)

        if with_nms:
            det_bboxes, det_labels = multiclass_nms(
                multi_lvl_bboxes,
                multi_lvl_cls_scores,
                cfg.score_thr,
                cfg.nms,
                cfg.max_per_img,
                score_factors=multi_lvl_conf_scores)
            return det_bboxes, det_labels
        else:
            return (multi_lvl_bboxes, multi_lvl_cls_scores,
                    multi_lvl_conf_scores)

    def loss(self,
             pred_maps,
             gt_bboxes,
             gt_labels,
             img_metas,
             gt_bboxes_ignore=None):
        """Compute loss of the head.

        Args:
            pred_maps (list[Tensor]): Prediction map for each scale level,
                shape (N, num_anchors * num_attrib, H, W)
            gt_bboxes (list[Tensor]): Ground truth bboxes for each image with
                shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format.
            gt_labels (list[Tensor]): class indices corresponding to each box
            img_metas (list[dict]): Meta information of each image, e.g.,
                image size, scaling factor, etc.
            gt_bboxes_ignore (None | list[Tensor]): specify which bounding
                boxes can be ignored when computing the loss.

        Returns:
            dict[str, Tensor]: A dictionary of loss components.
        """
        num_imgs = len(img_metas)
        device = pred_maps[0][0].device

        featmap_sizes = [
            pred_maps[i].shape[-2:] for i in range(self.num_levels)
        ]

        multi_level_anchors = self.anchor_generator.grid_anchors(
            featmap_sizes, device)
        anchor_list = [multi_level_anchors for _ in range(num_imgs)]

        responsible_flag_list = []
        for img_id in range(len(img_metas)):
            # 计算当前图片中的bbox中心在特征图上面的位置,以flag形式表示
            responsible_flag_list.append(
                self.anchor_generator.responsible_flags(
                    featmap_sizes, gt_bboxes[img_id], device))

        #### target_maps_list [batch_per, feat_size*feat_size*anchor_num_per_feat, 7], ex [8, 507, 7]
        #### neg_maps_list  [batch_per, feat_size*feat_size*anchor_num_per_feat], ex [8, 507]
        target_maps_list, neg_maps_list  = self.get_targets(
            anchor_list, responsible_flag_list, gt_bboxes, gt_labels, img_metas)

        ####### loss: [tensor0( feat level0 ), tensor1(feat_level1), tensor2]
        ####### pred_maps[0]:  21 * 13*13
        ####### target_maps_list[0]: batch_size*(13*13)*7, 7=4(xywh) +1(is obj) +2(class_num)
        ####### neg_maps_list[0]: batch_size*(13*13)
        losses_cls, losses_conf, losses_xy, losses_wh = multi_apply(
            self.loss_single, pred_maps, target_maps_list, neg_maps_list)

        if self.valid_loss_bbox is True:

            multi_lvl_anchors = self.anchor_generator.grid_anchors(featmap_sizes, pred_maps[0].device)
            pred_bboxes_encode_list = []
            target_bboxes_encode_list = []

            for i in range(self.num_levels):
                # get some key info for current scale
                pred_map = pred_maps[i]
                stride = self.featmap_strides[i]

                # (h, w, num_anchors*num_attrib) -> (h*w*num_anchors, num_attrib)
                pred_map = pred_map.permute(0, 2, 3, 1).reshape(pred_map.size()[0], -1, self.num_attrib)
                pred_map[..., :2] = torch.sigmoid(pred_map[..., :2])

                multi_level_anchors_batch = multi_lvl_anchors[i][None,:,:].repeat((pred_map.size()[0], 1, 1))
                bbox_pred = self.bbox_coder.decode(multi_level_anchors_batch, pred_map[..., :4], stride)

                bbox_target = self.bbox_coder.decode(multi_level_anchors_batch, torch.sigmoid(target_maps_list[i][..., :4]), stride)

                pred_bboxes_encode_list.append(bbox_pred)
                target_bboxes_encode_list.append(bbox_target)

            weights_list = []
            for i in range(self.num_levels):
                weights_list.append(target_maps_list[i][..., 4][:, :, None].repeat((1, 1, 4)))
                # if torch.where(target_maps_list[0][..., 4][:,:,None].repeat((1,1,4))[...,0]!=0)[0].size(0) ==0:
                #     factor =0
                # else:
                #     factor = float(weights_list[i].size(0)*weights_list[i].size(1))/float(torch.where(target_maps_list[0][..., 4][:,:,None].repeat((1,1,4))[...,0]!=0)[0].size(0))
                # weights_list[i] *= factor

            # weights_list = [target_maps_list[0][..., 4][:,:,None].repeat((1,1,4)),
            #                 target_maps_list[1][..., 4][:,:,None].repeat((1,1,4)), target_maps_list[2][..., 4][:,:,None].repeat((1,1,4))]
            # losses_bbox = multi_apply(self.loss_bbox,  pred_bboxes_encode_list, target_bboxes_encode_list, weights_list)

            losses_bbox = []
            losses_bbox.append(
                self.loss_bbox(target_bboxes_encode_list[0].view([-1, 4]), pred_bboxes_encode_list[0].view([-1, 4]), weights_list[0].view([-1, 4])))
            losses_bbox.append(
                self.loss_bbox(target_bboxes_encode_list[1].view([-1, 4]), pred_bboxes_encode_list[1].view([-1, 4]), weights_list[1].view([-1, 4])))
            losses_bbox.append(
                self.loss_bbox(target_bboxes_encode_list[2].view([-1, 4]), pred_bboxes_encode_list[2].view([-1, 4]), weights_list[2].view([-1, 4])))

            return dict(loss_cls=losses_cls,
                loss_conf=losses_conf,
                loss_xy=losses_xy,
                loss_wh=losses_wh,
                loss_bbox=losses_bbox)

        return dict(
            loss_cls=losses_cls,
            loss_conf=losses_conf,
            loss_xy=losses_xy,
            loss_wh=losses_wh)

    def loss_single(self, pred_map, target_map, neg_map):
        """Compute loss of a single image from a batch.

        Args:
            pred_map (Tensor): Raw predictions for a single level.
            target_map (Tensor): The Ground-Truth target for a single level.
            neg_map (Tensor): The negative masks for a single level.

        Returns:
            tuple:
                loss_cls (Tensor): Classification loss.
                loss_conf (Tensor): Confidence loss.
                loss_xy (Tensor): Regression loss of x, y coordinate.
                loss_wh (Tensor): Regression loss of w, h coordinate.
        """

        num_imgs = len(pred_map)
        pred_map = pred_map.permute(0, 2, 3, 1).reshape(num_imgs, -1, self.num_attrib)
        neg_mask = neg_map.float()  # 忽略位置图
        pos_mask = target_map[..., 4]  # confidence，也是衡量该位置有没有obj的标志
        pos_and_neg_mask = neg_mask + pos_mask
        pos_mask = pos_mask.unsqueeze(dim=-1)
        if torch.max(pos_and_neg_mask) > 1.:  # 说明有些正样本区域是要忽略的
            warnings.warn('There is overlap between pos and neg sample.')
            pos_and_neg_mask = pos_and_neg_mask.clamp(min=0., max=1.)

        pred_xy = pred_map[..., :2]
        pred_wh = pred_map[..., 2:4]
        pred_conf = pred_map[..., 4]
        pred_label = pred_map[..., 5:]

        target_xy = target_map[..., :2]
        target_wh = target_map[..., 2:4]
        target_conf = target_map[..., 4]
        target_label = target_map[..., 5:]

        #
        # import pdb
        # pdb.set_trace()
        loss_cls = self.loss_cls(pred_label, target_label, weight=pos_mask)
        loss_conf = self.loss_conf(
            pred_conf, target_conf, weight=pos_and_neg_mask)
        loss_xy = self.loss_xy(pred_xy, target_xy, weight=pos_mask)
        loss_wh = self.loss_wh(pred_wh, target_wh, weight=pos_mask)

        return loss_cls, loss_conf, loss_xy, loss_wh

    def get_targets(self, anchor_list, responsible_flag_list, gt_bboxes_list,
                    gt_labels_list, img_metas):
        """Compute target maps for anchors in multiple images.

        Args:
            anchor_list (list[list[Tensor]]): Multi level anchors of each
                image. The outer list indicates images, and the inner list
                corresponds to feature levels of the image. Each element of
                the inner list is a tensor of shape (num_total_anchors, 4).
            responsible_flag_list (list[list[Tensor]]): Multi level responsible
                flags of each image. Each element is a tensor of shape
                (num_total_anchors, )
            gt_bboxes_list (list[Tensor]): Ground truth bboxes of each image.
            gt_labels_list (list[Tensor]): Ground truth labels of each box.

        Returns:
            tuple: Usually returns a tuple containing learning targets.
                - target_map_list (list[Tensor]): Target map of each level.
                - neg_map_list (list[Tensor]): Negative map of each level.
        """
        num_imgs = len(anchor_list)

        # anchor number of multi levels 每一层anchor的个数，例如输出是10x10，那么10x10x3=300
        num_level_anchors = [anchors.size(0) for anchors in anchor_list[0]]

        results = multi_apply(self._get_targets_single, anchor_list,
                              responsible_flag_list, gt_bboxes_list,
                              gt_labels_list, img_metas, num_level_anchors=num_level_anchors)

        all_target_maps, all_neg_maps, all_pos_bboxes, all_pos_gt_bboxes, all_anchor_strides = results

        assert num_imgs == len(all_target_maps) == len(all_neg_maps)
        # all_target_maps是图片级别list结构，假设有8张图片则len()=8，需要转化为fpn级别的输出，才方便计算Loss,假设有三层输出，则len()=3
        #### target_maps_list: [batch_size, 3*13*13, 7]
        target_maps_list = images_to_levels(all_target_maps, num_level_anchors)
        neg_maps_list = images_to_levels(all_neg_maps, num_level_anchors)
        # pred_bboxes_list, gt_bboxes_list = bbox_to_levels(all_pos_bboxes, all_pos_gt_bboxes, all_anchor_strides)

        return target_maps_list, neg_maps_list

    def _get_targets_single(self, anchors, responsible_flags, gt_bboxes,
                            gt_labels, img_meta, num_level_anchors=None):
        """Generate matching bounding box prior and converted GT.

        Args:
            anchors (list[Tensor]): Multi-level anchors of the image.
            responsible_flags (list[Tensor]): Multi-level responsible flags of
                anchors
            gt_bboxes (Tensor): Ground truth bboxes of single image.
            gt_labels (Tensor): Ground truth labels of single image.

        Returns:
            tuple:
                target_map (Tensor): Predication target map of each
                    scale level, shape (num_total_anchors,
                    5+num_classes)
                neg_map (Tensor): Negative map of each scale level,
                    shape (num_total_anchors,)
        """

        anchor_strides = []
        for i in range(len(anchors)):
            anchor_strides.append(
                torch.tensor(self.featmap_strides[i],
                             device=gt_bboxes.device).repeat(len(anchors[i])))
        concat_anchors = torch.cat(anchors)  # 三个输出层的anchor合并
        concat_responsible_flags = torch.cat(responsible_flags)

        anchor_strides = torch.cat(anchor_strides)
        assert len(anchor_strides) == len(concat_anchors) == \
               len(concat_responsible_flags)
        assign_result = self.assigner.assign(concat_anchors,
                                             concat_responsible_flags,
                                             gt_bboxes)

        if self.debug:
            # 统计下正样本个数
            print('----anchor分配正负样本后，正样本anchor可视化，白色bbox是gt----')
            gt_inds = assign_result.gt_inds  # 0 1 -1
            index = gt_inds > 0
            gt_inds = gt_inds[index]
            gt_anchors = concat_anchors[index]
            print('单张图片中正样本anchor个数', len(gt_inds))
            if num_level_anchors is None:  # 不分层显示
                show_pos_anchor(img_meta, gt_anchors, gt_bboxes)
            else:
                imgs = []
                count = 0
                for num_level in num_level_anchors:
                    gt_inds = assign_result.gt_inds[count:count + num_level]
                    anchor = concat_anchors[count:count + num_level]
                    count += num_level
                    index = gt_inds > 0
                    gt_anchor = anchor[index]
                    img = show_pos_anchor(img_meta, gt_anchor, gt_bboxes, is_show=False)
                    imgs.append(img)
                print('从小特征图到大特征图顺序显示')
                show_img(imgs)

        # 相当于没有，只是为了不报错
        sampling_result = self.sampler.sample(assign_result, concat_anchors,
                                              gt_bboxes)

        # 转化为最终计算Loss所需要的格式
        target_map = concat_anchors.new_zeros(
            concat_anchors.size(0), self.num_attrib)  # 5+class_count

        # 正样本位置anchor bbox；对应的gt bbox；strides
        # target_map前4个是xywh在特征图尺度上面的转化后的label
        target_map[sampling_result.pos_inds, :4] = self.bbox_coder.encode(
            sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes,
            anchor_strides[sampling_result.pos_inds])

        target_map[sampling_result.pos_inds, 4] = 1  # confidence label

        gt_labels_one_hot = F.one_hot(
            gt_labels, num_classes=self.num_classes).float()
        if self.one_hot_smoother != 0:  # label smooth
            gt_labels_one_hot = gt_labels_one_hot * (
                    1 - self.one_hot_smoother
            ) + self.one_hot_smoother / self.num_classes
        target_map[sampling_result.pos_inds, 5:] = gt_labels_one_hot[
            sampling_result.pos_assigned_gt_inds]  # class one hot label

        neg_map = concat_anchors.new_zeros(
            concat_anchors.size(0), dtype=torch.uint8)
        neg_map[sampling_result.neg_inds] = 1

        return target_map, neg_map, \
               sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes, anchor_strides[sampling_result.pos_inds]

    def aug_test(self, feats, img_metas, rescale=False):
        """Test function with test time augmentation.

        Args:
            feats (list[Tensor]): the outer list indicates test-time
                augmentations and inner Tensor should have a shape NxCxHxW,
                which contains features for all images in the batch.
            img_metas (list[list[dict]]): the outer list indicates test-time
                augs (multiscale, flip, etc.) and the inner list indicates
                images in a batch. each dict has image information.
            rescale (bool, optional): Whether to rescale the results.
                Defaults to False.

        Returns:
            list[ndarray]: bbox results of each class
        """
        return self.aug_test_bboxes(feats, img_metas, rescale=rescale)
