
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_						: 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_						: Any         =  {
    """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        A__					(						lowerCamelCase						)  ->	Optional[Any]:
				UpperCamelCase_:							Dict	=						EfficientNetConfig()
				UpperCamelCase_:							Optional[Any]	=						CONFIG_MAP[model_name]["""hidden_dim"""]
				UpperCamelCase_:							Optional[int]	=						CONFIG_MAP[model_name]["""width_coef"""]
				UpperCamelCase_:							str	=						CONFIG_MAP[model_name]["""depth_coef"""]
				UpperCamelCase_:							int	=						CONFIG_MAP[model_name]["""image_size"""]
				UpperCamelCase_:							Optional[int]	=						CONFIG_MAP[model_name]["""dropout_rate"""]
				UpperCamelCase_:							Optional[Any]	=						CONFIG_MAP[model_name]["""dw_padding"""]

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

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







def        A__					(						)  ->	Optional[Any]:
				UpperCamelCase_:							Optional[Any]	=						"""http://images.cocodataset.org/val2017/000000039769.jpg"""
				UpperCamelCase_:							Any	=						Image.open(requests.get(lowerCamelCase     ,    stream=lowerCamelCase						).raw						)
				return im







def        A__					(						lowerCamelCase						)  ->	Tuple:
				UpperCamelCase_:							Dict	=						CONFIG_MAP[model_name]["""image_size"""]
				UpperCamelCase_:							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=lowerCamelCase     ,    )
				return preprocessor







def        A__					(						lowerCamelCase						)  ->	Optional[Any]:
				UpperCamelCase_:							Tuple	=						[v.split("""_"""						)[0].split("""block"""						)[1] for v in original_param_names if v.startswith("""block"""						)]
				UpperCamelCase_:							List[str]	=						sorted(set(lowerCamelCase						)						)
				UpperCamelCase_:							Optional[Any]	=						len(lowerCamelCase						)
				UpperCamelCase_:							Any	=						{b: str(lowerCamelCase						) for b, i in zip(lowerCamelCase     ,    range(lowerCamelCase						)						)}

				UpperCamelCase_:							List[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:
								UpperCamelCase_:							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""")						)

				UpperCamelCase_:							int	=						{}
				for item in rename_keys:
								if item[0] in original_param_names:
												UpperCamelCase_:							Optional[Any]	=						"""efficientnet.""" + item[1]

				UpperCamelCase_:							List[str]	=						"""classifier.weight"""
				UpperCamelCase_:							Dict	=						"""classifier.bias"""
				return key_mapping







def        A__					(						lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase						)  ->	List[Any]:
				for key, value in tf_params.items():
								if "normalization" in key:
												continue

								UpperCamelCase_:							Dict	=						key_mapping[key]
								if "_conv" in key and "kernel" in key:
												UpperCamelCase_:							List[Any]	=						torch.from_numpy(lowerCamelCase						).permute(3     ,    2     ,    0     ,    1						)
								elif "depthwise_kernel" in key:
												UpperCamelCase_:							int	=						torch.from_numpy(lowerCamelCase						).permute(2     ,    3     ,    0     ,    1						)
								elif "kernel" in key:
												UpperCamelCase_:							str	=						torch.from_numpy(np.transpose(lowerCamelCase						)						)
								else:
												UpperCamelCase_:							int	=						torch.from_numpy(lowerCamelCase						)

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







@torch.no_grad()
def        A__					(						lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase						)  ->	Any:
				UpperCamelCase_:							Optional[int]	=						model_classes[model_name](
				    include_top=lowerCamelCase     ,    weights="""imagenet"""     ,    input_tensor=lowerCamelCase     ,    input_shape=lowerCamelCase     ,    pooling=lowerCamelCase     ,    classes=10_00     ,    classifier_activation="""softmax"""     ,    )

				UpperCamelCase_:							Tuple	=						original_model.trainable_variables
				UpperCamelCase_:							str	=						original_model.non_trainable_variables
				UpperCamelCase_:							List[Any]	=						{param.name: param.numpy() for param in tf_params}
				for param in tf_non_train_params:
								UpperCamelCase_:							Any	=						param.numpy()
				UpperCamelCase_:							Union[str, Any]	=						list(tf_params.keys()						)

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

				# Create src-to-dst parameter name mapping dictionary
				print("""Converting parameters..."""						)
				UpperCamelCase_:							List[Any]	=						rename_keys(lowerCamelCase						)
				replace_params(lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase						)

				# Initialize preprocessor and preprocess input image
				UpperCamelCase_:							List[str]	=						convert_image_processor(lowerCamelCase						)
				UpperCamelCase_:							Dict	=						preprocessor(images=prepare_img()     ,    return_tensors="""pt"""						)

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

				# Original model inference
				UpperCamelCase_:							Any	=						False
				UpperCamelCase_:							Union[str, Any]	=						CONFIG_MAP[model_name]["""image_size"""]
				UpperCamelCase_:							Tuple	=						prepare_img().resize((image_size, image_size)     ,    resample=PIL.Image.NEAREST						)
				UpperCamelCase_:							List[Any]	=						image.img_to_array(lowerCamelCase						)
				UpperCamelCase_:							int	=						np.expand_dims(lowerCamelCase     ,    axis=0						)
				UpperCamelCase_:							List[Any]	=						original_model.predict(lowerCamelCase						)

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

				if push_to_hub:
								# Push model and image processor to hub
								print(F'''Pushing converted {model_name} to the hub...'''						)
								UpperCamelCase_:							int	=						F'''efficientnet-{model_name}'''
								preprocessor.push_to_hub(lowerCamelCase						)
								hf_model.push_to_hub(lowerCamelCase						)


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








def        A__					(						lowerCamelCase = 50						)  ->	int:
				UpperCamelCase_:							List[Any]	=						[[0] * 3 for _ in range(length + 1						)]

				for row_length in range(length + 1						):
								for tile_length in range(2     ,    5						):
												for tile_start in range(row_length - tile_length + 1						):
																different_colour_ways_number[row_length][tile_length - 2] += (
																    different_colour_ways_number[row_length - tile_start - tile_length][
																        tile_length - 2
																    ]
																    + 1
																)

				return sum(different_colour_ways_number[length]						)


if __name__ == "__main__":
							print(F"""{solution() = }""")






