


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

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

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

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

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







def     _UpperCamelCase      (		):
			UpperCAmelCase__  :					Optional[int]        =     """http://images.cocodataset.org/val2017/000000039769.jpg"""
			UpperCAmelCase__  :					Optional[Any]        =     Image.open(requests.get(UpperCamelCase__ ,      stream=UpperCamelCase__					).raw					)
			return im







def     _UpperCamelCase      (		UpperCamelCase__					):
			UpperCAmelCase__  :					Optional[int]        =     CONFIG_MAP[model_name]["""image_size"""]
			UpperCAmelCase__  :					int        =     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=UpperCamelCase__ ,      )
			return preprocessor







def     _UpperCamelCase      (		UpperCamelCase__					):
			UpperCAmelCase__  :					List[str]        =     [v.split("""_"""					)[0].split("""block"""					)[1] for v in original_param_names if v.startswith("""block"""					)]
			UpperCAmelCase__  :					Tuple        =     sorted(set(UpperCamelCase__					)					)
			UpperCAmelCase__  :					List[Any]        =     len(UpperCamelCase__					)
			UpperCAmelCase__  :					Union[str, Any]        =     {b: str(UpperCamelCase__					) for b, i in zip(UpperCamelCase__ ,      range(UpperCamelCase__					)					)}

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

			UpperCAmelCase__  :					Dict        =     """classifier.weight"""
			UpperCAmelCase__  :					Optional[int]        =     """classifier.bias"""
			return key_mapping







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

						UpperCAmelCase__  :					List[str]        =     key_mapping[key]
						if "_conv" in key and "kernel" in key:
									UpperCAmelCase__  :					int        =     torch.from_numpy(UpperCamelCase__					).permute(3 ,      2 ,      0 ,      1					)
						elif "depthwise_kernel" in key:
									UpperCAmelCase__  :					int        =     torch.from_numpy(UpperCamelCase__					).permute(2 ,      3 ,      0 ,      1					)
						elif "kernel" in key:
									UpperCAmelCase__  :					Optional[Any]        =     torch.from_numpy(np.transpose(UpperCamelCase__					)					)
						else:
									UpperCAmelCase__  :					Optional[int]        =     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     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					):
			UpperCAmelCase__  :					Any        =     model_classes[model_name](
			    include_top=UpperCamelCase__ ,      weights="""imagenet""" ,      input_tensor=UpperCamelCase__ ,      input_shape=UpperCamelCase__ ,      pooling=UpperCamelCase__ ,      classes=1_0_0_0 ,      classifier_activation="""softmax""" ,      )

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

			# Load HuggingFace model
			UpperCAmelCase__  :					Any        =     get_efficientnet_config(UpperCamelCase__					)
			UpperCAmelCase__  :					Dict        =     EfficientNetForImageClassification(UpperCamelCase__					).eval()
			UpperCAmelCase__  :					Optional[int]        =     hf_model.state_dict()

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

			# Initialize preprocessor and preprocess input image
			UpperCAmelCase__  :					Any        =     convert_image_processor(UpperCamelCase__					)
			UpperCAmelCase__  :					List[Any]        =     preprocessor(images=prepare_img() ,      return_tensors="""pt"""					)

			# HF model inference
			hf_model.eval()
			with torch.no_grad():
						UpperCAmelCase__  :					List[str]        =     hf_model(**UpperCamelCase__					)
			UpperCAmelCase__  :					Dict        =     outputs.logits.detach().numpy()

			# Original model inference
			UpperCAmelCase__  :					Union[str, Any]        =     False
			UpperCAmelCase__  :					Union[str, Any]        =     CONFIG_MAP[model_name]["""image_size"""]
			UpperCAmelCase__  :					Optional[int]        =     prepare_img().resize((image_size, image_size) ,      resample=PIL.Image.NEAREST					)
			UpperCAmelCase__  :					Optional[Any]        =     image.img_to_array(UpperCamelCase__					)
			UpperCAmelCase__  :					Any        =     np.expand_dims(UpperCamelCase__ ,      axis=0					)
			UpperCAmelCase__  :					Union[str, Any]        =     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...'''					)
						UpperCAmelCase__  :					Any        =     f'''efficientnet-{model_name}'''
						preprocessor.push_to_hub(UpperCamelCase__					)
						hf_model.push_to_hub(UpperCamelCase__					)


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

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



'''simple docstring'''



from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from ..image_utils import load_image


if is_torch_available():
		import torch

		from ..models.auto.modeling_auto import MODEL_FOR_OBJECT_DETECTION_MAPPING, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING

__A	=logging.get_logger(__name__)


__A	=Dict[str, Any]
__A	=List[Prediction]


@add_end_docstrings(a__	)
class 		_snake_case					(     a__	):




					def __init__(      self							,							*_lowerCamelCase							,							**_lowerCamelCase):
								super().__init__(*_lowerCamelCase							,							**_lowerCamelCase)

								if self.framework == "tf":
											raise ValueError(f'''The {self.__class__} is only available in PyTorch.''')

								requires_backends(self							,							"""vision""")
								self.check_model_type(
								    dict(MODEL_FOR_OBJECT_DETECTION_MAPPING.items() + MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items()))




					def 				snake_case__						(      self							,							**_lowerCamelCase):
								UpperCAmelCase__  :					int        =     {}
								if "threshold" in kwargs:
											UpperCAmelCase__  :					List[str]        =     kwargs["""threshold"""]
								return {}, {}, postprocess_kwargs




					def __call__(      self							,							*_lowerCamelCase							,							**_lowerCamelCase):
								return super().__call__(*_lowerCamelCase							,							**_lowerCamelCase)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					List[Any]        =     load_image(_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     torch.IntTensor([[image.height, image.width]])
								UpperCAmelCase__  :					str        =     self.image_processor(images=[image]							,							return_tensors="""pt""")
								if self.tokenizer is not None:
											UpperCAmelCase__  :					Tuple        =     self.tokenizer(text=inputs["""words"""]							,							boxes=inputs["""boxes"""]							,							return_tensors="""pt""")
								UpperCAmelCase__  :					str        =     target_size
								return inputs




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					Optional[int]        =     model_inputs.pop("""target_size""")
								UpperCAmelCase__  :					Optional[Any]        =     self.model(**_lowerCamelCase)
								UpperCAmelCase__  :					Union[str, Any]        =     outputs.__class__({"""target_size""": target_size, **outputs})
								if self.tokenizer is not None:
											UpperCAmelCase__  :					int        =     model_inputs["""bbox"""]
								return model_outputs




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase=0.9):
								UpperCAmelCase__  :					Optional[Any]        =     model_outputs["""target_size"""]
								if self.tokenizer is not None:
											# This is a LayoutLMForTokenClassification variant.
											# The OCR got the boxes and the model classified the words.
											UpperCAmelCase__   ,       UpperCAmelCase__  :					Optional[int]        =     target_size[0].tolist()

											def unnormalize(_lowerCamelCase):
														return self._get_bounding_box(
														    torch.Tensor(
														        [
														            (width * bbox[0] / 1000),
														            (height * bbox[1] / 1000),
														            (width * bbox[2] / 1000),
														            (height * bbox[3] / 1000),
														        ]))

											UpperCAmelCase__   ,       UpperCAmelCase__  :					Tuple        =     model_outputs["""logits"""].squeeze(0).softmax(dim=-1).max(dim=-1)
											UpperCAmelCase__  :					List[Any]        =     [self.model.config.idalabel[prediction] for prediction in classes.tolist()]
											UpperCAmelCase__  :					Tuple        =     [unnormalize(_lowerCamelCase) for bbox in model_outputs["""bbox"""].squeeze(0)]
											UpperCAmelCase__  :					Tuple        =     ["""score""", """label""", """box"""]
											UpperCAmelCase__  :					Dict        =     [dict(zip(_lowerCamelCase							,							_lowerCamelCase)) for vals in zip(scores.tolist()							,							_lowerCamelCase							,							_lowerCamelCase) if vals[0] > threshold]
								else:
											# This is a regular ForObjectDetectionModel
											UpperCAmelCase__  :					Union[str, Any]        =     self.image_processor.post_process_object_detection(_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase)
											UpperCAmelCase__  :					int        =     raw_annotations[0]
											UpperCAmelCase__  :					Any        =     raw_annotation["""scores"""]
											UpperCAmelCase__  :					List[str]        =     raw_annotation["""labels"""]
											UpperCAmelCase__  :					str        =     raw_annotation["""boxes"""]

											UpperCAmelCase__  :					Tuple        =     scores.tolist()
											UpperCAmelCase__  :					int        =     [self.model.config.idalabel[label.item()] for label in labels]
											UpperCAmelCase__  :					int        =     [self._get_bounding_box(_lowerCamelCase) for box in boxes]

											# {"scores": [...], ...} --> [{"score":x, ...}, ...]
											UpperCAmelCase__  :					Dict        =     ["""score""", """label""", """box"""]
											UpperCAmelCase__  :					Optional[int]        =     [
											    dict(zip(_lowerCamelCase							,							_lowerCamelCase))
											    for vals in zip(raw_annotation["""scores"""]							,							raw_annotation["""labels"""]							,							raw_annotation["""boxes"""])
											]

								return annotation




					def 				snake_case__						(      self							,							_lowerCamelCase):
								if self.framework != "pt":
											raise ValueError("""The ObjectDetectionPipeline is only available in PyTorch.""")
								UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__  :					Union[str, Any]        =     box.int().tolist()
								UpperCAmelCase__  :					Dict        =     {
								    """xmin""": xmin,
								    """ymin""": ymin,
								    """xmax""": xmax,
								    """ymax""": ymax,
								}
								return bbox