

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

lowerCAmelCase__    :							Tuple					    =    {
    """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__    :							List[str]					    =    {
    """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							(					__lowerCAmelCase		:   Tuple  ):

		"""simple docstring"""

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

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

		snake_case__    :							str							    =	idalabel
		snake_case__    :							Tuple							    =	{v: k for k, v in idalabel.items()}
		return config




def 						_a							(					):

		"""simple docstring"""

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




def 						_a							(					__lowerCAmelCase		:   Any  ):

		"""simple docstring"""

		snake_case__    :							int							    =	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.47853944, 0.4732864, 0.47434163] ,				do_center_crop=__a ,				)
		return preprocessor




def 						_a							(					__lowerCAmelCase		:   List[str]  ):

		"""simple docstring"""

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

		snake_case__    :							Union[str, 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:
				snake_case__    :							List[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__    :							Dict							    =	{}
		for item in rename_keys:
				if item[0] in original_param_names:
						snake_case__    :							List[str]							    =	'efficientnet.' + item[1]

		snake_case__    :							List[Any]							    =	'classifier.weight'
		snake_case__    :							Tuple							    =	'classifier.bias'
		return key_mapping




def 						_a							(					__lowerCAmelCase		:   Any ,				__lowerCAmelCase		:   Optional[int] ,				__lowerCAmelCase		:   Tuple  ):

		"""simple docstring"""

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

				snake_case__    :							Any							    =	key_mapping[key]
				if "_conv" in key and "kernel" in key:
						snake_case__    :							Optional[int]							    =	torch.from_numpy(__a  ).permute(3 ,				2 ,				0 ,				1  )
				elif "depthwise_kernel" in key:
						snake_case__    :							Optional[Any]							    =	torch.from_numpy(__a  ).permute(2 ,				3 ,				0 ,				1  )
				elif "kernel" in key:
						snake_case__    :							Dict							    =	torch.from_numpy(np.transpose(__a  )  )
				else:
						snake_case__    :							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							(					__lowerCAmelCase		:   Optional[int] ,				__lowerCAmelCase		:   Any ,				__lowerCAmelCase		:   Dict ,				__lowerCAmelCase		:   Optional[Any]  ):

		"""simple docstring"""

		snake_case__    :							Optional[int]							    =	model_classes[model_name](
		    include_top=__a ,				weights='''imagenet''' ,				input_tensor=__a ,				input_shape=__a ,				pooling=__a ,				classes=10_00 ,				classifier_activation='''softmax''' ,				)

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

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

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

		# Initialize preprocessor and preprocess input image
		snake_case__    :							Tuple							    =	convert_image_processor(__a  )
		snake_case__    :							int							    =	preprocessor(images=prepare_img() ,				return_tensors='''pt'''  )

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

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


if __name__ == "__main__":
					lowerCAmelCase__    :							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__    :							Optional[Any]					    =    parser.parse_args()
					convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)








"""simple docstring"""







def     UpperCAmelCase_					(				__a   :			int						):




			'''simple docstring'''



			_lowerCamelCase					:					Optional[Any]	       =    int(__a						)
			if decimal in (0, 1):  # Exit cases for the recursion
						return str(__a						)
			_lowerCamelCase       ,							_lowerCamelCase					:					Union[str, Any]	       =    divmod(__a      ,							2						)
			return binary_recursive(__a						) + str(__a						)

def     UpperCAmelCase_					(				__a   :			str						):




			'''simple docstring'''



			_lowerCamelCase					:					int	       =    str(__a						).strip()
			if not number:
						raise ValueError('No input value was provided'						)
			_lowerCamelCase					:					Tuple	       =    '-' if number.startswith('-'						) else ''
			_lowerCamelCase					:					List[Any]	       =    number.lstrip('-'						)
			if not number.isnumeric():
						raise ValueError('Input value is not an integer'						)
			return f"{negative}0b{binary_recursive(int(__a						)						)}"


if __name__ == "__main__":
	from doctest import testmod

	testmod()





