



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

lowerCamelCase	 =       {
    """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	 =       {
    """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					):



							"""simple docstring"""



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

							__lowercase       	="huggingface/label-files"
							__lowercase       	="imagenet-1k-id2label.json"
							__lowercase       	=1_000
							__lowercase       	=json.load(open(hf_hub_download(_lowerCamelCase  ,   _lowerCamelCase  ,   repo_type='dataset'					)  ,   'r'					)					)
							__lowercase       	={int(_lowerCamelCase					): v for k, v in idalabel.items()}

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



def   _A       (   ):



							"""simple docstring"""



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



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=CONFIG_MAP[model_name]["image_size"]
							__lowercase       	=EfficientNetImageProcessor(
							    size={'height': size, 'width': size}  ,   image_mean=[0.4_85, 0.4_56, 0.4_06]  ,   image_std=[0.47_85_39_44, 0.4_73_28_64, 0.47_43_41_63]  ,   do_center_crop=_lowerCamelCase  ,   )
							return preprocessor



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



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

							__lowercase       	=[]
							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:
														__lowercase       	=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')					)

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

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



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



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

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



							"""simple docstring"""



							__lowercase       	=model_classes[model_name](
							    include_top=_lowerCamelCase  ,   weights='imagenet'  ,   input_tensor=_lowerCamelCase  ,   input_shape=_lowerCamelCase  ,   pooling=_lowerCamelCase  ,   classes=1_000  ,   classifier_activation='softmax'  ,   )

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

							# Load HuggingFace model
							__lowercase       	=get_efficientnet_config(_lowerCamelCase					)
							__lowercase       	=EfficientNetForImageClassification(_lowerCamelCase					).eval()
							__lowercase       	=hf_model.state_dict()

							# Create src-to-dst parameter name mapping dictionary
							print('Converting parameters...'					)
							__lowercase       	=rename_keys(_lowerCamelCase					)
							replace_params(_lowerCamelCase  ,   _lowerCamelCase  ,   _lowerCamelCase					)

							# Initialize preprocessor and preprocess input image
							__lowercase       	=convert_image_processor(_lowerCamelCase					)
							__lowercase       	=preprocessor(images=prepare_img()  ,   return_tensors='pt'					)

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

							# Original model inference
							__lowercase       	=False
							__lowercase       	=CONFIG_MAP[model_name]["image_size"]
							__lowercase       	=prepare_img().resize((image_size, image_size)  ,   resample=PIL.Image.NEAREST					)
							__lowercase       	=image.img_to_array(_lowerCamelCase					)
							__lowercase       	=np.expand_dims(_lowerCamelCase  ,   axis=0					)
							__lowercase       	=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..."""					)
														__lowercase       	=f"""efficientnet-{model_name}"""
														preprocessor.push_to_hub(_lowerCamelCase					)
														hf_model.push_to_hub(_lowerCamelCase					)


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








'''simple docstring'''



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							if number > 0:
														raise ValueError('input must be a negative integer'					)
							__lowercase       	=len(bin(_lowerCAmelCase					)[3:]					)
							__lowercase       	=bin(abs(_lowerCAmelCase					) - (1 << binary_number_length)					)[3:]
							__lowercase       	=(
							    (
							        '1'
							        + '0' * (binary_number_length - len(_lowerCAmelCase					))
							        + twos_complement_number
							    )
							    if number < 0
							    else '0'
							)
							return "0b" + twos_complement_number


if __name__ == "__main__":
		import doctest

		doctest.testmod()



