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


logging.set_verbosity_info()
UpperCamelCase_						   =   logging.get_logger(__name__)
def    UpperCamelCase						(					)						->List[str]:

   """simple docstring"""



   a_						   =     DetaConfig(
       num_queries=900 ,   encoder_ffn_dim=2_048 ,   decoder_ffn_dim=2_048 ,   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    UpperCamelCase						(					UpperCAmelCase     )						->Optional[int]:

   """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    UpperCamelCase						(					UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase     )						->Optional[int]:

   """simple docstring"""



   a_						   =     dct.pop(UpperCAmelCase     )
   a_						   =     val
def    UpperCamelCase						(					UpperCAmelCase ,   UpperCAmelCase     )						->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    UpperCamelCase						(					)						->Union[str, Any]:

   """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    UpperCamelCase						(					UpperCAmelCase ,   UpperCAmelCase ,   UpperCAmelCase     )						->Optional[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.3978, -2.5406, -4.1668], [-8.2684, -3.9933, -3.8096], [-7.0515, -3.7973, -5.8516]]     )
      a_						   =     torch.tensor([[0.5043, 0.4973, 0.9998], [0.2542, 0.5489, 0.4748], [0.5490, 0.2765, 0.0570]]     )
   elif model_name == "deta-resnet-50-24-epochs":
      a_						   =     torch.tensor(
          [[-7.1688, -2.4857, -4.8669], [-7.8630, -3.8154, -4.2674], [-7.2730, -4.1865, -5.5323]]     )
      a_						   =     torch.tensor([[0.5021, 0.4971, 0.9994], [0.2546, 0.5486, 0.4731], [0.1686, 0.1986, 0.2142]]     )

   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__":
      UpperCamelCase_						   =   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.'
      )
      UpperCamelCase_						   =   parser.parse_args()
      convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_lowerCamelCase					:						Dict					     =							{
    'configuration_donut_swin': ['DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DonutSwinConfig'],
    'processing_donut': ['DonutProcessor'],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowerCamelCase					:						List[str]					     =							[
				    'DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'DonutSwinModel',
				    'DonutSwinPreTrainedModel',
				]

try:
				if not is_vision_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowerCamelCase					:						Any					     =							['DonutFeatureExtractor']
				_lowerCamelCase					:						Union[str, Any]					     =							['DonutImageProcessor']


if TYPE_CHECKING:
				from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
				from .processing_donut import DonutProcessor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_donut_swin import (
								    DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
								    DonutSwinModel,
								    DonutSwinPreTrainedModel,
								)

				try:
								if not is_vision_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .feature_extraction_donut import DonutFeatureExtractor
								from .image_processing_donut import DonutImageProcessor

else:
				import sys

				_lowerCamelCase					:						Optional[int]					     =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

