
from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
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, logging


__lowerCamelCase				=			logging.get_logger(__name__)






class     UpperCAmelCase						(		A_					):
   A__   :						Tuple			=	["pixel_values"]



   def __init__(self						:   List[Any]       ,		snake_case__						:   bool = True       ,		snake_case__						:   Optional[Dict[str, int]] = None       ,		snake_case__						:   PILImageResampling = PILImageResampling.BILINEAR       ,		snake_case__						:   bool = True       ,		snake_case__						:   Dict[str, int] = None       ,		snake_case__						:   bool = True       ,		snake_case__						:   Union[int, float] = 1 / 2_55       ,		snake_case__						:   bool = True       ,		snake_case__						:   Optional[Union[float, List[float]]] = None       ,		snake_case__						:   Optional[Union[float, List[float]]] = None       ,		**snake_case__						:   Dict       ,		)						->      None:


          '''simple docstring'''


          super().__init__(**snake_case__					)
          snake_case			:  List[str]     							= size if size is not None else {"shortest_edge": 2_56}
          snake_case			:  str     							= get_size_dict(snake_case__       ,		default_to_square=snake_case__					)
          snake_case			:  int     							= crop_size if crop_size is not None else {"height": 2_24, "width": 2_24}
          snake_case			:  Optional[int]     							= get_size_dict(snake_case__					)
          snake_case			:  Optional[int]     							= do_resize
          snake_case			:  Optional[int]     							= size
          snake_case			:  Tuple     							= resample
          snake_case			:  Tuple     							= do_center_crop
          snake_case			:  List[str]     							= crop_size
          snake_case			:  List[Any]     							= do_rescale
          snake_case			:  Optional[int]     							= rescale_factor
          snake_case			:  Dict     							= do_normalize
          snake_case			:  Dict     							= image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
          snake_case			:  Optional[int]     							= image_std if image_std is not None else IMAGENET_STANDARD_STD



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[int]       ,		snake_case__						:   np.ndarray       ,		snake_case__						:   Dict[str, int]       ,		snake_case__						:   PILImageResampling = PILImageResampling.BICUBIC       ,		snake_case__						:   Optional[Union[str, ChannelDimension]] = None       ,		**snake_case__						:   Tuple       ,		)						->      np.ndarray:


          '''simple docstring'''


          snake_case			:  Tuple     							= get_size_dict(snake_case__       ,		default_to_square=snake_case__					)
          if "shortest_edge" not in size:
                 raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""					)
          snake_case			:  int     							= get_resize_output_image_size(snake_case__       ,		size=size["shortest_edge"]       ,		default_to_square=snake_case__					)
          return resize(snake_case__       ,		size=snake_case__       ,		resample=snake_case__       ,		data_format=snake_case__       ,		**snake_case__					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   str       ,		snake_case__						:   np.ndarray       ,		snake_case__						:   Dict[str, int]       ,		snake_case__						:   Optional[Union[str, ChannelDimension]] = None       ,		**snake_case__						:   Union[str, Any]       ,		)						->      np.ndarray:


          '''simple docstring'''


          snake_case			:  Optional[int]     							= get_size_dict(snake_case__					)
          return center_crop(snake_case__       ,		size=(size["height"], size["width"])       ,		data_format=snake_case__       ,		**snake_case__					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[str]       ,		snake_case__						:   np.ndarray       ,		snake_case__						:   float       ,		snake_case__						:   Optional[Union[str, ChannelDimension]] = None       ,		**snake_case__						:   List[Any]					)						->      np.ndarray:


          '''simple docstring'''


          return rescale(snake_case__       ,		scale=snake_case__       ,		data_format=snake_case__       ,		**snake_case__					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[str]       ,		snake_case__						:   np.ndarray       ,		snake_case__						:   Union[float, List[float]]       ,		snake_case__						:   Union[float, List[float]]       ,		snake_case__						:   Optional[Union[str, ChannelDimension]] = None       ,		**snake_case__						:   int       ,		)						->      np.ndarray:


          '''simple docstring'''


          return normalize(snake_case__       ,		mean=snake_case__       ,		std=snake_case__       ,		data_format=snake_case__       ,		**snake_case__					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[int]       ,		snake_case__						:   ImageInput       ,		snake_case__						:   Optional[bool] = None       ,		snake_case__						:   Dict[str, int] = None       ,		snake_case__						:   PILImageResampling = None       ,		snake_case__						:   bool = None       ,		snake_case__						:   Dict[str, int] = None       ,		snake_case__						:   Optional[bool] = None       ,		snake_case__						:   Optional[float] = None       ,		snake_case__						:   Optional[bool] = None       ,		snake_case__						:   Optional[Union[float, List[float]]] = None       ,		snake_case__						:   Optional[Union[float, List[float]]] = None       ,		snake_case__						:   Optional[Union[str, TensorType]] = None       ,		snake_case__						:   Union[str, ChannelDimension] = ChannelDimension.FIRST       ,		**snake_case__						:   Union[str, Any]       ,		)						->      str:


          '''simple docstring'''


          snake_case			:  str     							= do_resize if do_resize is not None else self.do_resize
          snake_case			:  List[str]     							= size if size is not None else self.size
          snake_case			:  List[str]     							= get_size_dict(snake_case__       ,		default_to_square=snake_case__					)
          snake_case			:  Any     							= resample if resample is not None else self.resample
          snake_case			:  Optional[int]     							= do_center_crop if do_center_crop is not None else self.do_center_crop
          snake_case			:  str     							= crop_size if crop_size is not None else self.crop_size
          snake_case			:  int     							= get_size_dict(snake_case__					)
          snake_case			:  int     							= do_rescale if do_rescale is not None else self.do_rescale
          snake_case			:  Optional[Any]     							= rescale_factor if rescale_factor is not None else self.rescale_factor
          snake_case			:  List[Any]     							= do_normalize if do_normalize is not None else self.do_normalize
          snake_case			:  Dict     							= image_mean if image_mean is not None else self.image_mean
          snake_case			:  int     							= image_std if image_std is not None else self.image_std

          snake_case			:  Dict     							= make_list_of_images(snake_case__					)

          if not valid_images(snake_case__					):
                 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_center_crop and crop_size is None:
                 raise ValueError("Crop size must be specified if do_center_crop 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("Image mean and std must be specified if do_normalize is True."					)

          # All transformations expect numpy arrays.
          snake_case			:  Dict     							= [to_numpy_array(snake_case__					) for image in images]

          if do_resize:
                 snake_case			:  Union[str, Any]     							= [self.resize(image=snake_case__       ,		size=snake_case__       ,		resample=snake_case__					) for image in images]

          if do_center_crop:
                 snake_case			:  int     							= [self.center_crop(image=snake_case__       ,		size=snake_case__					) for image in images]

          if do_rescale:
                 snake_case			:  Dict     							= [self.rescale(image=snake_case__       ,		scale=snake_case__					) for image in images]

          if do_normalize:
                 snake_case			:  Optional[Any]     							= [self.normalize(image=snake_case__       ,		mean=snake_case__       ,		std=snake_case__					) for image in images]

          snake_case			:  Union[str, Any]     							= [to_channel_dimension_format(snake_case__       ,		snake_case__					) for image in images]

          snake_case			:  Optional[Any]     							= {"pixel_values": images}
          return BatchFeature(data=snake_case__       ,		tensor_type=snake_case__					)







import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


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

__lowerCamelCase				=			[
    ("""bert.bert""", """visual_bert"""),
    ("""bert.cls""", """cls"""),
    ("""bert.classifier""", """cls"""),
    ("""token_type_embeddings_visual""", """visual_token_type_embeddings"""),
    ("""position_embeddings_visual""", """visual_position_embeddings"""),
    ("""projection""", """visual_projection"""),
]

__lowerCamelCase				=			[
    """nlvr2_coco_pre_trained.th""",
    """nlvr2_fine_tuned.th""",
    """nlvr2_pre_trained.th""",
    """vcr_coco_pre_train.th""",
    """vcr_fine_tune.th""",
    """vcr_pre_train.th""",
    """vqa_coco_pre_trained.th""",
    """vqa_fine_tuned.th""",
    """vqa_pre_trained.th""",
]


def      UpperCamelCase   (       __lowerCamelCase	:					int       ):
       snake_case			:  List[Any]     							= torch.load(__lowerCamelCase							,   map_location="cpu"       )
       return sd


def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[int]							,   __lowerCamelCase	:					Optional[Any]=rename_keys_prefix       ):
       snake_case			:  Optional[Any]     							= OrderedDict()
       snake_case			:  Optional[int]     							= torch.arange(config.max_position_embeddings       ).expand((1, -1)       )
       # detector_d = OrderedDict()
       for key in d:
              if "detector" in key:
                     # detector_d[key.replace('detector.','')] = d[key]
                     continue
              snake_case			:  Any     							= key
              for name_pair in rename_keys_prefix:
                     snake_case			:  List[Any]     							= new_key.replace(name_pair[0]							,   name_pair[1]       )
              snake_case			:  List[str]     							= d[key]
              if key == "bert.cls.predictions.decoder.weight":
                     # Old bert code didn't have `decoder.bias`, but was added separately
                     snake_case			:  int     							= new_d["cls.predictions.bias"]
       return new_d


@torch.no_grad()
def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[Any]       ):
       assert (
           checkpoint_path.split("/"       )[-1] in ACCEPTABLE_CHECKPOINTS
       ), f"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

       # Get Config
       if "pre" in checkpoint_path:
              snake_case			:  Optional[Any]     							= "pretraining"
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[int]     							= {"visual_embedding_dim": 512}
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  int     							= {"visual_embedding_dim": 2048}
              elif "vqa" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
              elif "nlvr" in checkpoint_path:
                     snake_case			:  Tuple     							= {"visual_embedding_dim": 1024}
              else:
                     raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""       )
       else:
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 512}
                     snake_case			:  List[str]     							= "multichoice"
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
                     snake_case			:  Tuple     							= "vqa_advanced"
              elif "vqa" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 2048, "num_labels": 3129}
                     snake_case			:  Tuple     							= "vqa"
              elif "nlvr" in checkpoint_path:
                     snake_case			:  str     							= {
                         "visual_embedding_dim": 1024,
                         "num_labels": 2,
                     }
                     snake_case			:  str     							= "nlvr"

       snake_case			:  int     							= VisualBertConfig(**__lowerCamelCase       )

       # Load State Dict
       snake_case			:  Any     							= load_state_dict(__lowerCamelCase       )

       snake_case			:  List[Any]     							= get_new_dict(__lowerCamelCase							,   __lowerCamelCase       )

       if model_type == "pretraining":
              snake_case			:  str     							= VisualBertForPreTraining(__lowerCamelCase       )
       elif model_type == "vqa":
              snake_case			:  Any     							= VisualBertForQuestionAnswering(__lowerCamelCase       )
       elif model_type == "nlvr":
              snake_case			:  Optional[Any]     							= VisualBertForVisualReasoning(__lowerCamelCase       )
       elif model_type == "multichoice":
              snake_case			:  Optional[int]     							= VisualBertForMultipleChoice(__lowerCamelCase       )

       model.load_state_dict(__lowerCamelCase       )
       # Save Checkpoints
       Path(__lowerCamelCase       ).mkdir(exist_ok=__lowerCamelCase       )
       model.save_pretrained(__lowerCamelCase       )


if __name__ == "__main__":
   __lowerCamelCase				=			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument("""orig_checkpoint_path""", type=str, help="""A path to .th on local filesystem.""")
   parser.add_argument("""pytorch_dump_folder_path""", type=str, help="""Path to the output PyTorch model.""")
   __lowerCamelCase				=			parser.parse_args()
   convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)





