




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
from transformers.utils import logging


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





def 							lowerCamelCase    (							)->			Optional[Any]:

       """simple docstring"""




       a			   =DetaConfig(
           num_queries=900   ,       encoder_ffn_dim=2048   ,       decoder_ffn_dim=2048   ,       num_feature_levels=5   ,       assign_first_stage=UpperCAmelCase_   ,       with_box_refine=UpperCAmelCase_   ,       two_stage=UpperCAmelCase_   ,       )

       # set labels
       a			   =91
       a			   ="""huggingface/label-files"""
       a			   ="""coco-detection-id2label.json"""
       a			   =json.load(open(cached_download(hf_hub_url(UpperCAmelCase_   ,       UpperCAmelCase_   ,       repo_type="""dataset"""      )      )   ,       """r"""      )      )
       a			   ={int(UpperCAmelCase_      ): v for k, v in idalabel.items()}
       a			   =idalabel
       a			   ={v: k for k, v in idalabel.items()}

       return config





def 							lowerCamelCase    (							UpperCAmelCase_	:      Optional[Any]      )->			Optional[Any]:

       """simple docstring"""




       a			   =[]

       # stem
       # fmt: off
       rename_keys.append(("""backbone.0.body.conv1.weight""", """model.backbone.model.embedder.embedder.convolution.weight""")      )
       rename_keys.append(("""backbone.0.body.bn1.weight""", """model.backbone.model.embedder.embedder.normalization.weight""")      )
       rename_keys.append(("""backbone.0.body.bn1.bias""", """model.backbone.model.embedder.embedder.normalization.bias""")      )
       rename_keys.append(("""backbone.0.body.bn1.running_mean""", """model.backbone.model.embedder.embedder.normalization.running_mean""")      )
       rename_keys.append(("""backbone.0.body.bn1.running_var""", """model.backbone.model.embedder.embedder.normalization.running_var""")      )
       # stages
       for stage_idx in range(len(config.backbone_config.depths      )      ):
              for layer_idx in range(config.backbone_config.depths[stage_idx]      ):
                     # shortcut
                     if layer_idx == 0:
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var''',
                                )      )
                     # 3 convs
                     for i in range(3      ):
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean''',
                                )      )
                            rename_keys.append(
                                (
                                    F'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var''',
                                    F'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var''',
                                )      )
    # 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 							lowerCamelCase    (							UpperCAmelCase_	:      str   ,       UpperCAmelCase_	:      List[Any]   ,       UpperCAmelCase_	:      List[str]      )->			int:

       """simple docstring"""




       a			   =dct.pop(UpperCAmelCase_      )
       a			   =val





def 							lowerCamelCase    (							UpperCAmelCase_	:      List[Any]   ,       UpperCAmelCase_	:      Dict      )->			int:

       """simple docstring"""




       a			   =config.d_model
       for i in range(config.decoder_layers      ):
              # read in weights + bias of input projection layer of self-attention
              a			   =state_dict.pop(F'''transformer.decoder.layers.{i}.self_attn.in_proj_weight'''      )
              a			   =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
              a			   =in_proj_weight[:hidden_size, :]
              a			   =in_proj_bias[:hidden_size]
              a			   =in_proj_weight[
                  hidden_size : hidden_size * 2, :
              ]
              a			   =in_proj_bias[hidden_size : hidden_size * 2]
              a			   =in_proj_weight[-hidden_size:, :]
              a			   =in_proj_bias[-hidden_size:]





def 							lowerCamelCase    (							)->			List[str]:

       """simple docstring"""




       a			   ="""http://images.cocodataset.org/val2017/000000039769.jpg"""
       a			   =Image.open(requests.get(UpperCAmelCase_   ,       stream=UpperCAmelCase_      ).raw      )

       return im





@torch.no_grad()
def 							lowerCamelCase    (							UpperCAmelCase_	:      str   ,       UpperCAmelCase_	:      Optional[int]   ,       UpperCAmelCase_	:      List[Any]      )->			Any:

       """simple docstring"""




       a			   =get_deta_config()

       # load original state dict
       if model_name == "deta-resnet-50":
              a			   ="""adet_checkpoint0011.pth"""
       elif model_name == "deta-resnet-50-24-epochs":
              a			   ="""adet_2x_checkpoint0023.pth"""
       else:
              raise ValueError(F'''Model name {model_name} not supported'''      )
       a			   =hf_hub_download(repo_id="""nielsr/deta-checkpoints"""   ,       filename=UpperCAmelCase_      )
       a			   =torch.load(UpperCAmelCase_   ,       map_location="""cpu"""      )["""model"""]

       # rename keys
       a			   =create_rename_keys(UpperCAmelCase_      )
       for src, dest in rename_keys:
              rename_key(UpperCAmelCase_   ,       UpperCAmelCase_   ,       UpperCAmelCase_      )
       read_in_decoder_q_k_v(UpperCAmelCase_   ,       UpperCAmelCase_      )

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

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

       a			   ="""cuda""" if torch.cuda.is_available() else """cpu"""
       model.to(UpperCAmelCase_      )

       # load image processor
       a			   =DetaImageProcessor(format="""coco_detection"""      )

       # verify our conversion on image
       a			   =prepare_img()
       a			   =processor(images=UpperCAmelCase_   ,       return_tensors="""pt"""      )
       a			   =encoding["""pixel_values"""]
       a			   =model(pixel_values.to(UpperCAmelCase_      )      )

       # verify logits
       if model_name == "deta-resnet-50":
              a			   =torch.tensor(
                  [[-7.39_78, -2.54_06, -4.16_68], [-8.26_84, -3.99_33, -3.80_96], [-7.05_15, -3.79_73, -5.85_16]]      )
              a			   =torch.tensor([[0.50_43, 0.49_73, 0.99_98], [0.25_42, 0.54_89, 0.47_48], [0.54_90, 0.27_65, 0.05_70]]      )
       elif model_name == "deta-resnet-50-24-epochs":
              a			   =torch.tensor(
                  [[-7.16_88, -2.48_57, -4.86_69], [-7.86_30, -3.81_54, -4.26_74], [-7.27_30, -4.18_65, -5.53_23]]      )
              a			   =torch.tensor([[0.50_21, 0.49_71, 0.99_94], [0.25_46, 0.54_86, 0.47_31], [0.16_86, 0.19_86, 0.21_42]]      )

       assert torch.allclose(outputs.logits[0, :3, :3]   ,       expected_logits.to(UpperCAmelCase_      )   ,       atol=1e-4      )
       assert torch.allclose(outputs.pred_boxes[0, :3, :3]   ,       expected_boxes.to(UpperCAmelCase_      )   ,       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(UpperCAmelCase_      ).mkdir(exist_ok=UpperCAmelCase_      )
              model.save_pretrained(UpperCAmelCase_      )
              processor.save_pretrained(UpperCAmelCase_      )

       # 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-resnet-50''',
        choices=['''deta-resnet-50''', '''deta-resnet-50-24-epochs'''],
        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)






import math
from timeit import timeit





def 							lowerCamelCase    (							UpperCAmelCase_	:      int      )->			int:

       """simple docstring"""




       a			   =0
       a			   =abs(UpperCAmelCase_      )
       while True:
              a			   =n // 10
              digits += 1
              if n == 0:
                     break
       return digits





def 							lowerCamelCase    (							UpperCAmelCase_	:      int      )->			int:

       """simple docstring"""




       return 1 if n == 0 else math.floor(math.log(abs(UpperCAmelCase_      )   ,       10      ) + 1      )





def 							lowerCamelCase    (							UpperCAmelCase_	:      int      )->			int:

       """simple docstring"""




       return len(str(abs(UpperCAmelCase_      )      )      )





def 							lowerCamelCase    (							)->			None:

       """simple docstring"""




       from collections.abc import Callable

       def benchmark_a_function(UpperCAmelCase_	:      Callable   ,       UpperCAmelCase_	:      int      ) -> None:
              a			   =F'''{func.__name__}({value})'''
              a			   =timeit(F'''__main__.{call}'''   ,       setup="""import __main__"""      )
              print(F'''{call}: {func(UpperCAmelCase_      )} -- {timing} seconds'''      )

       for value in (26_2144, 1125_8999_0684_2624, 126_7650_6002_2822_9401_4967_0320_5376):
              for func in (num_digits, num_digits_fast, num_digits_faster):
                     benchmark_a_function(UpperCAmelCase_   ,       UpperCAmelCase_      )
              print()


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    benchmark()
