


'''simple docstring'''







import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase    :     str					=			logging.get_logger(__name__)

lowerCAmelCase    :     Union[str, Any]					=			{
    '''b0''': efficientnet.EfficientNetBa,
    '''b1''': efficientnet.EfficientNetBa,
    '''b2''': efficientnet.EfficientNetBa,
    '''b3''': efficientnet.EfficientNetBa,
    '''b4''': efficientnet.EfficientNetBa,
    '''b5''': efficientnet.EfficientNetBa,
    '''b6''': efficientnet.EfficientNetBa,
    '''b7''': efficientnet.EfficientNetBa,
}

lowerCAmelCase    :     Dict					=			{
    '''b0''': {
        '''hidden_dim''': 12_80,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.0,
        '''image_size''': 2_24,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [],
    },
    '''b1''': {
        '''hidden_dim''': 12_80,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.1,
        '''image_size''': 2_40,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [16],
    },
    '''b2''': {
        '''hidden_dim''': 14_08,
        '''width_coef''': 1.1,
        '''depth_coef''': 1.2,
        '''image_size''': 2_60,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 8, 16],
    },
    '''b3''': {
        '''hidden_dim''': 15_36,
        '''width_coef''': 1.2,
        '''depth_coef''': 1.4,
        '''image_size''': 3_00,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 18],
    },
    '''b4''': {
        '''hidden_dim''': 17_92,
        '''width_coef''': 1.4,
        '''depth_coef''': 1.8,
        '''image_size''': 3_80,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [6],
    },
    '''b5''': {
        '''hidden_dim''': 20_48,
        '''width_coef''': 1.6,
        '''depth_coef''': 2.2,
        '''image_size''': 4_56,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [13, 27],
    },
    '''b6''': {
        '''hidden_dim''': 23_04,
        '''width_coef''': 1.8,
        '''depth_coef''': 2.6,
        '''image_size''': 5_28,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [31],
    },
    '''b7''': {
        '''hidden_dim''': 25_60,
        '''width_coef''': 2.0,
        '''depth_coef''': 3.1,
        '''image_size''': 6_00,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [18],
    },
}






def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							Tuple   			=       EfficientNetConfig()
       _lowerCAmelCase    :							int   			=       CONFIG_MAP[model_name]['hidden_dim']
       _lowerCAmelCase    :							List[str]   			=       CONFIG_MAP[model_name]['width_coef']
       _lowerCAmelCase    :							Any   			=       CONFIG_MAP[model_name]['depth_coef']
       _lowerCAmelCase    :							Any   			=       CONFIG_MAP[model_name]['image_size']
       _lowerCAmelCase    :							int   			=       CONFIG_MAP[model_name]['dropout_rate']
       _lowerCAmelCase    :							Any   			=       CONFIG_MAP[model_name]['dw_padding']

       _lowerCAmelCase    :							Optional[Any]   			=       'huggingface/label-files'
       _lowerCAmelCase    :							Union[str, Any]   			=       'imagenet-1k-id2label.json'
       _lowerCAmelCase    :							Tuple   			=       1_0_0_0
       _lowerCAmelCase    :							Union[str, Any]   			=       json.load(open(hf_hub_download(_A   ,						_A   ,						repo_type='dataset'					)   ,						'r'					)					)
       _lowerCAmelCase    :							Optional[Any]   			=       {int(_A					): v for k, v in idalabel.items()}

       _lowerCAmelCase    :							Dict   			=       idalabel
       _lowerCAmelCase    :							Union[str, Any]   			=       {v: k for k, v in idalabel.items()}
       return config






def 			lowercase				():







       """simple docstring"""







       _lowerCAmelCase    :							Tuple   			=       'http://images.cocodataset.org/val2017/000000039769.jpg'
       _lowerCAmelCase    :							Any   			=       Image.open(requests.get(_A   ,						stream=_A					).raw					)
       return im






def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							Optional[int]   			=       CONFIG_MAP[model_name]['image_size']
       _lowerCAmelCase    :							Optional[Any]   			=       EfficientNetImageProcessor(
           size={'height': size, 'width': size}   ,						image_mean=[0.485, 0.456, 0.406]   ,						image_std=[0.47_853_944, 0.4_732_864, 0.47_434_163]   ,						do_center_crop=_A   ,						)
       return preprocessor






def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							Any   			=       [v.split('_'					)[0].split('block'					)[1] for v in original_param_names if v.startswith('block'					)]
       _lowerCAmelCase    :							List[str]   			=       sorted(set(_A					)					)
       _lowerCAmelCase    :							Optional[int]   			=       len(_A					)
       _lowerCAmelCase    :							str   			=       {b: str(_A					) for b, i in zip(_A   ,						range(_A					)					)}

       _lowerCAmelCase    :							Tuple   			=       []
       rename_keys.append(('stem_conv/kernel:0', 'embeddings.convolution.weight')					)
       rename_keys.append(('stem_bn/gamma:0', 'embeddings.batchnorm.weight')					)
       rename_keys.append(('stem_bn/beta:0', 'embeddings.batchnorm.bias')					)
       rename_keys.append(('stem_bn/moving_mean:0', 'embeddings.batchnorm.running_mean')					)
       rename_keys.append(('stem_bn/moving_variance:0', 'embeddings.batchnorm.running_var')					)

       for b in block_names:
              _lowerCAmelCase    :							Optional[Any]   			=       block_name_mapping[b]
              rename_keys.append((f'block{b}_expand_conv/kernel:0', f'encoder.blocks.{hf_b}.expansion.expand_conv.weight')					)
              rename_keys.append((f'block{b}_expand_bn/gamma:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.weight')					)
              rename_keys.append((f'block{b}_expand_bn/beta:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.bias')					)
              rename_keys.append(
                  (f'block{b}_expand_bn/moving_mean:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.running_mean')					)
              rename_keys.append(
                  (f'block{b}_expand_bn/moving_variance:0', f'encoder.blocks.{hf_b}.expansion.expand_bn.running_var')					)
              rename_keys.append(
                  (f'block{b}_dwconv/depthwise_kernel:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight')					)
              rename_keys.append((f'block{b}_bn/gamma:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight')					)
              rename_keys.append((f'block{b}_bn/beta:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias')					)
              rename_keys.append(
                  (f'block{b}_bn/moving_mean:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean')					)
              rename_keys.append(
                  (f'block{b}_bn/moving_variance:0', f'encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var')					)

              rename_keys.append((f'block{b}_se_reduce/kernel:0', f'encoder.blocks.{hf_b}.squeeze_excite.reduce.weight')					)
              rename_keys.append((f'block{b}_se_reduce/bias:0', f'encoder.blocks.{hf_b}.squeeze_excite.reduce.bias')					)
              rename_keys.append((f'block{b}_se_expand/kernel:0', f'encoder.blocks.{hf_b}.squeeze_excite.expand.weight')					)
              rename_keys.append((f'block{b}_se_expand/bias:0', f'encoder.blocks.{hf_b}.squeeze_excite.expand.bias')					)
              rename_keys.append(
                  (f'block{b}_project_conv/kernel:0', f'encoder.blocks.{hf_b}.projection.project_conv.weight')					)
              rename_keys.append((f'block{b}_project_bn/gamma:0', f'encoder.blocks.{hf_b}.projection.project_bn.weight')					)
              rename_keys.append((f'block{b}_project_bn/beta:0', f'encoder.blocks.{hf_b}.projection.project_bn.bias')					)
              rename_keys.append(
                  (f'block{b}_project_bn/moving_mean:0', f'encoder.blocks.{hf_b}.projection.project_bn.running_mean')					)
              rename_keys.append(
                  (f'block{b}_project_bn/moving_variance:0', f'encoder.blocks.{hf_b}.projection.project_bn.running_var')					)

       rename_keys.append(('top_conv/kernel:0', 'encoder.top_conv.weight')					)
       rename_keys.append(('top_bn/gamma:0', 'encoder.top_bn.weight')					)
       rename_keys.append(('top_bn/beta:0', 'encoder.top_bn.bias')					)
       rename_keys.append(('top_bn/moving_mean:0', 'encoder.top_bn.running_mean')					)
       rename_keys.append(('top_bn/moving_variance:0', 'encoder.top_bn.running_var')					)

       _lowerCAmelCase    :							Any   			=       {}
       for item in rename_keys:
              if item[0] in original_param_names:
                     _lowerCAmelCase    :							Optional[Any]   			=       'efficientnet.' + item[1]

       _lowerCAmelCase    :							Tuple   			=       'classifier.weight'
       _lowerCAmelCase    :							str   			=       'classifier.bias'
       return key_mapping






def 			lowercase				(_A   ,						_A   ,						_A					):







       """simple docstring"""







       for key, value in tf_params.items():
              if "normalization" in key:
                     continue

              _lowerCAmelCase    :							Optional[Any]   			=       key_mapping[key]
              if "_conv" in key and "kernel" in key:
                     _lowerCAmelCase    :							str   			=       torch.from_numpy(_A					).permute(3   ,						2   ,						0   ,						1					)
              elif "depthwise_kernel" in key:
                     _lowerCAmelCase    :							Any   			=       torch.from_numpy(_A					).permute(2   ,						3   ,						0   ,						1					)
              elif "kernel" in key:
                     _lowerCAmelCase    :							Dict   			=       torch.from_numpy(np.transpose(_A					)					)
              else:
                     _lowerCAmelCase    :							Union[str, Any]   			=       torch.from_numpy(_A					)

              # Replace HF parameters with original TF model parameters
              assert hf_params[hf_key].shape == new_hf_value.shape
              hf_params[hf_key].copy_(_A					)






@torch.no_grad()
def 			lowercase				(_A   ,						_A   ,						_A   ,						_A					):







       """simple docstring"""







       _lowerCAmelCase    :							Tuple   			=       model_classes[model_name](
           include_top=_A   ,						weights='imagenet'   ,						input_tensor=_A   ,						input_shape=_A   ,						pooling=_A   ,						classes=1_0_0_0   ,						classifier_activation='softmax'   ,						)

       _lowerCAmelCase    :							Tuple   			=       original_model.trainable_variables
       _lowerCAmelCase    :							Dict   			=       original_model.non_trainable_variables
       _lowerCAmelCase    :							int   			=       {param.name: param.numpy() for param in tf_params}
       for param in tf_non_train_params:
              _lowerCAmelCase    :							List[str]   			=       param.numpy()
       _lowerCAmelCase    :							int   			=       list(tf_params.keys()					)

       # Load HuggingFace model
       _lowerCAmelCase    :							List[str]   			=       get_efficientnet_config(_A					)
       _lowerCAmelCase    :							str   			=       EfficientNetForImageClassification(_A					).eval()
       _lowerCAmelCase    :							Optional[Any]   			=       hf_model.state_dict()

       # Create src-to-dst parameter name mapping dictionary
       print('Converting parameters...'					)
       _lowerCAmelCase    :							int   			=       rename_keys(_A					)
       replace_params(_A   ,						_A   ,						_A					)

       # Initialize preprocessor and preprocess input image
       _lowerCAmelCase    :							List[str]   			=       convert_image_processor(_A					)
       _lowerCAmelCase    :							Any   			=       preprocessor(images=prepare_img()   ,						return_tensors='pt'					)

       # HF model inference
       hf_model.eval()
       with torch.no_grad():
              _lowerCAmelCase    :							Any   			=       hf_model(**_A					)
       _lowerCAmelCase    :							Optional[int]   			=       outputs.logits.detach().numpy()

       # Original model inference
       _lowerCAmelCase    :							Optional[int]   			=       False
       _lowerCAmelCase    :							List[str]   			=       CONFIG_MAP[model_name]['image_size']
       _lowerCAmelCase    :							Dict   			=       prepare_img().resize((image_size, image_size)   ,						resample=PIL.Image.NEAREST					)
       _lowerCAmelCase    :							int   			=       image.img_to_array(_A					)
       _lowerCAmelCase    :							Tuple   			=       np.expand_dims(_A   ,						axis=0					)
       _lowerCAmelCase    :							Optional[Any]   			=       original_model.predict(_A					)

       # Check whether original and HF model outputs match  -> np.allclose
       assert np.allclose(_A   ,						_A   ,						atol=1E-3					), "The predicted logits are not the same."
       print('Model outputs match!'					)

       if save_model:
              # Create folder to save model
              if not os.path.isdir(_A					):
                     os.mkdir(_A					)
              # Save converted model and image processor
              hf_model.save_pretrained(_A					)
              preprocessor.save_pretrained(_A					)

       if push_to_hub:
              # Push model and image processor to hub
              print(f'Pushing converted {model_name} to the hub...'					)
              _lowerCAmelCase    :							Dict   			=       f'efficientnet-{model_name}'
              preprocessor.push_to_hub(_A					)
              hf_model.push_to_hub(_A					)


if __name__ == "__main__":
       lowerCAmelCase    :     Union[str, Any]					=			argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           """--model_name""",
           default="""b0""",
           type=str,
           help="""Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].""",
       )
       parser.add_argument(
           """--pytorch_dump_folder_path""",
           default="""hf_model""",
           type=str,
           help="""Path to the output PyTorch model directory.""",
       )
       parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
       parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image processor to the hub""")

       lowerCAmelCase    :     Any					=			parser.parse_args()
       convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)





'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


lowerCAmelCase    :     int					=						{
    """configuration_trocr""": ["""TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TrOCRConfig"""],
    """processing_trocr""": ["""TrOCRProcessor"""],
}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase    :     Optional[Any]					=						[
           """TROCR_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """TrOCRForCausalLM""",
           """TrOCRPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
       from .processing_trocr import TrOCRProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
       import sys

       lowerCAmelCase    :     Any					=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

