

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

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

_UpperCamelCase 			=		{
    "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 		_lowercase (     lowercase__						):
      __lowerCAmelCase					:						Any					=					EfficientNetConfig()
      __lowerCAmelCase					:						Tuple					=					CONFIG_MAP[model_name]['''hidden_dim''']
      __lowerCAmelCase					:						Optional[Any]					=					CONFIG_MAP[model_name]['''width_coef''']
      __lowerCAmelCase					:						Optional[int]					=					CONFIG_MAP[model_name]['''depth_coef''']
      __lowerCAmelCase					:						Optional[Any]					=					CONFIG_MAP[model_name]['''image_size''']
      __lowerCAmelCase					:						Optional[Any]					=					CONFIG_MAP[model_name]['''dropout_rate''']
      __lowerCAmelCase					:						Union[str, Any]					=					CONFIG_MAP[model_name]['''dw_padding''']

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

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







def 		_lowercase (     ):
      __lowerCAmelCase					:						str					=					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
      __lowerCAmelCase					:						List[str]					=					Image.open(requests.get(lowercase__   ,	stream=lowercase__						).raw						)
      return im







def 		_lowercase (     lowercase__						):
      __lowerCAmelCase					:						Dict					=					CONFIG_MAP[model_name]['''image_size''']
      __lowerCAmelCase					:						Tuple					=					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=lowercase__   ,	)
      return preprocessor







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

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

      __lowerCAmelCase					:						Optional[int]					=					'''classifier.weight'''
      __lowerCAmelCase					:						Optional[Any]					=					'''classifier.bias'''
      return key_mapping







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

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

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







@torch.no_grad()
def 		_lowercase (     lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__						):
      __lowerCAmelCase					:						List[Any]					=					model_classes[model_name](
          include_top=lowercase__   ,	weights='''imagenet'''   ,	input_tensor=lowercase__   ,	input_shape=lowercase__   ,	pooling=lowercase__   ,	classes=1_0_0_0   ,	classifier_activation='''softmax'''   ,	)

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

      # Load HuggingFace model
      __lowerCAmelCase					:						Dict					=					get_efficientnet_config(lowercase__						)
      __lowerCAmelCase					:						Union[str, Any]					=					EfficientNetForImageClassification(lowercase__						).eval()
      __lowerCAmelCase					:						Tuple					=					hf_model.state_dict()

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

      # Initialize preprocessor and preprocess input image
      __lowerCAmelCase					:						List[str]					=					convert_image_processor(lowercase__						)
      __lowerCAmelCase					:						Union[str, Any]					=					preprocessor(images=prepare_img()   ,	return_tensors='''pt'''						)

      # HF model inference
      hf_model.eval()
      with torch.no_grad():
            __lowerCAmelCase					:						List[Any]					=					hf_model(**lowercase__						)
      __lowerCAmelCase					:						Tuple					=					outputs.logits.detach().numpy()

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

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

      if push_to_hub:
            # Push model and image processor to hub
            print(f"""Pushing converted {model_name} to the hub..."""						)
            __lowerCAmelCase					:						Tuple					=					f"""efficientnet-{model_name}"""
            preprocessor.push_to_hub(lowercase__						)
            hf_model.push_to_hub(lowercase__						)


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

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





def 		_lowercase (     lowercase__   ,	lowercase__						):
      if a < 0 or b < 0:
            raise ValueError('''the value of both inputs must be positive'''						)
      __lowerCAmelCase					:						int					=					str(bin(lowercase__						)						)[2:]  # remove the leading "0b"
      __lowerCAmelCase					:						Any					=					str(bin(lowercase__						)						)[2:]
      __lowerCAmelCase					:						List[str]					=					max(len(lowercase__						)   ,	len(lowercase__						)						)
      return "0b" + "".join(
          str(int('''1''' in (char_a, char_b)						)						)
          for char_a, char_b in zip(a_binary.zfill(lowercase__						)   ,	b_binary.zfill(lowercase__						)						)						)


if __name__ == "__main__":
 import doctest

 doctest.testmod()


