

'''simple docstring'''






import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__      		=			logging.get_logger(__name__)


def 	_A  (					A__   ):






 """simple docstring"""



 __lowercase													=				SwinConfig(
     embed_dim=192  ,       depths=(2, 2, 18, 2)  ,       num_heads=(6, 12, 24, 48)  ,       window_size=12  ,       out_features=['''stage2''', '''stage3''', '''stage4''']  ,       )

 __lowercase													=				DetaConfig(
     backbone_config=A__  ,       num_queries=900  ,       encoder_ffn_dim=2048  ,       decoder_ffn_dim=2048  ,       num_feature_levels=5  ,       assign_first_stage=A__  ,       with_box_refine=A__  ,       two_stage=A__  ,       )

 # set labels
 __lowercase													=				'''huggingface/label-files'''
 if "o365" in model_name:
  __lowercase													=				366
  __lowercase													=				'''object365-id2label.json'''
 else:
  __lowercase													=				91
  __lowercase													=				'''coco-detection-id2label.json'''

 __lowercase													=				num_labels
 __lowercase													=				json.load(open(cached_download(hf_hub_url(A__  ,       A__  ,       repo_type='''dataset'''   )   )  ,       '''r'''   )   )
 __lowercase													=				{int(A__   ): v for k, v in idalabel.items()}
 __lowercase													=				idalabel
 __lowercase													=				{v: k for k, v in idalabel.items()}

 return config


def 	_A  (					A__   ):






 """simple docstring"""



 __lowercase													=				[]

 # stem
 # fmt: off
 rename_keys.append(('''backbone.0.body.patch_embed.proj.weight''', '''model.backbone.model.embeddings.patch_embeddings.projection.weight''')   )
 rename_keys.append(('''backbone.0.body.patch_embed.proj.bias''', '''model.backbone.model.embeddings.patch_embeddings.projection.bias''')   )
 rename_keys.append(('''backbone.0.body.patch_embed.norm.weight''', '''model.backbone.model.embeddings.norm.weight''')   )
 rename_keys.append(('''backbone.0.body.patch_embed.norm.bias''', '''model.backbone.model.embeddings.norm.bias''')   )
 # stages
 for i in range(len(config.backbone_config.depths   )   ):
  for j in range(config.backbone_config.depths[i]   ):
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm1.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm1.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm2.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm2.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias")   )

  if i < 3:
   rename_keys.append((F"backbone.0.body.layers.{i}.downsample.reduction.weight", F"model.backbone.model.encoder.layers.{i}.downsample.reduction.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.downsample.norm.weight", F"model.backbone.model.encoder.layers.{i}.downsample.norm.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.downsample.norm.bias", F"model.backbone.model.encoder.layers.{i}.downsample.norm.bias")   )

 rename_keys.append(('''backbone.0.body.norm1.weight''', '''model.backbone.model.hidden_states_norms.stage2.weight''')   )
 rename_keys.append(('''backbone.0.body.norm1.bias''', '''model.backbone.model.hidden_states_norms.stage2.bias''')   )
 rename_keys.append(('''backbone.0.body.norm2.weight''', '''model.backbone.model.hidden_states_norms.stage3.weight''')   )
 rename_keys.append(('''backbone.0.body.norm2.bias''', '''model.backbone.model.hidden_states_norms.stage3.bias''')   )
 rename_keys.append(('''backbone.0.body.norm3.weight''', '''model.backbone.model.hidden_states_norms.stage4.weight''')   )
 rename_keys.append(('''backbone.0.body.norm3.bias''', '''model.backbone.model.hidden_states_norms.stage4.bias''')   )

 # transformer encoder
 for i in range(config.encoder_layers   ):
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight", F"model.encoder.layers.{i}.self_attn.sampling_offsets.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias", F"model.encoder.layers.{i}.self_attn.sampling_offsets.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.attention_weights.weight", F"model.encoder.layers.{i}.self_attn.attention_weights.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.attention_weights.bias", F"model.encoder.layers.{i}.self_attn.attention_weights.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.value_proj.weight", F"model.encoder.layers.{i}.self_attn.value_proj.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.value_proj.bias", F"model.encoder.layers.{i}.self_attn.value_proj.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.output_proj.weight", F"model.encoder.layers.{i}.self_attn.output_proj.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.output_proj.bias", F"model.encoder.layers.{i}.self_attn.output_proj.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm1.weight", F"model.encoder.layers.{i}.self_attn_layer_norm.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm1.bias", F"model.encoder.layers.{i}.self_attn_layer_norm.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear1.weight", F"model.encoder.layers.{i}.fc1.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear1.bias", F"model.encoder.layers.{i}.fc1.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear2.weight", F"model.encoder.layers.{i}.fc2.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear2.bias", F"model.encoder.layers.{i}.fc2.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm2.weight", F"model.encoder.layers.{i}.final_layer_norm.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm2.bias", F"model.encoder.layers.{i}.final_layer_norm.bias")   )

 # transformer decoder
 for i in range(config.decoder_layers   ):
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight", F"model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias", F"model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.attention_weights.weight", F"model.decoder.layers.{i}.encoder_attn.attention_weights.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.attention_weights.bias", F"model.decoder.layers.{i}.encoder_attn.attention_weights.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.value_proj.weight", F"model.decoder.layers.{i}.encoder_attn.value_proj.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.value_proj.bias", F"model.decoder.layers.{i}.encoder_attn.value_proj.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.output_proj.weight", F"model.decoder.layers.{i}.encoder_attn.output_proj.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.output_proj.bias", F"model.decoder.layers.{i}.encoder_attn.output_proj.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm1.weight", F"model.decoder.layers.{i}.encoder_attn_layer_norm.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm1.bias", F"model.decoder.layers.{i}.encoder_attn_layer_norm.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.self_attn.out_proj.weight", F"model.decoder.layers.{i}.self_attn.out_proj.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.self_attn.out_proj.bias", F"model.decoder.layers.{i}.self_attn.out_proj.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm2.weight", F"model.decoder.layers.{i}.self_attn_layer_norm.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm2.bias", F"model.decoder.layers.{i}.self_attn_layer_norm.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear1.weight", F"model.decoder.layers.{i}.fc1.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear1.bias", F"model.decoder.layers.{i}.fc1.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear2.weight", F"model.decoder.layers.{i}.fc2.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear2.bias", F"model.decoder.layers.{i}.fc2.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm3.weight", F"model.decoder.layers.{i}.final_layer_norm.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm3.bias", F"model.decoder.layers.{i}.final_layer_norm.bias")   )

 # fmt: on

 return rename_keys


def 	_A  (					A__  ,       A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				dct.pop(A__   )
 __lowercase													=				val


def 	_A  (					A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				[int(backbone_config.embed_dim * 2**i   ) for i in range(len(backbone_config.depths   )   )]
 for i in range(len(backbone_config.depths   )   ):
  __lowercase													=				num_features[i]
  for j in range(backbone_config.depths[i]   ):
   # fmt: off
   # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
   __lowercase													=				state_dict.pop(F"backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight"   )
   __lowercase													=				state_dict.pop(F"backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias"   )
   # next, add query, keys and values (in that order) to the state dict
   __lowercase													=				in_proj_weight[:dim, :]
   __lowercase													=				in_proj_bias[: dim]
   __lowercase													=				in_proj_weight[
       dim : dim * 2, :
   ]
   __lowercase													=				in_proj_bias[
       dim : dim * 2
   ]
   __lowercase													=				in_proj_weight[
       -dim :, :
   ]
   __lowercase													=				in_proj_bias[-dim :]
   # fmt: on


def 	_A  (					A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				config.d_model
 for i in range(config.decoder_layers   ):
  # read in weights + bias of input projection layer of self-attention
  __lowercase													=				state_dict.pop(F"transformer.decoder.layers.{i}.self_attn.in_proj_weight"   )
  __lowercase													=				state_dict.pop(F"transformer.decoder.layers.{i}.self_attn.in_proj_bias"   )
  # next, add query, keys and values (in that order) to the state dict
  __lowercase													=				in_proj_weight[:hidden_size, :]
  __lowercase													=				in_proj_bias[:hidden_size]
  __lowercase													=				in_proj_weight[
      hidden_size : hidden_size * 2, :
  ]
  __lowercase													=				in_proj_bias[hidden_size : hidden_size * 2]
  __lowercase													=				in_proj_weight[-hidden_size:, :]
  __lowercase													=				in_proj_bias[-hidden_size:]


def 	_A  (					):






 """simple docstring"""



 __lowercase													=				'''http://images.cocodataset.org/val2017/000000039769.jpg'''
 __lowercase													=				Image.open(requests.get(A__  ,       stream=A__   ).raw   )

 return im


@torch.no_grad()
def 	_A  (					A__  ,       A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				get_deta_config(A__   )

 # load original state dict
 if model_name == "deta-swin-large":
  __lowercase													=				hf_hub_download(repo_id='''nielsr/deta-checkpoints'''  ,       filename='''adet_swin_ft.pth'''   )
 elif model_name == "deta-swin-large-o365":
  __lowercase													=				hf_hub_download(repo_id='''jozhang97/deta-swin-l-o365'''  ,       filename='''deta_swin_pt_o365.pth'''   )
 else:
  raise ValueError(F"Model name {model_name} not supported"   )

 __lowercase													=				torch.load(A__  ,       map_location='''cpu'''   )['''model''']

 # original state dict
 for name, param in state_dict.items():
  print(A__  ,       param.shape   )

 # rename keys
 __lowercase													=				create_rename_keys(A__   )
 for src, dest in rename_keys:
  rename_key(A__  ,       A__  ,       A__   )
 read_in_swin_q_k_v(A__  ,       config.backbone_config   )
 read_in_decoder_q_k_v(A__  ,       A__   )

 # fix some prefixes
 for key in state_dict.copy().keys():
  if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
   __lowercase													=				state_dict.pop(A__   )
   __lowercase													=				val
  if "input_proj" in key:
   __lowercase													=				state_dict.pop(A__   )
   __lowercase													=				val
  if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
   __lowercase													=				state_dict.pop(A__   )
   __lowercase													=				val

    # finally, create HuggingFace model and load state dict
 __lowercase													=				DetaForObjectDetection(A__   )
 model.load_state_dict(A__   )
 model.eval()

 __lowercase													=				'''cuda''' if torch.cuda.is_available() else '''cpu'''
 model.to(A__   )

 # load image processor
 __lowercase													=				DetaImageProcessor(format='''coco_detection'''   )

 # verify our conversion on image
 __lowercase													=				prepare_img()
 __lowercase													=				processor(images=A__  ,       return_tensors='''pt'''   )
 __lowercase													=				encoding['''pixel_values''']
 __lowercase													=				model(pixel_values.to(A__   )   )

 # verify logits
 print('''Logits:'''  ,       outputs.logits[0, :3, :3]   )
 print('''Boxes:'''  ,       outputs.pred_boxes[0, :3, :3]   )
 if model_name == "deta-swin-large":
  __lowercase													=				torch.tensor(
      [[-7.6_3_0_8, -2.8_4_8_5, -5.3_7_3_7], [-7.2_0_3_7, -4.5_5_0_5, -4.8_0_2_7], [-7.2_9_4_3, -4.2_6_1_1, -4.6_6_1_7]]   )
  __lowercase													=				torch.tensor([[0.4_9_8_7, 0.4_9_6_9, 0.9_9_9_9], [0.2_5_4_9, 0.5_4_9_8, 0.4_8_0_5], [0.5_4_9_8, 0.2_7_5_7, 0.0_5_6_9]]   )
 elif model_name == "deta-swin-large-o365":
  __lowercase													=				torch.tensor(
      [[-8.0_1_2_2, -3.5_7_2_0, -4.9_7_1_7], [-8.1_5_4_7, -3.6_8_8_6, -4.6_3_8_9], [-7.6_6_1_0, -3.6_1_9_4, -5.0_1_3_4]]   )
  __lowercase													=				torch.tensor([[0.2_5_2_3, 0.5_5_4_9, 0.4_8_8_1], [0.7_7_1_5, 0.4_1_4_9, 0.4_6_0_1], [0.5_5_0_3, 0.2_7_5_3, 0.0_5_7_5]]   )
 assert torch.allclose(outputs.logits[0, :3, :3]  ,       expected_logits.to(A__   )  ,       atol=1e-4   )
 assert torch.allclose(outputs.pred_boxes[0, :3, :3]  ,       expected_boxes.to(A__   )  ,       atol=1e-4   )
 print('''Everything ok!'''   )

 if pytorch_dump_folder_path:
  # Save model and processor
  logger.info(F"Saving PyTorch model and processor to {pytorch_dump_folder_path}..."   )
  Path(A__   ).mkdir(exist_ok=A__   )
  model.save_pretrained(A__   )
  processor.save_pretrained(A__   )

 # Push to hub
 if push_to_hub:
  print('''Pushing model and processor to hub...'''   )
  model.push_to_hub(F"jozhang97/{model_name}"   )
  processor.push_to_hub(F"jozhang97/{model_name}"   )


if __name__ == "__main__":
 lowerCAmelCase__      		=			argparse.ArgumentParser()

 parser.add_argument(
     '''--model_name''',
     type=str,
     default='''deta-swin-large''',
     choices=['''deta-swin-large''', '''deta-swin-large-o365'''],
     help='''Name of the model you\'d like to convert.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     help='''Path to the folder to output PyTorch model.''',
 )
 parser.add_argument(
     '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
 )
 lowerCAmelCase__      		=			parser.parse_args()
 convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








'''simple docstring'''






from argparse import ArgumentParser

from . import BaseTransformersCLICommand


def 	_A  (					A__   ):






 """simple docstring"""



 return DownloadCommand(args.model  ,       args.cache_dir  ,       args.force  ,       args.trust_remote_code   )






class     lowercase_			(lowerCamelCase__   ):




  """simple docstring"""




  @staticmethod
  def       SCREAMING_SNAKE_CASE   (   lowercase__     :  ArgumentParser    ):
   __lowercase													=				parser.add_parser('''download'''    )
   download_parser.add_argument(
       '''--cache-dir'''      ,type=lowercase__      ,default=lowercase__      ,help='''Path to location to store the models'''    )
   download_parser.add_argument(
       '''--force'''      ,action='''store_true'''      ,help='''Force the model to be download even if already in cache-dir'''    )
   download_parser.add_argument(
       '''--trust-remote-code'''      ,action='''store_true'''      ,help='''Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you\'ve reviewed the code as it will execute on your local machine'''      ,)
   download_parser.add_argument('''model'''      ,type=lowercase__      ,help='''Name of the model to download'''    )
   download_parser.set_defaults(func=lowercase__    )






  def __init__(   self     :  str      ,lowercase__     :  str      ,lowercase__     :  str      ,lowercase__     :  bool      ,lowercase__     :  bool    ):
   __lowercase													=				model
   __lowercase													=				cache
   __lowercase													=				force
   __lowercase													=				trust_remote_code






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   from ..models.auto import AutoModel, AutoTokenizer

   AutoModel.from_pretrained(
       self._model      ,cache_dir=self._cache      ,force_download=self._force      ,trust_remote_code=self._trust_remote_code    )
   AutoTokenizer.from_pretrained(
       self._model      ,cache_dir=self._cache      ,force_download=self._force      ,trust_remote_code=self._trust_remote_code    )

