
"""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()
_snake_case          =						logging.get_logger(__name__)






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					List[Any]						):






							"""simple docstring"""



							_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=SCREAMING_SNAKE_CASE     ,  num_queries=900     ,  encoder_ffn_dim=2048     ,  decoder_ffn_dim=2048     ,  num_feature_levels=5     ,  assign_first_stage=SCREAMING_SNAKE_CASE     ,  with_box_refine=SCREAMING_SNAKE_CASE     ,  two_stage=SCREAMING_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(SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE     ,  repo_type='dataset'						)						)     ,  'r'						)						)
							_lowerCAmelCase  	=			{int(SCREAMING_SNAKE_CASE						): v for k, v in idalabel.items()}
							_lowerCAmelCase  	=			idalabel
							_lowerCAmelCase  	=			{v: k for k, v in idalabel.items()}

							return config






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					Tuple						):






							"""simple docstring"""



							_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 	__snake_case  (    SCREAMING_SNAKE_CASE:					Dict     ,  SCREAMING_SNAKE_CASE:					Tuple     ,  SCREAMING_SNAKE_CASE:					Tuple						):






							"""simple docstring"""



							_lowerCAmelCase  	=			dct.pop(SCREAMING_SNAKE_CASE						)
							_lowerCAmelCase  	=			val






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					Any     ,  SCREAMING_SNAKE_CASE:					Optional[int]						):






							"""simple docstring"""



							_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 	__snake_case  (    SCREAMING_SNAKE_CASE:					Any     ,  SCREAMING_SNAKE_CASE:					List[str]						):






							"""simple docstring"""



							_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 	__snake_case  (    ):






							"""simple docstring"""



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

							return im






@torch.no_grad()
def 	__snake_case  (    SCREAMING_SNAKE_CASE:					List[str]     ,  SCREAMING_SNAKE_CASE:					str     ,  SCREAMING_SNAKE_CASE:					Dict						):






							"""simple docstring"""



							_lowerCAmelCase  	=			get_deta_config(SCREAMING_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(SCREAMING_SNAKE_CASE     ,  map_location='cpu'						)['model']

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

							# rename keys
							_lowerCAmelCase  	=			create_rename_keys(SCREAMING_SNAKE_CASE						)
							for src, dest in rename_keys:
														rename_key(SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE						)
							read_in_swin_q_k_v(SCREAMING_SNAKE_CASE     ,  config.backbone_config						)
							read_in_decoder_q_k_v(SCREAMING_SNAKE_CASE     ,  SCREAMING_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(SCREAMING_SNAKE_CASE						)
																					_lowerCAmelCase  	=			val
														if "input_proj" in key:
																					_lowerCAmelCase  	=			state_dict.pop(SCREAMING_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(SCREAMING_SNAKE_CASE						)
																					_lowerCAmelCase  	=			val

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

							_lowerCAmelCase  	=			'cuda' if torch.cuda.is_available() else 'cpu'
							model.to(SCREAMING_SNAKE_CASE						)

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

							# verify our conversion on image
							_lowerCAmelCase  	=			prepare_img()
							_lowerCAmelCase  	=			processor(images=SCREAMING_SNAKE_CASE     ,  return_tensors='pt'						)
							_lowerCAmelCase  	=			encoding['pixel_values']
							_lowerCAmelCase  	=			model(pixel_values.to(SCREAMING_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.63_08, -2.84_85, -5.37_37], [-7.20_37, -4.55_05, -4.80_27], [-7.29_43, -4.26_11, -4.66_17]]						)
														_lowerCAmelCase  	=			torch.tensor([[0.49_87, 0.49_69, 0.99_99], [0.25_49, 0.54_98, 0.48_05], [0.54_98, 0.27_57, 0.05_69]]						)
							elif model_name == "deta-swin-large-o365":
														_lowerCAmelCase  	=			torch.tensor(
														    [[-8.01_22, -3.57_20, -4.97_17], [-8.15_47, -3.68_86, -4.63_89], [-7.66_10, -3.61_94, -5.01_34]]						)
														_lowerCAmelCase  	=			torch.tensor([[0.25_23, 0.55_49, 0.48_81], [0.77_15, 0.41_49, 0.46_01], [0.55_03, 0.27_53, 0.05_75]]						)
							assert torch.allclose(outputs.logits[0, :3, :3]     ,  expected_logits.to(SCREAMING_SNAKE_CASE						)     ,  atol=1e-4						)
							assert torch.allclose(outputs.pred_boxes[0, :3, :3]     ,  expected_boxes.to(SCREAMING_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(SCREAMING_SNAKE_CASE						).mkdir(exist_ok=SCREAMING_SNAKE_CASE						)
														model.save_pretrained(SCREAMING_SNAKE_CASE						)
														processor.save_pretrained(SCREAMING_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__":
	_snake_case          =						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.'''
	)
	_snake_case          =						parser.parse_args()
	convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








"""simple docstring"""



def 	__snake_case  (    SCREAMING_SNAKE_CASE:					list[int]						):






							"""simple docstring"""



							_lowerCAmelCase  	=			[]
							if len(SCREAMING_SNAKE_CASE						) == 1:
														return [nums.copy()]
							for _ in range(len(SCREAMING_SNAKE_CASE						)						):
														_lowerCAmelCase  	=			nums.pop(0						)
														_lowerCAmelCase  	=			permute(SCREAMING_SNAKE_CASE						)
														for perm in permutations:
																					perm.append(SCREAMING_SNAKE_CASE						)
														result.extend(SCREAMING_SNAKE_CASE						)
														nums.append(SCREAMING_SNAKE_CASE						)
							return result






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					Any						):






							"""simple docstring"""



							def backtrack(SCREAMING_SNAKE_CASE:					Tuple						):
														if start == len(SCREAMING_SNAKE_CASE						) - 1:
																					output.append(nums[:]						)
														else:
																					for i in range(SCREAMING_SNAKE_CASE     ,  len(SCREAMING_SNAKE_CASE						)						):
																												_lowerCAmelCase	,		_lowerCAmelCase  	=			nums[i], nums[start]
																												backtrack(start + 1						)
																												_lowerCAmelCase	,		_lowerCAmelCase  	=			nums[i], nums[start]  # backtrack

							_lowerCAmelCase  	=			[]
							backtrack(0						)
							return output


if __name__ == "__main__":
	import doctest

	# use res to print the data in permute2 function
	_snake_case          =						permutea([1, 2, 3])
	print(res)
	doctest.testmod()






