





'''simple docstring'''

import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase=False   )    ->       Any:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							[]
   for i in range(config.num_hidden_layers   ):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append((f"""module.blocks.{i}.norm1.weight""", f"""vit.encoder.layer.{i}.layernorm_before.weight""")   )
      rename_keys.append((f"""module.blocks.{i}.norm1.bias""", f"""vit.encoder.layer.{i}.layernorm_before.bias""")   )
      rename_keys.append(
          (f"""module.blocks.{i}.attn.proj.weight""", f"""vit.encoder.layer.{i}.attention.output.dense.weight""")   )
      rename_keys.append((f"""module.blocks.{i}.attn.proj.bias""", f"""vit.encoder.layer.{i}.attention.output.dense.bias""")   )
      rename_keys.append((f"""module.blocks.{i}.norm2.weight""", f"""vit.encoder.layer.{i}.layernorm_after.weight""")   )
      rename_keys.append((f"""module.blocks.{i}.norm2.bias""", f"""vit.encoder.layer.{i}.layernorm_after.bias""")   )
      rename_keys.append((f"""module.blocks.{i}.mlp.fc1.weight""", f"""vit.encoder.layer.{i}.intermediate.dense.weight""")   )
      rename_keys.append((f"""module.blocks.{i}.mlp.fc1.bias""", f"""vit.encoder.layer.{i}.intermediate.dense.bias""")   )
      rename_keys.append((f"""module.blocks.{i}.mlp.fc2.weight""", f"""vit.encoder.layer.{i}.output.dense.weight""")   )
      rename_keys.append((f"""module.blocks.{i}.mlp.fc2.bias""", f"""vit.encoder.layer.{i}.output.dense.bias""")   )

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ("""module.cls_token""", """vit.embeddings.cls_token"""),
           ("""module.patch_embed.proj.weight""", """vit.embeddings.patch_embeddings.projection.weight"""),
           ("""module.patch_embed.proj.bias""", """vit.embeddings.patch_embeddings.projection.bias"""),
           ("""module.pos_embed""", """vit.embeddings.position_embeddings"""),
       ]   )

   if base_model:
      # layernorm + pooler
      rename_keys.extend(
          [
              ("""module.norm.weight""", """layernorm.weight"""),
              ("""module.norm.bias""", """layernorm.bias"""),
          ]   )

      # if just the base model, we should remove "vit" from all keys that start with "vit"
      __SCREAMING_SNAKE_CASE   				=							[(pair[0], pair[1][4:]) if pair[1].startswith("""vit"""   ) else pair for pair in rename_keys]
   else:
      # layernorm + classification head
      rename_keys.extend(
          [
              ("""norm.weight""", """vit.layernorm.weight"""),
              ("""norm.bias""", """vit.layernorm.bias"""),
              ("""head.weight""", """classifier.weight"""),
              ("""head.bias""", """classifier.bias"""),
          ]   )

   return rename_keys

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase=False   )    ->       str:

   '''simple docstring'''







   for i in range(config.num_hidden_layers   ):
      if base_model:
         __SCREAMING_SNAKE_CASE   				=							""""""
      else:
         __SCREAMING_SNAKE_CASE   				=							"""vit."""
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      __SCREAMING_SNAKE_CASE   				=							state_dict.pop(f"""module.blocks.{i}.attn.qkv.weight"""   )
      __SCREAMING_SNAKE_CASE   				=							state_dict.pop(f"""module.blocks.{i}.attn.qkv.bias"""   )
      # next, add query, keys and values (in that order) to the state dict
      __SCREAMING_SNAKE_CASE   				=							in_proj_weight[
          : config.hidden_size, :
      ]
      __SCREAMING_SNAKE_CASE   				=							in_proj_bias[: config.hidden_size]
      __SCREAMING_SNAKE_CASE   				=							in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      __SCREAMING_SNAKE_CASE   				=							in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      __SCREAMING_SNAKE_CASE   				=							in_proj_weight[
          -config.hidden_size :, :
      ]
      __SCREAMING_SNAKE_CASE   				=							in_proj_bias[-config.hidden_size :]

def     __magic_name__							(				__UpperCAmelCase   )    ->       int:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(__UpperCAmelCase				,						__UpperCAmelCase   )

def     __magic_name__							(				__UpperCAmelCase   )    ->       Tuple:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							[
       """module.fc.fc1.weight""",
       """module.fc.fc1.bias""",
       """module.fc.bn1.weight""",
       """module.fc.bn1.bias""",
       """module.fc.bn1.running_mean""",
       """module.fc.bn1.running_var""",
       """module.fc.bn1.num_batches_tracked""",
       """module.fc.fc2.weight""",
       """module.fc.fc2.bias""",
       """module.fc.bn2.weight""",
       """module.fc.bn2.bias""",
       """module.fc.bn2.running_mean""",
       """module.fc.bn2.running_var""",
       """module.fc.bn2.num_batches_tracked""",
       """module.fc.fc3.weight""",
       """module.fc.fc3.bias""",
   ]
   for k in ignore_keys:
      state_dict.pop(__UpperCAmelCase				,						__UpperCAmelCase   )

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase   )    ->       str:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							dct.pop(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							val

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       Any:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							ViTMSNConfig()
   __SCREAMING_SNAKE_CASE   				=							1000

   __SCREAMING_SNAKE_CASE   				=							"""datasets/huggingface/label-files"""
   __SCREAMING_SNAKE_CASE   				=							"""imagenet-1k-id2label.json"""
   __SCREAMING_SNAKE_CASE   				=							json.load(open(hf_hub_download(__UpperCAmelCase				,						__UpperCAmelCase   )				,						"""r"""   )   )
   __SCREAMING_SNAKE_CASE   				=							{int(__UpperCAmelCase   ): v for k, v in idalabel.items()}
   __SCREAMING_SNAKE_CASE   				=							idalabel
   __SCREAMING_SNAKE_CASE   				=							{v: k for k, v in idalabel.items()}

   if "s16" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							384
      __SCREAMING_SNAKE_CASE   				=							1536
      __SCREAMING_SNAKE_CASE   				=							6
   elif "l16" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							1024
      __SCREAMING_SNAKE_CASE   				=							4096
      __SCREAMING_SNAKE_CASE   				=							24
      __SCREAMING_SNAKE_CASE   				=							16
      __SCREAMING_SNAKE_CASE   				=							0.1
   elif "b4" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							4
   elif "l7" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							7
      __SCREAMING_SNAKE_CASE   				=							1024
      __SCREAMING_SNAKE_CASE   				=							4096
      __SCREAMING_SNAKE_CASE   				=							24
      __SCREAMING_SNAKE_CASE   				=							16
      __SCREAMING_SNAKE_CASE   				=							0.1

   __SCREAMING_SNAKE_CASE   				=							ViTMSNModel(__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							torch.hub.load_state_dict_from_url(__UpperCAmelCase				,						map_location="""cpu"""   )["""target_encoder"""]

   __SCREAMING_SNAKE_CASE   				=							ViTImageProcessor(size=config.image_size   )

   remove_projection_head(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							create_rename_keys(__UpperCAmelCase				,						base_model=__UpperCAmelCase   )

   for src, dest in rename_keys:
      rename_key(__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase   )
   read_in_q_k_v(__UpperCAmelCase				,						__UpperCAmelCase				,						base_model=__UpperCAmelCase   )

   model.load_state_dict(__UpperCAmelCase   )
   model.eval()

   __SCREAMING_SNAKE_CASE   				=							"""http://images.cocodataset.org/val2017/000000039769.jpg"""

   __SCREAMING_SNAKE_CASE   				=							Image.open(requests.get(__UpperCAmelCase				,						stream=__UpperCAmelCase   ).raw   )
   __SCREAMING_SNAKE_CASE   				=							ViTImageProcessor(
       size=config.image_size				,						image_mean=__UpperCAmelCase				,						image_std=__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							image_processor(images=__UpperCAmelCase				,						return_tensors="""pt"""   )

   # forward pass
   torch.manual_seed(2   )
   __SCREAMING_SNAKE_CASE   				=							model(**__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							outputs.last_hidden_state

   # The following Colab Notebook was used to generate these outputs:
   # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
   if "s16" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]]   )
   elif "b16" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							torch.tensor([[1_4.2_8_8_9, -1_8.9_0_4_5, 1_1.7_2_8_1]]   )
   elif "l16" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							torch.tensor([[4_1.5_0_2_8, -2_2.8_6_8_1, 4_5.6_4_7_5]]   )
   elif "b4" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]]   )
   else:
      __SCREAMING_SNAKE_CASE   				=							torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]]   )

   # verify logits
   assert torch.allclose(last_hidden_state[:, 0, :3]				,						__UpperCAmelCase				,						atol=1e-4   )

   print(f"""Saving model to {pytorch_dump_folder_path}"""   )
   model.save_pretrained(__UpperCAmelCase   )

   print(f"""Saving image processor to {pytorch_dump_folder_path}"""   )
   image_processor.save_pretrained(__UpperCAmelCase   )


if __name__ == "__main__":
    a   	=    argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--checkpoint_url",
        default="https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar",
        type=str,
        help="URL of the checkpoint you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
    )

    a   	=    parser.parse_args()
    convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''
from typing import Dict, Iterable, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_pytesseract_available, is_vision_available, logging, requires_backends


if is_vision_available():
							import PIL

# soft dependency
if is_pytesseract_available():
							import pytesseract

__SCREAMING_SNAKE_CASE					:				List[str]									=		logging.get_logger(__name__)







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							return [
							    int(1_0_0_0 * (box[0] / width) ),
							    int(1_0_0_0 * (box[1] / height) ),
							    int(1_0_0_0 * (box[2] / width) ),
							    int(1_0_0_0 * (box[3] / height) ),
							]







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		to_pil_image(UpperCamelCase_ )
							A_    ,					A_			    =		pil_image.size
							A_			    =		pytesseract.image_to_data(UpperCamelCase_						,  lang=UpperCamelCase_						,  output_type="dict"						,  config=UpperCamelCase_ )
							A_    ,					A_    ,					A_    ,					A_    ,					A_			    =		data["text"], data["left"], data["top"], data["width"], data["height"]

							# filter empty words and corresponding coordinates
							A_			    =		[idx for idx, word in enumerate(UpperCamelCase_ ) if not word.strip()]
							A_			    =		[word for idx, word in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]

							# turn coordinates into (left, top, left+width, top+height) format
							A_			    =		[]
							for x, y, w, h in zip(UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
														A_			    =		[x, y, x + w, y + h]
														actual_boxes.append(UpperCamelCase_ )

							# finally, normalize the bounding boxes
							A_			    =		[]
							for box in actual_boxes:
														normalized_boxes.append(normalize_box(UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ) )

							assert len(UpperCamelCase_ ) == len(UpperCamelCase_ ), "Not as many words as there are bounding boxes"

							return words, normalized_boxes







class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Any   						=["pixel_values"]






						def __init__(  self							:      List[Any] ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Dict[str, int] = None ,  lowerCAmelCase							:      PILImageResampling = PILImageResampling.BILINEAR ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      float = 1 / 2_55 ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Optional[str] = None ,  lowerCAmelCase							:      Optional[str] = "" ,  **lowerCAmelCase							:      int ,  ):
													super().__init__(**lowerCAmelCase	)
													A_			    =		size if size is not None else {"height": 2_24, "width": 2_24}
													A_			    =		get_size_dict(lowerCAmelCase	)

													A_			    =		do_resize
													A_			    =		size
													A_			    =		resample
													A_			    =		do_rescale
													A_			    =		rescale_value
													A_			    =		do_normalize
													A_			    =		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
													A_			    =		image_std if image_std is not None else IMAGENET_STANDARD_STD
													A_			    =		apply_ocr
													A_			    =		ocr_lang
													A_			    =		tesseract_config






						def 				_UpperCAmelCase       (  self							:      List[str] ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Dict[str, int] ,  lowerCAmelCase							:      PILImageResampling = PILImageResampling.BILINEAR ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      List[str] ,  ):
													A_			    =		get_size_dict(lowerCAmelCase	)
													if "height" not in size or "width" not in size:
																				raise ValueError(F"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"	)
													A_			    =		(size["height"], size["width"])
													return resize(lowerCAmelCase ,  size=lowerCAmelCase ,  resample=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      List[Any] ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Union[int, float] ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      Tuple ,  ):
													return rescale(lowerCAmelCase ,  scale=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      int ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Union[float, Iterable[float]] ,  lowerCAmelCase							:      Union[float, Iterable[float]] ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      int ,  ):
													return normalize(lowerCAmelCase ,  mean=lowerCAmelCase ,  std=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Optional[Any] ,  lowerCAmelCase							:      ImageInput ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      Dict[str, int] = None ,  lowerCAmelCase							:      Any=None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      float = None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      Optional[str] = None ,  lowerCAmelCase							:      Optional[str] = None ,  lowerCAmelCase							:      Optional[Union[str, TensorType]] = None ,  lowerCAmelCase							:      ChannelDimension = ChannelDimension.FIRST ,  **lowerCAmelCase							:      int ,  ):
													A_			    =		do_resize if do_resize is not None else self.do_resize
													A_			    =		size if size is not None else self.size
													A_			    =		get_size_dict(lowerCAmelCase	)
													A_			    =		resample if resample is not None else self.resample
													A_			    =		do_rescale if do_rescale is not None else self.do_rescale
													A_			    =		rescale_factor if rescale_factor is not None else self.rescale_factor
													A_			    =		do_normalize if do_normalize is not None else self.do_normalize
													A_			    =		image_mean if image_mean is not None else self.image_mean
													A_			    =		image_std if image_std is not None else self.image_std
													A_			    =		apply_ocr if apply_ocr is not None else self.apply_ocr
													A_			    =		ocr_lang if ocr_lang is not None else self.ocr_lang
													A_			    =		tesseract_config if tesseract_config is not None else self.tesseract_config

													A_			    =		make_list_of_images(lowerCAmelCase	)

													if not valid_images(lowerCAmelCase	):
																				raise ValueError(
																				    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
																				    "torch.Tensor, tf.Tensor or jax.ndarray."	)

													if do_resize and size is None:
																				raise ValueError("Size must be specified if do_resize is True."	)

													if do_rescale and rescale_factor is None:
																				raise ValueError("Rescale factor must be specified if do_rescale is True."	)

													if do_normalize and (image_mean is None or image_std is None):
																				raise ValueError("If do_normalize is True, image_mean and image_std must be specified."	)

													# All transformations expect numpy arrays.
													A_			    =		[to_numpy_array(lowerCAmelCase	) for image in images]

													# Tesseract OCR to get words + normalized bounding boxes
													if apply_ocr:
																				requires_backends(self ,  "pytesseract"	)
																				A_			    =		[]
																				A_			    =		[]
																				for image in images:
																											A_    ,					A_			    =		apply_tesseract(lowerCAmelCase ,  lowerCAmelCase ,  lowerCAmelCase	)
																											words_batch.append(lowerCAmelCase	)
																											boxes_batch.append(lowerCAmelCase	)

													if do_resize:
																				A_			    =		[self.resize(image=lowerCAmelCase ,  size=lowerCAmelCase ,  resample=lowerCAmelCase	) for image in images]

													if do_rescale:
																				A_			    =		[self.rescale(image=lowerCAmelCase ,  scale=lowerCAmelCase	) for image in images]

													if do_normalize:
																				A_			    =		[self.normalize(image=lowerCAmelCase ,  mean=lowerCAmelCase ,  std=lowerCAmelCase	) for image in images]

													A_			    =		[to_channel_dimension_format(lowerCAmelCase ,  lowerCAmelCase	) for image in images]

													A_			    =		BatchFeature(data={"pixel_values": images} ,  tensor_type=lowerCAmelCase	)

													if apply_ocr:
																				A_			    =		words_batch
																				A_			    =		boxes_batch
													return data





