


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

__snake_case											=					{
    '''b0''': efficientnet.EfficientNetBa,
    '''b1''': efficientnet.EfficientNetBa,
    '''b2''': efficientnet.EfficientNetBa,
    '''b3''': efficientnet.EfficientNetBa,
    '''b4''': efficientnet.EfficientNetBa,
    '''b5''': efficientnet.EfficientNetBa,
    '''b6''': efficientnet.EfficientNetBa,
    '''b7''': efficientnet.EfficientNetBa,
}

__snake_case											=					{
    '''b0''': {
        '''hidden_dim''': 1280,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.0,
        '''image_size''': 224,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [],
    },
    '''b1''': {
        '''hidden_dim''': 1280,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.1,
        '''image_size''': 240,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [16],
    },
    '''b2''': {
        '''hidden_dim''': 1408,
        '''width_coef''': 1.1,
        '''depth_coef''': 1.2,
        '''image_size''': 260,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 8, 16],
    },
    '''b3''': {
        '''hidden_dim''': 1536,
        '''width_coef''': 1.2,
        '''depth_coef''': 1.4,
        '''image_size''': 300,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 18],
    },
    '''b4''': {
        '''hidden_dim''': 1792,
        '''width_coef''': 1.4,
        '''depth_coef''': 1.8,
        '''image_size''': 380,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [6],
    },
    '''b5''': {
        '''hidden_dim''': 2048,
        '''width_coef''': 1.6,
        '''depth_coef''': 2.2,
        '''image_size''': 456,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [13, 27],
    },
    '''b6''': {
        '''hidden_dim''': 2304,
        '''width_coef''': 1.8,
        '''depth_coef''': 2.6,
        '''image_size''': 528,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [31],
    },
    '''b7''': {
        '''hidden_dim''': 2560,
        '''width_coef''': 2.0,
        '''depth_coef''': 3.1,
        '''image_size''': 600,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [18],
    },
}


def       a				(      __a       )  ->     Tuple:



					'''simple docstring'''






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

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

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


def       a				(      )  ->     Optional[Any]:



					'''simple docstring'''






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


def       a				(      __a       )  ->     int:



					'''simple docstring'''






					UpperCamelCase__       :Optional[int]							 =			CONFIG_MAP[model_name]['''image_size''']
					UpperCamelCase__       :Any							 =			EfficientNetImageProcessor(
					    size={'''height''': size, '''width''': size}			, image_mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]			, image_std=[0.4_7_8_5_3_9_4_4, 0.4_7_3_2_8_6_4, 0.4_7_4_3_4_1_6_3]			, do_center_crop=__a			, )
					return preprocessor


def       a				(      __a       )  ->     int:



					'''simple docstring'''






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

					UpperCamelCase__       :List[Any]							 =			[]
					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:
										UpperCamelCase__       :Optional[int]							 =			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''')       )

					UpperCamelCase__       :List[str]							 =			{}
					for item in rename_keys:
										if item[0] in original_param_names:
															UpperCamelCase__       :Optional[int]							 =			'''efficientnet.''' + item[1]

					UpperCamelCase__       :List[str]							 =			'''classifier.weight'''
					UpperCamelCase__       :Optional[int]							 =			'''classifier.bias'''
					return key_mapping


def       a				(      __a			, __a			, __a       )  ->     List[str]:



					'''simple docstring'''






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

										UpperCamelCase__       :Optional[Any]							 =			key_mapping[key]
										if "_conv" in key and "kernel" in key:
															UpperCamelCase__       :List[str]							 =			torch.from_numpy(__a       ).permute(3			, 2			, 0			, 1       )
										elif "depthwise_kernel" in key:
															UpperCamelCase__       :List[Any]							 =			torch.from_numpy(__a       ).permute(2			, 3			, 0			, 1       )
										elif "kernel" in key:
															UpperCamelCase__       :str							 =			torch.from_numpy(np.transpose(__a       )       )
										else:
															UpperCamelCase__       :Tuple							 =			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       a				(      __a			, __a			, __a			, __a       )  ->     Union[str, Any]:



					'''simple docstring'''






					UpperCamelCase__       :List[Any]							 =			model_classes[model_name](
					    include_top=__a			, weights='''imagenet'''			, input_tensor=__a			, input_shape=__a			, pooling=__a			, classes=1000			, classifier_activation='''softmax'''			, )

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

					# Load HuggingFace model
					UpperCamelCase__       :Optional[Any]							 =			get_efficientnet_config(__a       )
					UpperCamelCase__       :Optional[int]							 =			EfficientNetForImageClassification(__a       ).eval()
					UpperCamelCase__       :Optional[Any]							 =			hf_model.state_dict()

					# Create src-to-dst parameter name mapping dictionary
					print('''Converting parameters...'''       )
					UpperCamelCase__       :Optional[Any]							 =			rename_keys(__a       )
					replace_params(__a			, __a			, __a       )

					# Initialize preprocessor and preprocess input image
					UpperCamelCase__       :Dict							 =			convert_image_processor(__a       )
					UpperCamelCase__       :Optional[int]							 =			preprocessor(images=prepare_img()			, return_tensors='''pt'''       )

					# HF model inference
					hf_model.eval()
					with torch.no_grad():
										UpperCamelCase__       :List[str]							 =			hf_model(**__a       )
					UpperCamelCase__       :Dict							 =			outputs.logits.detach().numpy()

					# Original model inference
					UpperCamelCase__       :int							 =			False
					UpperCamelCase__       :str							 =			CONFIG_MAP[model_name]['''image_size''']
					UpperCamelCase__       :Optional[Any]							 =			prepare_img().resize((image_size, image_size)			, resample=PIL.Image.NEAREST       )
					UpperCamelCase__       :str							 =			image.img_to_array(__a       )
					UpperCamelCase__       :Union[str, Any]							 =			np.expand_dims(__a			, axis=0       )
					UpperCamelCase__       :str							 =			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...'''       )
										UpperCamelCase__       :Tuple							 =			f'''efficientnet-{model_name}'''
										preprocessor.push_to_hub(__a       )
										hf_model.push_to_hub(__a       )


if __name__ == "__main__":
		__snake_case											=					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''')

		__snake_case											=					parser.parse_args()
		convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)



'''simple docstring'''
# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser


def       a				(      __a=None       )  ->     List[str]:



					'''simple docstring'''






					UpperCamelCase__       :Optional[Any]							 =			argparse.ArgumentParser(add_help=__a			, allow_abbrev=__a       )
					# The main config parser
					UpperCamelCase__       :str							 =			config_command_parser(__a       )
					# The subparser to add commands to
					UpperCamelCase__       :Union[str, Any]							 =			config_parser.add_subparsers(title='''subcommands'''			, dest='''subcommand'''       )

					# Then add other parsers with the parent parser
					default_command_parser(__a			, parents=[parent_parser]       )
					update_command_parser(__a			, parents=[parent_parser]       )

					return config_parser


def       a				(      )  ->     Any:



					'''simple docstring'''






					UpperCamelCase__       :int							 =			get_config_parser()
					UpperCamelCase__       :List[Any]							 =			config_parser.parse_args()

					if not hasattr(__a			, '''func'''       ):
										config_parser.print_help()
										exit(1       )

					# Run
					args.func(__a       )


if __name__ == "__main__":
		main()