





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

lowerCAmelCase_    =				{
    '''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_    =				{
    '''b0''': {
        '''hidden_dim''': 1_2_8_0,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.0,
        '''image_size''': 2_2_4,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [],
    },
    '''b1''': {
        '''hidden_dim''': 1_2_8_0,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.1,
        '''image_size''': 2_4_0,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [1_6],
    },
    '''b2''': {
        '''hidden_dim''': 1_4_0_8,
        '''width_coef''': 1.1,
        '''depth_coef''': 1.2,
        '''image_size''': 2_6_0,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 8, 1_6],
    },
    '''b3''': {
        '''hidden_dim''': 1_5_3_6,
        '''width_coef''': 1.2,
        '''depth_coef''': 1.4,
        '''image_size''': 3_0_0,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 1_8],
    },
    '''b4''': {
        '''hidden_dim''': 1_7_9_2,
        '''width_coef''': 1.4,
        '''depth_coef''': 1.8,
        '''image_size''': 3_8_0,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [6],
    },
    '''b5''': {
        '''hidden_dim''': 2_0_4_8,
        '''width_coef''': 1.6,
        '''depth_coef''': 2.2,
        '''image_size''': 4_5_6,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [1_3, 2_7],
    },
    '''b6''': {
        '''hidden_dim''': 2_3_0_4,
        '''width_coef''': 1.8,
        '''depth_coef''': 2.6,
        '''image_size''': 5_2_8,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [3_1],
    },
    '''b7''': {
        '''hidden_dim''': 2_5_6_0,
        '''width_coef''': 2.0,
        '''depth_coef''': 3.1,
        '''image_size''': 6_0_0,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [1_8],
    },
}






def      lowerCamelCase_						(				_UpperCamelCase )		->						List[Any]:






				"""simple docstring"""


				snake_case_						:    List[str]				=  EfficientNetConfig()
				snake_case_						:    Optional[int]				=  CONFIG_MAP[model_name]['''hidden_dim''']
				snake_case_						:    Dict				=  CONFIG_MAP[model_name]['''width_coef''']
				snake_case_						:    Optional[Any]				=  CONFIG_MAP[model_name]['''depth_coef''']
				snake_case_						:    Dict				=  CONFIG_MAP[model_name]['''image_size''']
				snake_case_						:    Union[str, Any]				=  CONFIG_MAP[model_name]['''dropout_rate''']
				snake_case_						:    Union[str, Any]				=  CONFIG_MAP[model_name]['''dw_padding''']

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

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






def      lowerCamelCase_						(				)		->						Union[str, Any]:






				"""simple docstring"""


				snake_case_						:    Optional[Any]				=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
				snake_case_						:    int				=  Image.open(requests.get(_UpperCamelCase ,				stream=_UpperCamelCase ).raw )
				return im






def      lowerCamelCase_						(				_UpperCamelCase )		->						str:






				"""simple docstring"""


				snake_case_						:    Any				=  CONFIG_MAP[model_name]['''image_size''']
				snake_case_						:    int				=  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=_UpperCamelCase ,				)
				return preprocessor






def      lowerCamelCase_						(				_UpperCamelCase )		->						str:






				"""simple docstring"""


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

				snake_case_						:    str				=  []
				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:
								snake_case_						:    Union[str, 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''') )

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

				snake_case_						:    str				=  '''classifier.weight'''
				snake_case_						:    Any				=  '''classifier.bias'''
				return key_mapping






def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase )		->						Optional[Any]:






				"""simple docstring"""


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

								snake_case_						:    Optional[Any]				=  key_mapping[key]
								if "_conv" in key and "kernel" in key:
												snake_case_						:    Any				=  torch.from_numpy(_UpperCamelCase ).permute(3 ,				2 ,				0 ,				1 )
								elif "depthwise_kernel" in key:
												snake_case_						:    List[Any]				=  torch.from_numpy(_UpperCamelCase ).permute(2 ,				3 ,				0 ,				1 )
								elif "kernel" in key:
												snake_case_						:    Tuple				=  torch.from_numpy(np.transpose(_UpperCamelCase ) )
								else:
												snake_case_						:    str				=  torch.from_numpy(_UpperCamelCase )

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






@torch.no_grad()
def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase )		->						List[str]:






				"""simple docstring"""


				snake_case_						:    Dict				=  model_classes[model_name](
				    include_top=_UpperCamelCase ,				weights='''imagenet''' ,				input_tensor=_UpperCamelCase ,				input_shape=_UpperCamelCase ,				pooling=_UpperCamelCase ,				classes=1_000 ,				classifier_activation='''softmax''' ,				)

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

				# Load HuggingFace model
				snake_case_						:    Tuple				=  get_efficientnet_config(_UpperCamelCase )
				snake_case_						:    Tuple				=  EfficientNetForImageClassification(_UpperCamelCase ).eval()
				snake_case_						:    Dict				=  hf_model.state_dict()

				# Create src-to-dst parameter name mapping dictionary
				print('''Converting parameters...''' )
				snake_case_						:    Dict				=  rename_keys(_UpperCamelCase )
				replace_params(_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase )

				# Initialize preprocessor and preprocess input image
				snake_case_						:    Any				=  convert_image_processor(_UpperCamelCase )
				snake_case_						:    Any				=  preprocessor(images=prepare_img() ,				return_tensors='''pt''' )

				# HF model inference
				hf_model.eval()
				with torch.no_grad():
								snake_case_						:    Any				=  hf_model(**_UpperCamelCase )
				snake_case_						:    int				=  outputs.logits.detach().numpy()

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

				# Check whether original and HF model outputs match  -> np.allclose
				assert np.allclose(_UpperCamelCase ,				_UpperCamelCase ,				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(_UpperCamelCase ):
												os.mkdir(_UpperCamelCase )
								# Save converted model and image processor
								hf_model.save_pretrained(_UpperCamelCase )
								preprocessor.save_pretrained(_UpperCamelCase )

				if push_to_hub:
								# Push model and image processor to hub
								print(f'''Pushing converted {model_name} to the hub...''' )
								snake_case_						:    Optional[Any]				=  f'''efficientnet-{model_name}'''
								preprocessor.push_to_hub(_UpperCamelCase )
								hf_model.push_to_hub(_UpperCamelCase )


if __name__ == "__main__":
	lowerCAmelCase_    =				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_    =				parser.parse_args()
	convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)







def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase )		->						str:






				"""simple docstring"""


				return "\n".join(
				    f'''{number} * {i} = {number * i}''' for i in range(1 ,				number_of_terms + 1 ) )


if __name__ == "__main__":
	print(multiplication_table(number=5, number_of_terms=1_0))
