






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

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

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




def        _lowerCAmelCase						(							)						->Tuple:
						A__  :					Tuple											=					"""http://images.cocodataset.org/val2017/000000039769.jpg"""
						A__  :					Optional[Any]											=					Image.open(requests.get(UpperCAmelCase__,  stream=UpperCAmelCase__ ).raw )
						return im




def        _lowerCAmelCase						(							UpperCAmelCase__    :   Optional[int] )						->List[Any]:
						A__  :					List[str]											=					CONFIG_MAP[model_name]["""image_size"""]
						A__  :					str											=					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        _lowerCAmelCase						(							UpperCAmelCase__    :   Tuple )						->Any:
						A__  :					List[Any]											=					[v.split("""_""" )[0].split("""block""" )[1] for v in original_param_names if v.startswith("""block""" )]
						A__  :					Optional[Any]											=					sorted(set(UpperCAmelCase__ ) )
						A__  :					List[str]											=					len(UpperCAmelCase__ )
						A__  :					Dict											=					{b: str(UpperCAmelCase__ ) for b, i in zip(UpperCAmelCase__,  range(UpperCAmelCase__ ) )}

						A__  :					Union[str, 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:
												A__  :					Tuple											=					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""") )

						A__  :					Dict											=					{}
						for item in rename_keys:
												if item[0] in original_param_names:
																		A__  :					Union[str, Any]											=					"""efficientnet.""" + item[1]

						A__  :					str											=					"""classifier.weight"""
						A__  :					Dict											=					"""classifier.bias"""
						return key_mapping




def        _lowerCAmelCase						(							UpperCAmelCase__    :   List[Any],  UpperCAmelCase__    :   str,  UpperCAmelCase__    :   List[str] )						->List[str]:
						for key, value in tf_params.items():
												if "normalization" in key:
																		continue

												A__  :					Union[str, Any]											=					key_mapping[key]
												if "_conv" in key and "kernel" in key:
																		A__  :					Optional[int]											=					torch.from_numpy(UpperCAmelCase__ ).permute(3,  2,  0,  1 )
												elif "depthwise_kernel" in key:
																		A__  :					int											=					torch.from_numpy(UpperCAmelCase__ ).permute(2,  3,  0,  1 )
												elif "kernel" in key:
																		A__  :					Optional[Any]											=					torch.from_numpy(np.transpose(UpperCAmelCase__ ) )
												else:
																		A__  :					List[str]											=					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        _lowerCAmelCase						(							UpperCAmelCase__    :   Tuple,  UpperCAmelCase__    :   Optional[int],  UpperCAmelCase__    :   Optional[int],  UpperCAmelCase__    :   Dict )						->Union[str, Any]:
						A__  :					Dict											=					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""",  )

						A__  :					str											=					original_model.trainable_variables
						A__  :					str											=					original_model.non_trainable_variables
						A__  :					str											=					{param.name: param.numpy() for param in tf_params}
						for param in tf_non_train_params:
												A__  :					Tuple											=					param.numpy()
						A__  :					int											=					list(tf_params.keys() )

						# Load HuggingFace model
						A__  :					Optional[Any]											=					get_efficientnet_config(UpperCAmelCase__ )
						A__  :					Union[str, Any]											=					EfficientNetForImageClassification(UpperCAmelCase__ ).eval()
						A__  :					Optional[int]											=					hf_model.state_dict()

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

						# Initialize preprocessor and preprocess input image
						A__  :					Optional[int]											=					convert_image_processor(UpperCAmelCase__ )
						A__  :					str											=					preprocessor(images=prepare_img(),  return_tensors="""pt""" )

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

						# Original model inference
						A__  :					Dict											=					False
						A__  :					Tuple											=					CONFIG_MAP[model_name]["""image_size"""]
						A__  :					Optional[Any]											=					prepare_img().resize((image_size, image_size),  resample=PIL.Image.NEAREST )
						A__  :					int											=					image.img_to_array(UpperCAmelCase__ )
						A__  :					List[Any]											=					np.expand_dims(UpperCAmelCase__,  axis=0 )
						A__  :					Dict											=					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...' )
												A__  :					Union[str, 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"""



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# 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.
from typing import TYPE_CHECKING

from ..models.auto import AutoModelForVisionaSeq
from ..utils import requires_backends
from .base import PipelineTool


if TYPE_CHECKING:
						from PIL import Image






class 							__SCREAMING_SNAKE_CASE  (			UpperCamelCase						):
							snake_case_       =							'Salesforce/blip-image-captioning-base'
							snake_case_       =							(
							    'This is a tool that generates a description of an image. It takes an input named `image` which should be the '
							    'image to caption, and returns a text that contains the description in English.'
							)
							snake_case_       =							'image_captioner'
							snake_case_       =							AutoModelForVisionaSeq

							snake_case_       =							['image']
							snake_case_       =							['text']


							def __init__(     self :  int     ,      *snake_case :  Optional[int]     ,      **snake_case :  Optional[int]							):
													'''simple docstring'''





													requires_backends(self     ,      ["""vision"""]							)
													super().__init__(*snake_case     ,      **snake_case							)


							def   _UpperCamelCase    (     self :  int     ,      snake_case :  "Image"							):
													'''simple docstring'''





													return self.pre_processor(images=snake_case     ,      return_tensors="""pt"""							)


							def   _UpperCamelCase    (     self :  int     ,      snake_case :  List[Any]							):
													'''simple docstring'''





													return self.model.generate(**snake_case							)


							def   _UpperCamelCase    (     self :  Optional[int]     ,      snake_case :  Any							):
													'''simple docstring'''





													return self.pre_processor.batch_decode(snake_case     ,      skip_special_tokens=snake_case							)[0].strip()






