import copy
import itertools
from collections import defaultdict
from typing import List, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from detectron2.data.detection_utils import convert_image_to_rgb
from detectron2.structures import Boxes, ImageList, Instances, pairwise_iou
from detectron2.utils.events import get_event_storage
from detrex.data.datasets import clean_words_or_phrase, template_meta
from detrex.layers import MLP, ClassEmbed, box_cxcywh_to_xyxy, box_xyxy_to_cxcywh
from detrex.utils import inverse_sigmoid
from .utils import cosine_loss

from torch import Tensor
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union
import warnings
from transformers.modeling_outputs import ModelOutput


@dataclass
class BaseModelOutputWithPoolingAndCrossAttentions(ModelOutput):
    last_hidden_state: torch.FloatTensor = None
    pooler_output: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None

class BertlangEmbedding(nn.Module):
    def __init__(self, language_backbone, prompt_mode, category_names, num_classes, num_prompt, eval_only):
        super().__init__()
        self.tokenizer = language_backbone.tokenizer
        self.config = language_backbone.lang_model.config
        self.embeddings = language_backbone.lang_model.embeddings
        self.dtype = language_backbone.lang_model.dtype
        
        # prompt
        self.n_ctx = 12
        self.text_num_prompt = 3
        self.ctx_dim = 768
        self.init_flag = False
        self.prompt_prefix = ' '.join(['x'] * self.n_ctx * self.text_num_prompt) # prompt基本模版，x * 上下文长度 * 单prompt个数

        self.num_classes = num_classes

        self.prompt_mode = prompt_mode
        if prompt_mode == "spbank":
            text_prompt = torch.empty(num_prompt, self.n_ctx, self.ctx_dim, dtype=torch.float).cuda()
            nn.init.normal_(text_prompt, std=0.02)
            self.text_prompt = nn.Parameter(text_prompt)
            
            # import pdb; pdb.set_trace()
            self.eval_only = eval_only
            if not self.eval_only:
                prompts = [self.prompt_prefix + ' ' + name + '.' for name in category_names]
                tokenized_batch = self.tokenizer(prompts, return_mask=True)
                input_ids = tokenized_batch["input_ids"]

                input_shape = input_ids.size()
                device = input_ids.device
                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
                
                with torch.no_grad():
                    pre_embedding_output = self.embeddings(
                        input_ids=input_ids,
                        position_ids=None,
                        token_type_ids=token_type_ids,
                        inputs_embeds=None,
                        past_key_values_length=0,
                    )           
                self.register_buffer('token_prefix', pre_embedding_output[:, :1, :])  
                self.register_buffer('token_suffix', pre_embedding_output[:, 1+(self.n_ctx * self.text_num_prompt):, :])  
            


    def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor:
        """
        Invert an attention mask (e.g., switches 0. and 1.).

        Args:
            encoder_attention_mask (`torch.Tensor`): An attention mask.

        Returns:
            `torch.Tensor`: The inverted attention mask.
        """
        if encoder_attention_mask.dim() == 3:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
        if encoder_attention_mask.dim() == 2:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
        # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
        # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
        # /transformer/transformer_layers.py#L270
        # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
        # encoder_extended_attention_mask.transpose(-1, -2))
        encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
        encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min

        return encoder_extended_attention_mask
    
    @staticmethod
    def create_extended_attention_mask_for_decoder(input_shape, attention_mask, device=None):
        if device is not None:
            warnings.warn(
                "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning
            )
        else:
            device = attention_mask.device
        batch_size, seq_length = input_shape
        seq_ids = torch.arange(seq_length, device=device)
        causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
        # in case past_key_values are used we need to add a prefix ones mask to the causal mask
        # causal and attention masks must have same type with pytorch version < 1.3
        causal_mask = causal_mask.to(attention_mask.dtype)

        if causal_mask.shape[1] < attention_mask.shape[1]:
            prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
            causal_mask = torch.cat(
                [
                    torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),
                    causal_mask,
                ],
                axis=-1,
            )

        extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
        return extended_attention_mask
    
    def get_extended_attention_mask(
        self, attention_mask: Tensor, input_shape: Tuple[int], device: torch.device = None, dtype: torch.float = None 
    ) -> Tensor:
        """
        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.

        Arguments:
            attention_mask (`torch.Tensor`):
                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
            input_shape (`Tuple[int]`):
                The shape of the input to the model.

        Returns:
            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
        """
        if dtype is None:
            dtype = self.dtype

        if not (attention_mask.dim() == 2 and self.config.is_decoder):
            # show warning only if it won't be shown in `create_extended_attention_mask_for_decoder`
            if device is not None:
                warnings.warn(
                    "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning
                )
        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]
        elif attention_mask.dim() == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length]
            # - if the model is a decoder, apply a causal mask in addition to the padding mask
            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
            # self.config.is_decoder == False
            if self.config.is_decoder:
                extended_attention_mask = self.create_extended_attention_mask_for_decoder(
                    input_shape, attention_mask, device
                )
            else:
                extended_attention_mask = attention_mask[:, None, None, :]
        else:
            raise ValueError(
                f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
            )

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and the dtype's smallest value for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = extended_attention_mask.to(dtype=dtype)  # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
        return extended_attention_mask

    def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
        """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""
        if head_mask.dim() == 1:
            head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
        elif head_mask.dim() == 2:
            head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
        assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
        head_mask = head_mask.to(dtype=self.dtype)  # switch to float if need + fp16 compatibility
        return head_mask
    
    def get_head_mask(
        self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False
    ) -> Tensor:
        """
        Prepare the head mask if needed.

        Args:
            head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*):
                The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard).
            num_hidden_layers (`int`):
                The number of hidden layers in the model.
            is_attention_chunked (`bool`, *optional*, defaults to `False`):
                Whether or not the attentions scores are computed by chunks or not.

        Returns:
            `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with
            `[None]` for each layer.
        """
        if head_mask is not None:
            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
            if is_attention_chunked is True:
                head_mask = head_mask.unsqueeze(-1)
        else:
            head_mask = [None] * num_hidden_layers

        return head_mask
    
    def forward(
        self,
        # input_ids: Optional[torch.Tensor] = None,
        input: Optional[List[str]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        indices: Optional[torch.Tensor] = None,
        training: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        
        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        # tokenize
        # import pdb; pdb.set_trace()
        prompt_prefix = self.prompt_prefix  # 固定模板
        prompts = [prompt_prefix + ' ' + name + '.' for name in input]
        tokenized_batch = self.tokenizer(prompts, return_mask=True)
        input_ids = tokenized_batch["input_ids"].cuda()


        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        batch_size, seq_length = input_shape
        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)

        if token_type_ids is None:
            if hasattr(self.embeddings, "token_type_ids"):
                buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
                token_type_ids = buffered_token_type_ids_expanded
            else:
                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        if not training:
            # import pdb; pdb.set_trace()
            with torch.no_grad():    
                embedding_output = self.embeddings(
                    input_ids=input_ids,
                    position_ids=position_ids,
                    token_type_ids=token_type_ids,
                    inputs_embeds=inputs_embeds,
                    past_key_values_length=past_key_values_length,
                )
            self.register_buffer('test_token_prefix', embedding_output[:, :1, :]) # SOS, [n_cls, 1, ctx_dim]
            self.register_buffer('test_token_suffix', embedding_output[:, 1+(self.n_ctx * self.text_num_prompt):,:]) # CLS, EOS, [n_cls, -1, ctx_dim]

        if self.prompt_mode == "spbank":
            # import pdb; pdb.set_trace()
            if not training:
                batch = 1
                prefix = self.test_token_prefix.unsqueeze(0).repeat(batch,1,1,1) # 32,80,1,768
                suffix = self.test_token_suffix.unsqueeze(0).repeat(batch,1,1,1)
            else:
                batch = len(input)//self.num_classes
                prefix = self.token_prefix.unsqueeze(0).repeat(batch,1,1,1) # 32,80,1,768
                suffix = self.token_suffix.unsqueeze(0).repeat(batch,1,1,1)

            prompt = self.text_prompt[indices].view(batch, self.n_ctx * self.text_num_prompt, self.ctx_dim)
            if training:
                prompt = prompt.unsqueeze(1).repeat(1, self.num_classes, 1, 1)# 32, 80,36,768
            else:
                prompt = prompt.unsqueeze(1).repeat(1, len(input), 1, 1) 

            embedding_output = torch.cat([prefix, prompt, suffix], dim=2) # [16, 7, 44, 768])
            # import pdb; pdb.set_trace()
            embedding_output = embedding_output.view(-1, embedding_output.shape[2],embedding_output.shape[3]) # [112, 44, 768]

            
            nc_text_features = self.text_prompt.mean(dim=1)  # [16, 256]
            nc_text_features = nc_text_features / nc_text_features.norm(dim=-1, keepdim=True) # [16, 256]
            sim_matrix = nc_text_features @ nc_text_features.T  # [10, 10]
            off_diag_mask = ~torch.eye(self.text_prompt.shape[0], dtype=torch.bool, device=sim_matrix.device)
            loss_orth = sim_matrix[off_diag_mask].abs().mean()
            
    
        return  (embedding_output, extended_attention_mask, head_mask, encoder_hidden_states, encoder_extended_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict,
                input_ids, loss_orth)

        
class BertlangEncoder(nn.Module):
    def __init__(self, lang_model):
        super().__init__()
        self.encoder = lang_model.encoder
        self.pooler = lang_model.pooler
        self.dtype = lang_model.dtype

    def forward(
        self,
        embedding_output: Optional[torch.Tensor] = None, 
        extended_attention_mask: Optional[torch.Tensor] = None, 
        head_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None, 
        encoder_extended_attention_mask: Optional[torch.Tensor] = None, 
        past_key_values: Optional[torch.Tensor] = None, 
        use_cache: Optional[torch.Tensor] = None,
        output_attentions: Optional[torch.Tensor] = None, 
        output_hidden_states: Optional[torch.Tensor] = None, 
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
        
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )

class TextwithPromptBackbone(nn.Module):
    def __init__(self, language_backbone, prompt_mode, category_names, num_classes, num_prompt, eval_only):
        super().__init__()
        self.pooling_mode = language_backbone.pooling_mode
        self.lang_Embedding = BertlangEmbedding(language_backbone, prompt_mode, category_names, num_classes, num_prompt, eval_only)
        self.lang_Encoder = BertlangEncoder(language_backbone.lang_model)
        self.num_prompt = num_prompt

    def forward(self, x, indices, training = False):
        (
            embedding_output, extended_attention_mask, head_mask, encoder_hidden_states, 
            encoder_extended_attention_mask, past_key_values, use_cache, 
            output_attentions, output_hidden_states, return_dict, input_ids,
            loss_orth
        ) = self.lang_Embedding(
            input=x, 
            output_hidden_states=True,
            indices = indices,
            training = training,
        )
        # import pdb; pdb.set_trace()
        output = self.lang_Encoder(
                embedding_output, extended_attention_mask, 
                head_mask, encoder_hidden_states, 
                encoder_extended_attention_mask, 
                past_key_values, use_cache, output_attentions, output_hidden_states, return_dict
            )["last_hidden_state"]

        # import pdb; pdb.set_trace()
        # output = torch.stack(output, dim=0) 

        # import pdb; pdb.set_trace()
             
        if self.pooling_mode == "mean":
            # import pdb; pdb.set_trace()
            output = torch.mean(output, dim=1, keepdim=False)

        elif self.pooling_mode == "max":
            # take features from the eos_token embedding
            eot_indices = torch.nonzero(
                torch.eq(input_ids, self.tokenizer.eos_token_id)
            )
            output = output[torch.arange(output.shape[0]), eot_indices[:, 1]]
        else:
            raise NotImplementedError("Only support pooling_mode: [max, mean].")

        # return output
        return output, loss_orth


class OVDINOspbank2(nn.Module):
    """Implement OV-DINO in `OV-DINO: Unified Open-Vocabulary Detection with Language-Aware Selective Fusion`.

    Args:
        backbone (nn.Module): backbone module
        language_backbone (nn.Module): language backbone module
        position_embedding (nn.Module): position embedding module
        neck (nn.Module): neck module to handle the intermediate outputs features
        transformer (nn.Module): transformer module
        embed_dim (int): dimension of embedding
        num_classes (int): Number of total categories.
        num_queries (int): Number of proposal dynamic anchor boxes in Transformer
        criterion (nn.Module): Criterion for calculating the total losses.
        pixel_mean (List[float]): Pixel mean value for image normalization.
            Default: [123.675, 116.280, 103.530].
        pixel_std (List[float]): Pixel std value for image normalization.
            Default: [58.395, 57.120, 57.375].
        aux_loss (bool): Whether to calculate auxiliary loss in criterion. Default: True.
        select_box_nums_for_evaluation (int): the number of topk candidates
            slected at postprocess for evaluation. Default: 300.
        device (str): Training device. Default: "cuda".
    """

    def __init__(
        self,
        backbone: nn.Module,
        language_backbone: nn.Module,
        position_embedding: nn.Module,
        neck: nn.Module,
        transformer: nn.Module,
        embed_dim: int,
        num_classes: int,
        test_num_classes: int,
        num_queries: int,
        criterion: nn.Module,
        pixel_mean: List[float] = [123.675, 116.280, 103.530],
        pixel_std: List[float] = [58.395, 57.120, 57.375],
        aux_loss: bool = True,
        select_box_nums_for_evaluation: int = 300,
        device="cuda",
        dn_number: int = 100,
        label_noise_ratio: float = 0.2,
        box_noise_scale: float = 1.0,
        input_format: Optional[str] = "RGB",
        vis_period: int = 0,
        text_embed_dim: int = 768,
        inference_template: str = "identity",
        app_mode: bool = False,
        num_prompt: int = 10, 
        text_num_prompt: int = 3,
        prompt_mode = "spbank", 
        category_names = [],
        eval_only = False,
        router_mode = "ddas",
        router_name = "artaxor",
    ):
        super().__init__()
        # define backbone and position embedding module
        self.backbone = backbone
        self.language_backbone = language_backbone
        self.mylanguage_backbone = TextwithPromptBackbone(self.language_backbone, prompt_mode, category_names, num_classes, num_prompt, eval_only)
        self.position_embedding = position_embedding

        # define neck module
        self.neck = neck

        # number of dynamic anchor boxes and embedding dimension
        self.num_queries = num_queries
        self.embed_dim = embed_dim

        # define transformer module
        self.transformer = transformer

        # define classification head and box head
        self.class_embed = ClassEmbed(text_embed_dim, embed_dim)
        self.bbox_embed = MLP(embed_dim, embed_dim, 4, 3)
        self.num_classes = num_classes
        self.test_num_classes = test_num_classes

        # where to calculate auxiliary loss in criterion
        self.aux_loss = aux_loss
        self.criterion = criterion

        # denoising
        self.label_enc = nn.Embedding(num_classes, embed_dim)
        self.dn_number = dn_number
        self.label_noise_ratio = label_noise_ratio
        self.box_noise_scale = box_noise_scale

        # normalizer for input raw images
        self.device = device
        pixel_mean = torch.Tensor(pixel_mean).to(self.device).view(3, 1, 1)
        pixel_std = torch.Tensor(pixel_std).to(self.device).view(3, 1, 1)
        self.normalizer = lambda x: (x - pixel_mean) / pixel_std

        # initialize weights
        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)

        # for _, neck_layer in self.neck.named_modules():
        #     if isinstance(neck_layer, nn.Conv2d):
        #         nn.init.xavier_uniform_(neck_layer.weight, gain=1)
        #         nn.init.constant_(neck_layer.bias, 0)

        # if two-stage, the last class_embed and bbox_embed is for region proposal generation
        num_pred = (
            transformer.decoder.num_layers + 1
            if not self.criterion.two_stage_binary_cls
            else transformer.decoder.num_layers + 2
        )
        self.class_embed = nn.ModuleList(
            [copy.deepcopy(self.class_embed) for i in range(num_pred)]
        )

        self.bbox_embed = nn.ModuleList(
            [copy.deepcopy(self.bbox_embed) for i in range(num_pred)]
        )
        nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)

        # two-stage
        self.transformer.decoder.class_embed = self.class_embed
        self.transformer.decoder.bbox_embed = self.bbox_embed

        # hack implementation for two-stage
        for bbox_embed_layer in self.bbox_embed:
            nn.init.constant_(bbox_embed_layer.layers[-1].bias.data[2:], 0.0)

        # set topk boxes selected for inference
        self.select_box_nums_for_evaluation = select_box_nums_for_evaluation

        # the period for visualizing training samples
        self.input_format = input_format
        self.vis_period = vis_period
        if vis_period > 0:
            assert (
                input_format is not None
            ), "input_format is required for visualization!"

        #  store the same text_embed during evaluation to speed up.
        self.text_embed_dict = defaultdict()
        # store the laste_state of the model during training.
        self.last_state = None
        # the template for inference.
        self.inference_template = inference_template
        self.app_mode = app_mode

        # 初始化 text_key
        self.prompt_mode = prompt_mode
        self.num_prompt = num_prompt

        img_dim = 256
        text_key = torch.empty(num_prompt, img_dim, dtype=torch.float).cuda()
        nn.init.normal_(text_key, std=0.02)
        self.text_key = nn.Parameter(text_key)

        self.text_num_prompt = text_num_prompt 

    def reset_text_embed_dict(self):
        self.text_embed_dict = defaultdict()

    def forward(self, batched_inputs, ae_test = False):
        """Forward function of `DINO` which excepts a list of dict as inputs.

        Args:
            batched_inputs (List[dict]): A list of instance dict, and each instance dict must consists of:
                - dict["image"] (torch.Tensor): The unnormalized image tensor.
                - dict["height"] (int): The original image height.
                - dict["width"] (int): The original image width.
                - dict["instance"] (detectron2.structures.Instances):
                    Image meta informations and ground truth boxes and labels during training.
                    Please refer to
                    https://detectron2.readthedocs.io/en/latest/modules/structures.html#detectron2.structures.Instances
                    for the basic usage of Instances.

        Returns:
            dict: Returns a dict with the following elements:
                - dict["pred_logits"]: the classification logits for all queries (anchor boxes in DAB-DETR).
                            with shape ``[batch_size, num_queries, num_classes]``
                - dict["pred_boxes"]: The normalized boxes coordinates for all queries in format
                    ``(x, y, w, h)``. These values are normalized in [0, 1] relative to the size of
                    each individual image (disregarding possible padding). See PostProcess for information
                    on how to retrieve the unnormalized bounding box.
                - dict["aux_outputs"]: Optional, only returned when auxilary losses are activated. It is a list of
                            dictionnaries containing the two above keys for each decoder layer.
        """
        # store the previous state of the model.
        last_state = self.last_state
        
        use_lora = not ae_test
        images, names = self.preprocess_image_and_name(batched_inputs, use_lora) # ,320

        if self.training:
            batch_size, _, H, W = images.tensor.shape
            img_masks = images.tensor.new_ones(batch_size, H, W)
            for img_id in range(batch_size):
                img_h, img_w = batched_inputs[img_id]["instances"].image_size
                img_masks[img_id, :img_h, :img_w] = 0
        else:
            batch_size, _, H, W = images.tensor.shape
            img_masks = images.tensor.new_zeros(batch_size, H, W)
            assert (
                batch_size == 1
            ), f"batch_size={batch_size}, only support single image inference now..."

            category_start_index = batched_inputs[0].get("category_start_index", 0)
            chunk_index = batched_inputs[0].get("chunk_index", 0)
            self.test_num_classes = len(batched_inputs[0]["category_names"])

        # original features
        features = self.backbone(images.tensor, use_lora)  # output feature dict
        
        # project backbone features to the reuired dimension of transformer
        # we use multi-scale features in DINO
        multi_level_feats = self.neck(features)

        # multi_level_feats = list(multi_level_feats)

        '''
        # baozha
        # for i, feat in enumerate(multi_level_feats):
        #     print(f"Shape of feature {i}:", feat.shape)
        # Shape of feature 0: torch.Size([1, 256, 100, 163])
        # Shape of feature 1: torch.Size([1, 256, 50, 82])
        # Shape of feature 2: torch.Size([1, 256, 25, 41])
        # Shape of feature 3: torch.Size([1, 256, 13, 21])
        # for i, feat in enumerate(multi_level_feats):
        #     B, C, H, W = feat.shape
        #     feat_reshaped = feat.view(B, C, -1) # [B, C, H*W]
        #     feat_reshaped = feat_reshaped.view(C, -1) # [C, B*H*W]
        #     feat_reshaped = feat_reshaped.permute(1, 0) # [B*H*W, C]
        #     feat_processed = self.visual_linear_layers[i](feat_reshaped)
        #     feat_processed = feat_processed.view(B, C, H, W)
        #     multi_level_feats[i] = feat_processed
        '''

        # attribank多尺度，每个尺度选top1
        if self.prompt_mode == "spbank" and ae_test == False:
            # with torch.no_grad():
            query_key = []
            
            for i, feat in enumerate(multi_level_feats):
                pooled_feature = feat.mean(dim=[2, 3])  # (batch_size, dim) 
                query_key.append(pooled_feature)
                if i == 2:
                    break
            
            query_key = torch.stack(query_key, dim=1)  # (batch_size, num_scales, dim)
            # with torch.no_grad():
            image_features = query_key / query_key.norm(dim=-1, keepdim=True) 
            image_features = image_features[:, :3, :]

            probability = query_key @ self.text_key.t()  # [16,3,12]

            # import pdb;pdb.set_trace()

            _, indices = probability.topk(k=1, dim=2, largest=True)
            indices = indices[:, :3, :]  # (batch_size, 3) 
            indices = indices.squeeze(dim=2)
            # import pdb; pdb.set_trace()
            key_choose = self.text_key[indices]
            self.loss_k = cosine_loss(image_features, key_choose)

            # text_embed = self.mylanguage_backbone(names, indices, self.training)
            text_embed, loss_orth = self.mylanguage_backbone(names, indices, self.training)
            
        else:
            if self.text_embed_dict.get(chunk_index, None) is None:
                if isinstance(names[0], list):
                    text_embed = torch.stack(
                        [self.language_backbone(name) for name in names], dim=0
                    )  # [bs*num_classes, num_templates, embed_dim]
                    text_embed = text_embed.mean(1)
                else:
                    text_embed = self.language_backbone(names)
            
                self.text_embed_dict[chunk_index] = text_embed
            else:
                text_embed = self.text_embed_dict[chunk_index]
            # text_embed = self.language_backbone(names)

        num_classes = self.num_classes if self.training else self.test_num_classes
        # text_embed = text_embed.view(batch_size, num_classes, -1)
        # print("after_text_embed: ", text_embed.shape) # torch.Size([4, 80, 768]) batch,class,embedding  
        text_embed = text_embed.view(batch_size, num_classes, -1)

        # import pdb; pdb.set_trace()

        multi_level_masks = []
        multi_level_position_embeddings = []
        for feat in multi_level_feats:
            multi_level_masks.append(
                F.interpolate(img_masks[None], size=feat.shape[-2:])
                .to(torch.bool)
                .squeeze(0)
            )
            multi_level_position_embeddings.append(
                self.position_embedding(multi_level_masks[-1])
            )
        '''
        print("multi_level_masks:", multi_level_masks[-1].shape)
        print("multi_level_position_embeddings:", multi_level_position_embeddings[-1].shape)
        multi_level_masks: torch.Size([1, 100, 134])                                                                      
        multi_level_position_embeddings: torch.Size([1, 256, 100, 134])                                                   
        multi_level_masks: torch.Size([1, 50, 67])                                                                        
        multi_level_position_embeddings: torch.Size([1, 256, 50, 67])                                                     
        multi_level_masks: torch.Size([1, 25, 34])                                                                        
        multi_level_position_embeddings: torch.Size([1, 256, 25, 34])                                                     
        multi_level_masks: torch.Size([1, 13, 17])                                                                        
        multi_level_position_embeddings: torch.Size([1, 256, 13, 17]) 
        '''
        
        # denoising preprocessing
        # prepare label query embedding
        if self.training:
            gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
            targets = self.prepare_targets(gt_instances)
            # target = {"labels": gt_classes, "boxes": gt_boxes}
            input_query_label, input_query_bbox, attn_mask, dn_meta = (
                self.prepare_for_cdn(
                    targets,
                    dn_number=self.dn_number,
                    label_noise_ratio=self.label_noise_ratio,
                    box_noise_scale=self.box_noise_scale,
                    num_queries=self.num_queries,
                    num_classes=self.num_classes,
                    hidden_dim=self.embed_dim,
                    label_enc=self.label_enc,
                )
            )
        else:
            input_query_label, input_query_bbox, attn_mask, dn_meta = (
                None,
                None,
                None,
                None,
            )
        query_embeds = (input_query_label, input_query_bbox)

        # feed into transformer
        (
            inter_states,
            init_reference,
            inter_references,
            enc_state,
            enc_reference,  
        ) = self.transformer(
            multi_level_feats,
            multi_level_masks,
            multi_level_position_embeddings,
            query_embeds,
            attn_masks=[attn_mask, None, None],
            class_embed=text_embed,
        )

        inter_states[0] += self.label_enc.weight[0, 0] * 0.0

        # Calculate output coordinates and classes.
        outputs_classes = []
        outputs_coords = []

        for lvl in range(inter_states.shape[0]):
            if lvl == 0:
                reference = init_reference
            else:
                reference = inter_references[lvl - 1]
            reference = inverse_sigmoid(reference)
            outputs_class = self.class_embed[lvl](inter_states[lvl], text_embed)
            tmp = self.bbox_embed[lvl](inter_states[lvl])
            if reference.shape[-1] == 4:
                tmp += reference
            else:
                assert reference.shape[-1] == 2
                tmp[..., :2] += reference
            outputs_coord = tmp.sigmoid()
            outputs_classes.append(outputs_class)
            outputs_coords.append(outputs_coord)

        outputs_class = torch.stack(outputs_classes)
        # tensor shape: [num_decoder_layers, bs, num_query, num_classes]
        outputs_coord = torch.stack(outputs_coords)
        # tensor shape: [num_decoder_layers, bs, num_query, 4]

        # denoising postprocessing
        if dn_meta is not None:
            outputs_class, outputs_coord = self.dn_post_process(
                outputs_class, outputs_coord, dn_meta
            )

        # prepare for loss computation
        output = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord[-1]}
        if self.aux_loss:
            output["aux_outputs"] = self._set_aux_loss(outputs_class, outputs_coord)

        # prepare two stage output
        interm_class = self.transformer.decoder.class_embed[
            self.transformer.decoder.num_layers
        ](enc_state, text_embed)
        interm_coord = enc_reference
        output["enc_outputs"] = {
            "pred_logits": interm_class,
            "pred_boxes": interm_coord,
        }

        if self.training:
            # visualize training samples
            if self.vis_period > 0:
                storage = get_event_storage()
                if storage.iter % self.vis_period == 0:
                    box_cls = output["pred_logits"]
                    box_pred = output["pred_boxes"]
                    results = self.inference(box_cls, box_pred, images.image_sizes)
                    self.visualize_training(batched_inputs, results)

            # compute loss
            loss_dict = self.criterion(output, targets, dn_meta)


            loss_dict["loss_p"] = 0.3 * loss_orth
            loss_dict["loss_m"] = 0.7 * self.loss_k
            

            weight_dict = self.criterion.weight_dict
            for k in loss_dict.keys():
                if k in weight_dict:
                    loss_dict[k] *= weight_dict[k]
            return loss_dict
        else:
            box_cls = output["pred_logits"]
            box_pred = output["pred_boxes"]
            results = self.inference(
                box_cls, box_pred, images.image_sizes, category_start_index
            )
            processed_results = []
            for results_per_image, input_per_image, image_size in zip(
                results, batched_inputs, images.image_sizes
            ):
                height = input_per_image.get("height", image_size[0])
                width = input_per_image.get("width", image_size[1])
                r = detector_postprocess(results_per_image, height, width)
                processed_results.append({"instances": r})
            return processed_results

    def visualize_training(self, batched_inputs, results):
        from detectron2.utils.visualizer import Visualizer

        storage = get_event_storage()
        max_vis_box = 20

        for input, results_per_image in zip(batched_inputs, results):
            img = input["image"]
            img = convert_image_to_rgb(img.permute(1, 2, 0), self.input_format)
            v_gt = Visualizer(img, None)
            v_gt = v_gt.overlay_instances(boxes=input["instances"].gt_boxes)
            anno_img = v_gt.get_image()
            v_pred = Visualizer(img, None)
            v_pred = v_pred.overlay_instances(
                boxes=results_per_image.pred_boxes[:max_vis_box]
                .tensor.detach()
                .cpu()
                .numpy()
            )
            pred_img = v_pred.get_image()
            vis_img = np.concatenate((anno_img, pred_img), axis=1)
            vis_img = vis_img.transpose(2, 0, 1)
            vis_name = "Left: GT bounding boxes;  Right: Predicted boxes"
            storage.put_image(vis_name, vis_img)
            break  # only visualize one image in a batch

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # 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}
            for a, b in zip(outputs_class[:-1], outputs_coord[:-1])
        ]

    def prepare_for_cdn(
        self,
        targets,
        dn_number,
        label_noise_ratio,
        box_noise_scale,
        num_queries,
        num_classes,
        hidden_dim,
        label_enc,
    ):
        """
        A major difference of DINO from DN-DETR is that the author process pattern embedding pattern embedding
            in its detector
        forward function and use learnable tgt embedding, so we change this function a little bit.
        :param dn_args: targets, dn_number, label_noise_ratio, box_noise_scale
        :param training: if it is training or inference
        :param num_queries: number of queires
        :param num_classes: number of classes
        :param hidden_dim: transformer hidden dim
        :param label_enc: encode labels in dn
        :return:
        """
        # print("targets", targets)
        if dn_number <= 0:
            return None, None, None, None
            # positive and negative dn queries
        # 初始化正负样本数量
        dn_number = dn_number * 2
        known = [(torch.ones_like(t["labels"])).cuda() for t in targets] # 每个目标的正样本掩码
        # know: [tensor([1], device='cuda:1'), tensor([1, 1, 1], device='cuda:1'), tensor([1, 1, 1, 1, 1, 1, 1, 1], device='
        # cuda:1'), tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], device='cuda:1'), tensor([1, 1, 1, 1, 1], device='cuda:1'), tenso
        # r([1], device='cuda:1'), tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], device='cuda:1'), tensor([1, 1], device='cuda:1
        # ')] 
        batch_size = len(known) # 有多少样本
        known_num = [sum(k) for k in known] # 每个样本包含的目标标注数量
        # print("known_num:", known_num)
        if int(max(known_num)) == 0:
            return None, None, None, None

        # 动态噪声样本数量调整
        dn_number = dn_number // (int(max(known_num) * 2))
        #（确保dn_number与known_num相匹配)
        if dn_number == 0:
            dn_number = 1
        
        # 生成正样本和负样本的索引和标签
        unmask_bbox = unmask_label = torch.cat(known)
        labels = torch.cat([t["labels"] for t in targets])
        boxes = torch.cat([t["boxes"] for t in targets])
        batch_idx = torch.cat(
            [torch.full_like(t["labels"].long(), i) for i, t in enumerate(targets)]
        ) 


        known_indice = torch.nonzero(unmask_label + unmask_bbox) 
        known_indice = known_indice.view(-1) #  [81]

        known_indice = known_indice.repeat(2 * dn_number, 1).view(-1) # [486]
        known_labels = labels.repeat(2 * dn_number, 1).view(-1)
        known_bid = batch_idx.repeat(2 * dn_number, 1).view(-1)
        known_bboxs = boxes.repeat(2 * dn_number, 1)
        known_labels_expaned = known_labels.clone() 
        known_bbox_expand = known_bboxs.clone()


        if label_noise_ratio > 0:
            p = torch.rand_like(known_labels_expaned.float()) 
            chosen_indice = torch.nonzero(p < (label_noise_ratio * 0.5)).view(
                -1
            )  # half of bbox prob
            new_label = torch.randint_like(
                chosen_indice, 0, num_classes
            )  # randomly put a new one here
            known_labels_expaned.scatter_(0, chosen_indice, new_label) 
        
        single_padding = int(max(known_num)) 
        pad_size = int(single_padding * 2 * dn_number) 
        
        positive_idx = (
            torch.tensor(range(len(boxes)))
            .long()
            .cuda()
            .unsqueeze(0)
            .repeat(dn_number, 1)
        )
        positive_idx += (
            (torch.tensor(range(dn_number)) * len(boxes) * 2).long().cuda().unsqueeze(1)
        )
        positive_idx = positive_idx.flatten()
        negative_idx = positive_idx + len(boxes)


        if box_noise_scale > 0:
            known_bbox_ = torch.zeros_like(known_bboxs)
            known_bbox_[:, :2] = known_bboxs[:, :2] - known_bboxs[:, 2:] / 2
            known_bbox_[:, 2:] = known_bboxs[:, :2] + known_bboxs[:, 2:] / 2

            diff = torch.zeros_like(known_bboxs)
            diff[:, :2] = known_bboxs[:, 2:] / 2
            diff[:, 2:] = known_bboxs[:, 2:] / 2

            rand_sign = (
                torch.randint_like(known_bboxs, low=0, high=2, dtype=torch.float32)
                * 2.0
                - 1.0
            )
            rand_part = torch.rand_like(known_bboxs)
            rand_part[negative_idx] += 1.0
            rand_part *= rand_sign
            known_bbox_ = (
                known_bbox_ + torch.mul(rand_part, diff).cuda() * box_noise_scale
            )
            known_bbox_ = known_bbox_.clamp(min=0.0, max=1.0)
            known_bbox_expand[:, :2] = (known_bbox_[:, :2] + known_bbox_[:, 2:]) / 2
            known_bbox_expand[:, 2:] = known_bbox_[:, 2:] - known_bbox_[:, :2]

        m = known_labels_expaned.long().to("cuda") # [1608] （0-79）
        # print(m)
        input_label_embed = label_enc(m) 
        input_bbox_embed = inverse_sigmoid(known_bbox_expand) 

        padding_label = torch.zeros(pad_size, hidden_dim).cuda() 
        padding_bbox = torch.zeros(pad_size, 4).cuda()
        input_query_label = padding_label.repeat(batch_size, 1, 1) 
        input_query_bbox = padding_bbox.repeat(batch_size, 1, 1)


        map_known_indice = torch.tensor([]).to("cuda")
        if len(known_num):
            map_known_indice = torch.cat(
                [torch.tensor(range(num)) for num in known_num]
            )  # [1,2, 1,2,3]
            map_known_indice = torch.cat(
                [map_known_indice + single_padding * i for i in range(2 * dn_number)]
            ).long()

        if len(known_bid):
            input_query_label[(known_bid.long(), map_known_indice)] = input_label_embed
            input_query_bbox[(known_bid.long(), map_known_indice)] = input_bbox_embed

        tgt_size = pad_size + num_queries
        attn_mask = torch.ones(tgt_size, tgt_size).to("cuda") < 0
        # match query cannot see the reconstruct
        attn_mask[pad_size:, :pad_size] = True
        # reconstruct cannot see each other
        for i in range(dn_number):
            if i == 0:
                attn_mask[
                    single_padding * 2 * i : single_padding * 2 * (i + 1),
                    single_padding * 2 * (i + 1) : pad_size,
                ] = True
            if i == dn_number - 1:
                attn_mask[
                    single_padding * 2 * i : single_padding * 2 * (i + 1),
                    : single_padding * i * 2,
                ] = True
            else:
                attn_mask[
                    single_padding * 2 * i : single_padding * 2 * (i + 1),
                    single_padding * 2 * (i + 1) : pad_size,
                ] = True
                attn_mask[
                    single_padding * 2 * i : single_padding * 2 * (i + 1),
                    : single_padding * 2 * i,
                ] = True

        dn_meta = {
            "single_padding": single_padding * 2,
            "dn_num": dn_number,
        }
        # input_query_label torch.Size([8, 190, 256])                                                                       
        # input_query_bbox torch.Size([8, 190, 4])  
        return input_query_label, input_query_bbox, attn_mask, dn_meta

    def dn_post_process(self, outputs_class, outputs_coord, dn_metas):
        if dn_metas and dn_metas["single_padding"] > 0:
            padding_size = dn_metas["single_padding"] * dn_metas["dn_num"]
            output_known_class = outputs_class[:, :, :padding_size, :]
            output_known_coord = outputs_coord[:, :, :padding_size, :]
            outputs_class = outputs_class[:, :, padding_size:, :]
            outputs_coord = outputs_coord[:, :, padding_size:, :]

            out = {
                "pred_logits": output_known_class[-1],
                "pred_boxes": output_known_coord[-1],
            }
            if self.aux_loss:
                out["aux_outputs"] = self._set_aux_loss(
                    output_known_class, output_known_coord
                )
            dn_metas["output_known_lbs_bboxes"] = out
        return outputs_class, outputs_coord

    def preprocess_image_and_name(self, batched_inputs, use_lora):
        images = [self.normalizer(x["image"].to(self.device)) for x in batched_inputs]
        images = ImageList.from_tensors(images)

        # process category_names
        category_names = [
            x["category_names"] for x in batched_inputs
        ]  # list[list[str]]
        if not self.training and self.inference_template != "identity":
            category_names = [
                [
                    [
                        template.format(clean_words_or_phrase(cat_name))
                        for template in template_meta[self.inference_template]
                    ]
                    for cat_name in batch_cat_names
                ]
                for batch_cat_names in category_names
            ]

        if not self.training and use_lora == False:
            category_names = [
                [
                    [
                        template.format(clean_words_or_phrase(cat_name))
                        for template in template_meta["full"]
                    ]
                    for cat_name in batch_cat_names
                ]
                for batch_cat_names in category_names
            ]

        category_names = list(itertools.chain(*category_names))

        return images, category_names

    def inference(self, box_cls, box_pred, image_sizes, category_start_index=0):
        """
        Arguments:
            box_cls (Tensor): tensor of shape (batch_size, num_queries, K).
                The tensor predicts the classification probability for each query.
            box_pred (Tensor): tensors of shape (batch_size, num_queries, 4).
                The tensor predicts 4-vector (x,y,w,h) box
                regression values for every queryx
            image_sizes (List[torch.Size]): the input image sizes.
            category_start_index (Int): the shift index of category.

        Returns:
            results (List[Instances]): a list of #images elements.
        """
        assert len(box_cls) == len(image_sizes)
        results = []

        # box_cls.shape: 1, 300, 80
        # box_pred.shape: 1, 300, 4
        prob = box_cls.sigmoid()
        # follow glip(ref: https://github.com/microsoft/GLIP/blob/a5f302bfd4c5c67010e29f779e3b0bde94e89985/maskrcnn_benchmark/modeling/rpn/inference.py#L510).
        # We perform sqrt of cls_prob to obtain more larger probility for visulization.
        prob = torch.sqrt(prob)
        topk_values, topk_indexes = torch.topk(
            prob.view(box_cls.shape[0], -1), self.select_box_nums_for_evaluation, dim=1
        )
        scores = topk_values
        topk_boxes = torch.div(topk_indexes, box_cls.shape[2], rounding_mode="floor")
        labels = topk_indexes % box_cls.shape[2] + category_start_index

        boxes = torch.gather(box_pred, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))

        # For each box we assign the best class or the second best if the best on is `no_object`.
        # scores, labels = F.softmax(box_cls, dim=-1)[:, :, :-1].max(-1)

        for i, (
            scores_per_image,
            labels_per_image,
            box_pred_per_image,
            image_size,
        ) in enumerate(zip(scores, labels, boxes, image_sizes)):
            result = Instances(image_size)
            result.pred_boxes = Boxes(box_cxcywh_to_xyxy(box_pred_per_image))

            result.pred_boxes.scale(scale_x=image_size[1], scale_y=image_size[0])
            result.scores = scores_per_image
            result.pred_classes = labels_per_image
            results.append(result)
        return results

    def prepare_targets(self, targets):
        new_targets = []
        for targets_per_image in targets:
            h, w = targets_per_image.image_size
            image_size_xyxy = torch.as_tensor(
                [w, h, w, h], dtype=torch.float, device=self.device
            )
            gt_classes = targets_per_image.gt_classes
            gt_boxes = targets_per_image.gt_boxes.tensor / image_size_xyxy
            gt_boxes = box_xyxy_to_cxcywh(gt_boxes)
            new_targets.append({"labels": gt_classes, "boxes": gt_boxes})
        return new_targets


# we reimplement the detector_postprocess here to meet our demands:
# 1. filter large boxes are almost the image_size.
# TODO: filter large boxes are almost the image_size.
def detector_postprocess(
    results: Instances,
    output_height: int,
    output_width: int,
    iou_threshold: float = 0.95,
):
    """
    Resize the output instances.
    The input images are often resized when entering an object detector.
    As a result, we often need the outputs of the detector in a different
    resolution from its inputs.

    This function will resize the raw outputs of an R-CNN detector
    to produce outputs according to the desired output resolution.

    Args:
        results (Instances): the raw outputs from the detector.
            `results.image_size` contains the input image resolution the detector sees.
            This object might be modified in-place.
        output_height, output_width: the desired output resolution.
    Returns:
        Instances: the resized output from the model, based on the output resolution
    """
    if isinstance(output_width, torch.Tensor):

        output_width_tmp = output_width.float()
        output_height_tmp = output_height.float()
        new_size = torch.stack([output_height, output_width])  
    else:

        new_size = (output_height, output_width)
        output_width_tmp = output_width
        output_height_tmp = output_height


    scale_x, scale_y = (
        output_width_tmp / results.image_size[1], 
        output_height_tmp / results.image_size[0], 
    )


    results = Instances(new_size, **results.get_fields())


    if results.has("pred_boxes"):
        output_boxes = results.pred_boxes  
    elif results.has("proposal_boxes"):
        output_boxes = results.proposal_boxes  
    else:
        output_boxes = None
    assert output_boxes is not None, "Predictions must contain boxes!" 


    output_boxes.scale(scale_x, scale_y)

    output_boxes.clip(results.image_size)


    results = results[output_boxes.nonempty()]


    image_width, image_height = results.image_size
    image_boxes = Boxes(
        torch.tensor([[0, 0, image_width, image_height]], device=output_boxes.device)
    )  
    ious = pairwise_iou(output_boxes, image_boxes)[:, 0]  
    results = results[ious < iou_threshold]  

    return results
