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    :				int							     =	{
    '''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    :				Union[str, 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 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: List[Any]			)							->      Optional[int]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   EfficientNetConfig()
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""hidden_dim"""]
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""width_coef"""]
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""depth_coef"""]
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""image_size"""]
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""dropout_rate"""]
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""dw_padding"""]

  SCREAMING_SNAKE_CASE__    							=   """huggingface/label-files"""
  SCREAMING_SNAKE_CASE__    							=   """imagenet-1k-id2label.json"""
  SCREAMING_SNAKE_CASE__    							=   10_00
  SCREAMING_SNAKE_CASE__    							=   json.load(open(hf_hub_download(__UpperCamelCase				,  __UpperCamelCase				,  repo_type="""dataset"""			)				,  """r"""			)			)
  SCREAMING_SNAKE_CASE__    							=   {int(__UpperCamelCase			): v for k, v in idalabel.items()}

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



def 			__SCREAMING_SNAKE_CASE						(      )							->      Any:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   """http://images.cocodataset.org/val2017/000000039769.jpg"""
  SCREAMING_SNAKE_CASE__    							=   Image.open(requests.get(__UpperCamelCase				,  stream=__UpperCamelCase			).raw			)
  return im



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: int			)							->      str:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""image_size"""]
  SCREAMING_SNAKE_CASE__    							=   EfficientNetImageProcessor(
      size={"""height""": size, """width""": size}				,  image_mean=[0.485, 0.456, 0.406]				,  image_std=[0.4785_3944, 0.473_2864, 0.4743_4163]				,  do_center_crop=__UpperCamelCase				,  )
  return preprocessor



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Tuple			)							->      Any:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   [v.split("""_"""			)[0].split("""block"""			)[1] for v in original_param_names if v.startswith("""block"""			)]
  SCREAMING_SNAKE_CASE__    							=   sorted(set(__UpperCamelCase			)			)
  SCREAMING_SNAKE_CASE__    							=   len(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   {b: str(__UpperCamelCase			) for b, i in zip(__UpperCamelCase				,  range(__UpperCamelCase			)			)}

  SCREAMING_SNAKE_CASE__    							=   []
  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:
    SCREAMING_SNAKE_CASE__    							=   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""")			)

  SCREAMING_SNAKE_CASE__    							=   {}
  for item in rename_keys:
    if item[0] in original_param_names:
      SCREAMING_SNAKE_CASE__    							=   """efficientnet.""" + item[1]

  SCREAMING_SNAKE_CASE__    							=   """classifier.weight"""
  SCREAMING_SNAKE_CASE__    							=   """classifier.bias"""
  return key_mapping



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Union[str, Any]				,  __UpperCamelCase					: Dict				,  __UpperCamelCase					: List[str]			)							->      Dict:







  """simple docstring"""





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

    SCREAMING_SNAKE_CASE__    							=   key_mapping[key]
    if "_conv" in key and "kernel" in key:
      SCREAMING_SNAKE_CASE__    							=   torch.from_numpy(__UpperCamelCase			).permute(3				,  2				,  0				,  1			)
    elif "depthwise_kernel" in key:
      SCREAMING_SNAKE_CASE__    							=   torch.from_numpy(__UpperCamelCase			).permute(2				,  3				,  0				,  1			)
    elif "kernel" in key:
      SCREAMING_SNAKE_CASE__    							=   torch.from_numpy(np.transpose(__UpperCamelCase			)			)
    else:
      SCREAMING_SNAKE_CASE__    							=   torch.from_numpy(__UpperCamelCase			)

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



@torch.no_grad()
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Tuple				,  __UpperCamelCase					: Any				,  __UpperCamelCase					: Tuple				,  __UpperCamelCase					: Tuple			)							->      List[Any]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   model_classes[model_name](
      include_top=__UpperCamelCase				,  weights="""imagenet"""				,  input_tensor=__UpperCamelCase				,  input_shape=__UpperCamelCase				,  pooling=__UpperCamelCase				,  classes=10_00				,  classifier_activation="""softmax"""				,  )

  SCREAMING_SNAKE_CASE__    							=   original_model.trainable_variables
  SCREAMING_SNAKE_CASE__    							=   original_model.non_trainable_variables
  SCREAMING_SNAKE_CASE__    							=   {param.name: param.numpy() for param in tf_params}
  for param in tf_non_train_params:
    SCREAMING_SNAKE_CASE__    							=   param.numpy()
  SCREAMING_SNAKE_CASE__    							=   list(tf_params.keys()			)

  # Load HuggingFace model
  SCREAMING_SNAKE_CASE__    							=   get_efficientnet_config(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   EfficientNetForImageClassification(__UpperCamelCase			).eval()
  SCREAMING_SNAKE_CASE__    							=   hf_model.state_dict()

  # Create src-to-dst parameter name mapping dictionary
  print("""Converting parameters..."""			)
  SCREAMING_SNAKE_CASE__    							=   rename_keys(__UpperCamelCase			)
  replace_params(__UpperCamelCase				,  __UpperCamelCase				,  __UpperCamelCase			)

  # Initialize preprocessor and preprocess input image
  SCREAMING_SNAKE_CASE__    							=   convert_image_processor(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   preprocessor(images=prepare_img()				,  return_tensors="""pt"""			)

  # HF model inference
  hf_model.eval()
  with torch.no_grad():
    SCREAMING_SNAKE_CASE__    							=   hf_model(**__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   outputs.logits.detach().numpy()

  # Original model inference
  SCREAMING_SNAKE_CASE__    							=   False
  SCREAMING_SNAKE_CASE__    							=   CONFIG_MAP[model_name]["""image_size"""]
  SCREAMING_SNAKE_CASE__    							=   prepare_img().resize((image_size, image_size)				,  resample=PIL.Image.NEAREST			)
  SCREAMING_SNAKE_CASE__    							=   image.img_to_array(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   np.expand_dims(__UpperCamelCase				,  axis=0			)
  SCREAMING_SNAKE_CASE__    							=   original_model.predict(__UpperCamelCase			)

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

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


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

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

# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team.
#
# 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.

# this script dumps information about the environment

import os
import sys

import transformers


__lowerCamelCase    :				Tuple							     =	'''3'''

print('''Python version:''', sys.version)
print('''transformers version:''', transformers.__version__)

try:
   import torch

   print('''Torch version:''', torch.__version__)
   print('''Cuda available:''', torch.cuda.is_available())
   print('''Cuda version:''', torch.version.cuda)
   print('''CuDNN version:''', torch.backends.cudnn.version())
   print('''Number of GPUs available:''', torch.cuda.device_count())
   print('''NCCL version:''', torch.cuda.nccl.version())
except ImportError:
   print('''Torch version:''', None)

try:
   import deepspeed

   print('''DeepSpeed version:''', deepspeed.__version__)
except ImportError:
   print('''DeepSpeed version:''', None)

try:
   import tensorflow as tf

   print('''TensorFlow version:''', tf.__version__)
   print('''TF GPUs available:''', bool(tf.config.list_physical_devices('''GPU''')))
   print('''Number of TF GPUs available:''', len(tf.config.list_physical_devices('''GPU''')))
except ImportError:
   print('''TensorFlow version:''', None)
