YOLO-Tutorial-v2/odlab/models/detectors/detr/detr.py

import math
import torch
import torch.nn as nn

from ...backbone          import build_backbone
from ...basic.mlp         import MLP
from ...basic.conv        import BasicConv, UpSampleWrapper
from ...basic.transformer import TransformerEncoder, PlainDETRTransformer, get_clones

from utils.misc import multiclass_nms


# DETR
class DETR(nn.Module):
    def __init__(self,
                 cfg,
                 num_classes = 80,
                 conf_thresh = 0.1,
                 nms_thresh  = 0.5,
                 topk        = 300,
                 use_nms     = False,
                 ca_nms      = False,
                 ):
        super().__init__()
        # ---------------- Basic setting ----------------
        self.stride = cfg['out_stride']
        self.upsample_factor = cfg['max_stride'] // cfg['out_stride']
        self.num_classes = num_classes
        ## Transformer parameters
        self.num_queries_one2one = cfg['num_queries_one2one']
        self.num_queries_one2many = cfg['num_queries_one2many']
        self.num_queries = self.num_queries_one2one + self.num_queries_one2many
        ## Post-process parameters
        self.ca_nms = ca_nms
        self.use_nms = use_nms
        self.num_topk = topk
        self.nms_thresh = nms_thresh
        self.conf_thresh = conf_thresh

        # ---------------- Network setting ----------------
        ## Backbone Network
        self.backbone, feat_dims = build_backbone(cfg)

        ## Input projection
        self.input_proj = BasicConv(feat_dims[-1], cfg['hidden_dim'], kernel_size=1, act_type=None, norm_type='GN')

        ## Transformer Encoder
        self.transformer_encoder = TransformerEncoder(d_model    = cfg['hidden_dim'],
                                                      num_heads  = cfg['en_num_heads'],
                                                      num_layers = cfg['en_num_layers'],
                                                      ffn_dim    = cfg['en_ffn_dim'],
                                                      dropout    = cfg['en_dropout'],
                                                      act_type   = cfg['en_act'],
                                                      pre_norm   = cfg['en_pre_norm'],
                                                      )

        ## Upsample layer
        self.upsample = UpSampleWrapper(cfg['hidden_dim'], self.upsample_factor)
        
        ## Output projection
        self.output_proj = BasicConv(cfg['hidden_dim'], cfg['hidden_dim'], kernel_size=3, padding=1, act_type='silu', norm_type='BN')
        
        ## Transformer
        self.query_embed = nn.Embedding(self.num_queries, cfg['hidden_dim'])
        self.transformer = PlainDETRTransformer(d_model             = cfg['hidden_dim'],
                                                num_heads           = cfg['de_num_heads'],
                                                ffn_dim             = cfg['de_ffn_dim'],
                                                dropout             = cfg['de_dropout'],
                                                act_type            = cfg['de_act'],
                                                pre_norm            = cfg['de_pre_norm'],
                                                rpe_hidden_dim      = cfg['rpe_hidden_dim'],
                                                feature_stride      = cfg['out_stride'],
                                                num_layers          = cfg['de_num_layers'],
                                                use_checkpoint      = cfg['use_checkpoint'],
                                                num_queries_one2one = cfg['num_queries_one2one'],
                                                num_queries_one2many    = cfg['num_queries_one2many'],
                                                proposal_feature_levels = cfg['proposal_feature_levels'],
                                                proposal_in_stride      = cfg['out_stride'],
                                                proposal_tgt_strides    = cfg['proposal_tgt_strides'],
                                                return_intermediate = True,
                                                )
    
        ## Detect Head
        class_embed = nn.Linear(cfg['hidden_dim'], num_classes)
        bbox_embed = MLP(cfg['hidden_dim'], cfg['hidden_dim'], 4, 3)

        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        class_embed.bias.data = torch.ones(num_classes) * bias_value
        nn.init.constant_(bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(bbox_embed.layers[-1].bias.data, 0)

        self.class_embed = get_clones(class_embed, cfg['de_num_layers'] + 1)
        self.bbox_embed  = get_clones(bbox_embed, cfg['de_num_layers'] + 1)
        nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)

        self.transformer.decoder.bbox_embed = self.bbox_embed
        self.transformer.decoder.class_embed = self.class_embed

    def get_posembed(self, d_model, mask, temperature=10000, normalize=False):
        not_mask = ~mask
        scale = 2 * torch.pi
        num_pos_feats = d_model // 2

        # -------------- Generate XY coords --------------
        ## [B, H, W]
        y_embed = not_mask.cumsum(1, dtype=torch.float32)
        x_embed = not_mask.cumsum(2, dtype=torch.float32)
        ## Normalize coords
        if normalize:
            y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + 1e-6)
            x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + 1e-6)
        else:
            y_embed = y_embed - 0.5
            x_embed = x_embed - 0.5
        # [H, W] -> [B, H, W, 2]
        pos = torch.stack([x_embed, y_embed], dim=-1)

        # -------------- Sine-PosEmbedding --------------
        dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos.device)
        dim_t_ = torch.div(dim_t, 2, rounding_mode='floor') / num_pos_feats
        dim_t = temperature ** (2 * dim_t_)

        x_embed = pos[..., 0] * scale
        y_embed = pos[..., 1] * scale
        pos_x = x_embed[..., None] / dim_t
        pos_y = y_embed[..., None] / dim_t
        pos_x = torch.stack((pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=-1).flatten(-2)
        pos_y = torch.stack((pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=-1).flatten(-2)
        pos_embed = torch.cat((pos_y, pos_x), dim=-1)
        
        # [B, H, W, C] -> [B, C, H, W]
        pos_embed = pos_embed.permute(0, 3, 1, 2)
        
        return pos_embed

    def post_process(self, box_pred, cls_pred):
        # Top-k select
        cls_pred = cls_pred[0].flatten().sigmoid_()
        box_pred = box_pred[0]

        # Keep top k top scoring indices only.
        num_topk = min(self.num_topk, box_pred.size(0))

        # Topk candidates
        predicted_prob, topk_idxs = cls_pred.sort(descending=True)
        topk_scores = predicted_prob[:num_topk]
        topk_idxs = topk_idxs[:self.num_topk]

        # Filter out the proposals with low confidence score
        keep_idxs = topk_scores > self.conf_thresh
        topk_scores = topk_scores[keep_idxs]
        topk_idxs = topk_idxs[keep_idxs]
        topk_box_idxs = torch.div(topk_idxs, self.num_classes, rounding_mode='floor')

        ## Top-k results
        topk_labels = topk_idxs % self.num_classes
        topk_bboxes = box_pred[topk_box_idxs]

        topk_scores = topk_scores.cpu().numpy()
        topk_labels = topk_labels.cpu().numpy()
        topk_bboxes = topk_bboxes.cpu().numpy()

        # nms
        if self.use_nms:
            topk_scores, topk_labels, topk_bboxes = multiclass_nms(
                topk_scores, topk_labels, topk_bboxes, self.nms_thresh, self.num_classes, self.nms_class_agnostic)

        return topk_bboxes, topk_scores, topk_labels

    def resize_mask(self, src, mask=None):
        bs, c, h, w = src.shape
        if mask is not None:
            # [B, H, W]
            mask = nn.functional.interpolate(mask[None].float(), size=[h, w]).bool()[0]
        else:
            mask = torch.zeros([bs, h, w], device=src.device, dtype=torch.bool)

        return mask
    
    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord, outputs_coord_old, outputs_deltas):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [
            {"pred_logits": a, "pred_boxes": b, "pred_boxes_old": c, "pred_deltas": d, }
            for a, b, c, d in zip(outputs_class[:-1], outputs_coord[:-1], outputs_coord_old[:-1], outputs_deltas[:-1])
        ]

    def inference_single_image(self, x):
        # ----------- Image Encoder -----------
        pyramid_feats = self.backbone(x)
        src = self.input_proj(pyramid_feats[-1])
        src = self.transformer_encoder(src)
        src = self.upsample(src)
        src = self.output_proj(src)

        # ----------- Prepare inputs for Transformer -----------
        mask = self.resize_mask(src)
        pos_embed = self.get_posembed(src.shape[1], mask, normalize=False)
        query_embeds = self.query_embed.weight[:self.num_queries_one2one]
        self_attn_mask = None

        # -----------Transformer -----------
        (
            hs,
            init_reference,
            inter_references,
            _,
            _,
            _,
            _,
            max_shape
        ) = self.transformer(src, mask, pos_embed, query_embeds, self_attn_mask)

        # ----------- Process outputs -----------
        outputs_classes_one2one = []
        outputs_coords_one2one = []
        outputs_deltas_one2one = []

        for lid in range(hs.shape[0]):
            if lid == 0:
                reference = init_reference
            else:
                reference = inter_references[lid - 1]
            outputs_class = self.class_embed[lid](hs[lid])
            tmp = self.bbox_embed[lid](hs[lid])
            outputs_coord = self.transformer.decoder.delta2bbox(reference, tmp, max_shape)  # xyxy

            outputs_classes_one2one.append(outputs_class[:, :self.num_queries_one2one])
            outputs_coords_one2one.append(outputs_coord[:, :self.num_queries_one2one])
            outputs_deltas_one2one.append(tmp[:, :self.num_queries_one2one])

        outputs_classes_one2one = torch.stack(outputs_classes_one2one)
        outputs_coords_one2one = torch.stack(outputs_coords_one2one)

        # ------------ Post process ------------
        cls_pred = outputs_classes_one2one[-1]
        box_pred = outputs_coords_one2one[-1]
        
        # post-process
        bboxes, scores, labels = self.post_process(box_pred, cls_pred)
        # normalize bbox
        bboxes[..., 0::2] /= x.shape[-1]
        bboxes[..., 1::2] /= x.shape[-2]
        bboxes = bboxes.clip(0., 1.)

        return bboxes, scores, labels
        
    def forward(self, x, src_mask=None, targets=None):
        if not self.training:
            return self.inference_single_image(x)

        # ----------- Image Encoder -----------
        pyramid_feats = self.backbone(x)
        src = self.input_proj(pyramid_feats[-1])
        src = self.transformer_encoder(src)
        src = self.upsample(src)
        src = self.output_proj(src)

        # ----------- Prepare inputs for Transformer -----------
        mask = self.resize_mask(src, src_mask)
        pos_embed = self.get_posembed(src.shape[1], mask, normalize=False)
        query_embeds = self.query_embed.weight
        self_attn_mask = torch.zeros(
            [self.num_queries, self.num_queries, ]).bool().to(src.device)
        self_attn_mask[self.num_queries_one2one:, 0: self.num_queries_one2one, ] = True
        self_attn_mask[0: self.num_queries_one2one, self.num_queries_one2one:, ] = True

        # -----------Transformer -----------
        (
            hs,
            init_reference,
            inter_references,
            enc_outputs_class,
            enc_outputs_coord_unact,
            enc_outputs_delta,
            output_proposals,
            max_shape
        ) = self.transformer(src, mask, pos_embed, query_embeds, self_attn_mask)

        # ----------- Process outputs -----------
        outputs_classes_one2one = []
        outputs_coords_one2one = []
        outputs_coords_old_one2one = []
        outputs_deltas_one2one = []

        outputs_classes_one2many = []
        outputs_coords_one2many = []
        outputs_coords_old_one2many = []
        outputs_deltas_one2many = []

        for lid in range(hs.shape[0]):
            if lid == 0:
                reference = init_reference
            else:
                reference = inter_references[lid - 1]
            outputs_class = self.class_embed[lid](hs[lid])
            tmp = self.bbox_embed[lid](hs[lid])
            outputs_coord = self.transformer.decoder.box_xyxy_to_cxcywh(
                self.transformer.decoder.delta2bbox(reference, tmp, max_shape))

            outputs_classes_one2one.append(outputs_class[:, 0: self.num_queries_one2one])
            outputs_classes_one2many.append(outputs_class[:, self.num_queries_one2one:])

            outputs_coords_one2one.append(outputs_coord[:, 0: self.num_queries_one2one])
            outputs_coords_one2many.append(outputs_coord[:, self.num_queries_one2one:])

            outputs_coords_old_one2one.append(reference[:, :self.num_queries_one2one])
            outputs_coords_old_one2many.append(reference[:, self.num_queries_one2one:])
            outputs_deltas_one2one.append(tmp[:, :self.num_queries_one2one])
            outputs_deltas_one2many.append(tmp[:, self.num_queries_one2one:])

        outputs_classes_one2one = torch.stack(outputs_classes_one2one)
        outputs_coords_one2one = torch.stack(outputs_coords_one2one)

        outputs_classes_one2many = torch.stack(outputs_classes_one2many)
        outputs_coords_one2many = torch.stack(outputs_coords_one2many)

        out = {
            "pred_logits": outputs_classes_one2one[-1],
            "pred_boxes": outputs_coords_one2one[-1],
            "pred_logits_one2many": outputs_classes_one2many[-1],
            "pred_boxes_one2many": outputs_coords_one2many[-1],

            "pred_boxes_old": outputs_coords_old_one2one[-1],
            "pred_deltas": outputs_deltas_one2one[-1],
            "pred_boxes_old_one2many": outputs_coords_old_one2many[-1],
            "pred_deltas_one2many": outputs_deltas_one2many[-1],
        }

        out["aux_outputs"] = self._set_aux_loss(
            outputs_classes_one2one, outputs_coords_one2one, outputs_coords_old_one2one, outputs_deltas_one2one
        )
        out["aux_outputs_one2many"] = self._set_aux_loss(
            outputs_classes_one2many, outputs_coords_one2many, outputs_coords_old_one2many, outputs_deltas_one2many
        )

        out["enc_outputs"] = {
            "pred_logits": enc_outputs_class,
            "pred_boxes": enc_outputs_coord_unact,
            "pred_boxes_old": output_proposals,
            "pred_deltas": enc_outputs_delta,
        }

        return out