




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()
SCREAMING_SNAKE_CASE__       :      List[str]  =   logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__       :      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,
}

SCREAMING_SNAKE_CASE__       :      Any  =   {
    '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      (      _SCREAMING_SNAKE_CASE    )	-> str:
       lowerCamelCase   :      int			 =					EfficientNetConfig()
       lowerCamelCase   :      List[str]			 =					CONFIG_MAP[model_name]["hidden_dim"]
       lowerCamelCase   :      List[str]			 =					CONFIG_MAP[model_name]["width_coef"]
       lowerCamelCase   :      Any			 =					CONFIG_MAP[model_name]["depth_coef"]
       lowerCamelCase   :      Union[str, Any]			 =					CONFIG_MAP[model_name]["image_size"]
       lowerCamelCase   :      Optional[int]			 =					CONFIG_MAP[model_name]["dropout_rate"]
       lowerCamelCase   :      str			 =					CONFIG_MAP[model_name]["dw_padding"]

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

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

def      A      (      )	-> int:
       lowerCamelCase   :      str			 =					"http://images.cocodataset.org/val2017/000000039769.jpg"
       lowerCamelCase   :      Tuple			 =					Image.open(requests.get(_SCREAMING_SNAKE_CASE ,stream=_SCREAMING_SNAKE_CASE    ).raw    )
       return im

def      A      (      _SCREAMING_SNAKE_CASE    )	-> str:
       lowerCamelCase   :      List[Any]			 =					CONFIG_MAP[model_name]["image_size"]
       lowerCamelCase   :      str			 =					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=_SCREAMING_SNAKE_CASE ,)
       return preprocessor

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

       lowerCamelCase   :      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:
              lowerCamelCase   :      Dict			 =					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   :      Optional[int]			 =					{}
       for item in rename_keys:
              if item[0] in original_param_names:
                     lowerCamelCase   :      List[str]			 =					"efficientnet." + item[1]

       lowerCamelCase   :      int			 =					"classifier.weight"
       lowerCamelCase   :      Union[str, Any]			 =					"classifier.bias"
       return key_mapping

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> Dict:
       for key, value in tf_params.items():
              if "normalization" in key:
                     continue

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

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

@torch.no_grad()
def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> Optional[int]:
       lowerCamelCase   :      Optional[int]			 =					model_classes[model_name](
           include_top=_SCREAMING_SNAKE_CASE ,weights="imagenet" ,input_tensor=_SCREAMING_SNAKE_CASE ,input_shape=_SCREAMING_SNAKE_CASE ,pooling=_SCREAMING_SNAKE_CASE ,classes=1000 ,classifier_activation="softmax" ,)

       lowerCamelCase   :      List[Any]			 =					original_model.trainable_variables
       lowerCamelCase   :      Tuple			 =					original_model.non_trainable_variables
       lowerCamelCase   :      Union[str, Any]			 =					{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   :      Union[str, Any]			 =					get_efficientnet_config(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Optional[int]			 =					EfficientNetForImageClassification(_SCREAMING_SNAKE_CASE    ).eval()
       lowerCamelCase   :      Tuple			 =					hf_model.state_dict()

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

       # Initialize preprocessor and preprocess input image
       lowerCamelCase   :      int			 =					convert_image_processor(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      int			 =					preprocessor(images=prepare_img() ,return_tensors="pt"    )

       # HF model inference
       hf_model.eval()
       with torch.no_grad():
              lowerCamelCase   :      Optional[Any]			 =					hf_model(**_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      str			 =					outputs.logits.detach().numpy()

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

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

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


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

       SCREAMING_SNAKE_CASE__       :      Tuple  =   parser.parse_args()
       convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)














def 						_a    (      a :float	,							a :float )							->			float:
 return price * (1 + tax_rate)


if __name__ == "__main__":
     print(f"""{price_plus_tax(100, 0.25) = }""")
     print(f"""{price_plus_tax(125.50, 0.05) = }""")


