






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      :       Union[str, Any]    =						logging.get_logger(__name__)





def    __lowerCAmelCase   (    __snake_case							):
     __lowerCAmelCase             =       SwinConfig(
         embed_dim=192       ,			depths=(2, 2, 18, 2)       ,			num_heads=(6, 12, 24, 48)       ,			window_size=12       ,			out_features=["stage2", "stage3", "stage4"]       ,			)

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

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

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

     return config





def    __lowerCAmelCase   (    __snake_case							):
     __lowerCAmelCase             =       []

     # 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    __lowerCAmelCase   (    __snake_case       ,			__snake_case       ,			__snake_case							):
     __lowerCAmelCase             =       dct.pop(__snake_case							)
     __lowerCAmelCase             =       val





def    __lowerCAmelCase   (    __snake_case       ,			__snake_case							):
     __lowerCAmelCase             =       [int(backbone_config.embed_dim * 2**i							) for i in range(len(backbone_config.depths							)							)]
     for i in range(len(backbone_config.depths							)							):
          __lowerCAmelCase             =       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)
               __lowerCAmelCase             =       state_dict.pop(F"""backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight"""							)
               __lowerCAmelCase             =       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
               __lowerCAmelCase             =       in_proj_weight[:dim, :]
               __lowerCAmelCase             =       in_proj_bias[: dim]
               __lowerCAmelCase             =       in_proj_weight[
                   dim : dim * 2, :
               ]
               __lowerCAmelCase             =       in_proj_bias[
                   dim : dim * 2
               ]
               __lowerCAmelCase             =       in_proj_weight[
                   -dim :, :
               ]
               __lowerCAmelCase             =       in_proj_bias[-dim :]
               # fmt: on





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





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

     return im





@torch.no_grad()
def    __lowerCAmelCase   (    __snake_case       ,			__snake_case       ,			__snake_case							):
     __lowerCAmelCase             =       get_deta_config(__snake_case							)

     # load original state dict
     if model_name == "deta-swin-large":
          __lowerCAmelCase             =       hf_hub_download(repo_id="nielsr/deta-checkpoints"       ,			filename="adet_swin_ft.pth"							)
     elif model_name == "deta-swin-large-o365":
          __lowerCAmelCase             =       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"""							)

     __lowerCAmelCase             =       torch.load(__snake_case       ,			map_location="cpu"							)["model"]

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

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

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

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

     __lowerCAmelCase             =       "cuda" if torch.cuda.is_available() else "cpu"
     model.to(__snake_case							)

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

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

     # verify logits
     print("Logits:"       ,			outputs.logits[0, :3, :3]							)
     print("Boxes:"       ,			outputs.pred_boxes[0, :3, :3]							)
     if model_name == "deta-swin-large":
          __lowerCAmelCase             =       torch.tensor(
              [[-7.6_308, -2.8_485, -5.3_737], [-7.2_037, -4.5_505, -4.8_027], [-7.2_943, -4.2_611, -4.6_617]]							)
          __lowerCAmelCase             =       torch.tensor([[0.4_987, 0.4_969, 0.9_999], [0.2_549, 0.5_498, 0.4_805], [0.5_498, 0.2_757, 0.0_569]]							)
     elif model_name == "deta-swin-large-o365":
          __lowerCAmelCase             =       torch.tensor(
              [[-8.0_122, -3.5_720, -4.9_717], [-8.1_547, -3.6_886, -4.6_389], [-7.6_610, -3.6_194, -5.0_134]]							)
          __lowerCAmelCase             =       torch.tensor([[0.2_523, 0.5_549, 0.4_881], [0.7_715, 0.4_149, 0.4_601], [0.5_503, 0.2_753, 0.0_575]]							)
     assert torch.allclose(outputs.logits[0, :3, :3]       ,			expected_logits.to(__snake_case							)       ,			atol=1E-4							)
     assert torch.allclose(outputs.pred_boxes[0, :3, :3]       ,			expected_boxes.to(__snake_case							)       ,			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(__snake_case							).mkdir(exist_ok=__snake_case							)
          model.save_pretrained(__snake_case							)
          processor.save_pretrained(__snake_case							)

     # 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      :       List[Any]    =						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      :       List[str]    =						parser.parse_args()
 convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)














from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowerCamelCase      :       List[Any]    =						{
    '''configuration_owlvit''': [
        '''OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''OwlViTConfig''',
        '''OwlViTOnnxConfig''',
        '''OwlViTTextConfig''',
        '''OwlViTVisionConfig''',
    ],
    '''processing_owlvit''': ['''OwlViTProcessor'''],
}


try:
 if not is_vision_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase      :       List[Any]    =						['''OwlViTFeatureExtractor''']
 lowerCamelCase      :       str    =						['''OwlViTImageProcessor''']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase      :       List[str]    =						[
     '''OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''OwlViTModel''',
     '''OwlViTPreTrainedModel''',
     '''OwlViTTextModel''',
     '''OwlViTVisionModel''',
     '''OwlViTForObjectDetection''',
 ]

if TYPE_CHECKING:
 from .configuration_owlvit import (
     OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
     OwlViTConfig,
     OwlViTOnnxConfig,
     OwlViTTextConfig,
     OwlViTVisionConfig,
 )
 from .processing_owlvit import OwlViTProcessor

 try:
  if not is_vision_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_owlvit import OwlViTFeatureExtractor
  from .image_processing_owlvit import OwlViTImageProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_owlvit import (
      OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
      OwlViTForObjectDetection,
      OwlViTModel,
      OwlViTPreTrainedModel,
      OwlViTTextModel,
      OwlViTVisionModel,
  )

else:
 import sys

 lowerCamelCase      :       List[str]    =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






