





from typing import Optional, Union

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_mobilenet_va import MobileNetVaConfig


lowerCAmelCase_    =				logging.get_logger(__name__)


# General docstring
lowerCAmelCase_    =				'''MobileNetV1Config'''

# Base docstring
lowerCAmelCase_    =				'''google/mobilenet_v1_1.0_224'''
lowerCAmelCase_    =				[1, 1_0_2_4, 7, 7]

# Image classification docstring
lowerCAmelCase_    =				'''google/mobilenet_v1_1.0_224'''
lowerCAmelCase_    =				'''tabby, tabby cat'''


lowerCAmelCase_    =				[
    '''google/mobilenet_v1_1.0_224''',
    '''google/mobilenet_v1_0.75_192''',
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]






def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase=None )		->						str:






				"""simple docstring"""


				snake_case_						:    Union[str, Any]				=  {}

				if isinstance(_UpperCamelCase ,				_UpperCamelCase ):
								snake_case_						:    Optional[Any]				=  model.mobilenet_va
				else:
								snake_case_						:    Optional[Any]				=  model

				snake_case_						:    Optional[int]				=  '''MobilenetV1/Conv2d_0/'''
				snake_case_						:    List[str]				=  backbone.conv_stem.convolution.weight
				snake_case_						:    Any				=  backbone.conv_stem.normalization.bias
				snake_case_						:    List[str]				=  backbone.conv_stem.normalization.weight
				snake_case_						:    Dict				=  backbone.conv_stem.normalization.running_mean
				snake_case_						:    Optional[int]				=  backbone.conv_stem.normalization.running_var

				for i in range(13 ):
								snake_case_						:    str				=  i + 1
								snake_case_						:    List[Any]				=  i * 2

								snake_case_						:    Tuple				=  backbone.layer[pt_index]
								snake_case_						:    List[str]				=  f'''MobilenetV1/Conv2d_{tf_index}_depthwise/'''
								snake_case_						:    int				=  pointer.convolution.weight
								snake_case_						:    Tuple				=  pointer.normalization.bias
								snake_case_						:    Optional[int]				=  pointer.normalization.weight
								snake_case_						:    int				=  pointer.normalization.running_mean
								snake_case_						:    List[Any]				=  pointer.normalization.running_var

								snake_case_						:    List[Any]				=  backbone.layer[pt_index + 1]
								snake_case_						:    Dict				=  f'''MobilenetV1/Conv2d_{tf_index}_pointwise/'''
								snake_case_						:    List[Any]				=  pointer.convolution.weight
								snake_case_						:    Optional[int]				=  pointer.normalization.bias
								snake_case_						:    List[str]				=  pointer.normalization.weight
								snake_case_						:    Union[str, Any]				=  pointer.normalization.running_mean
								snake_case_						:    Optional[Any]				=  pointer.normalization.running_var

				if isinstance(_UpperCamelCase ,				_UpperCamelCase ):
								snake_case_						:    Optional[int]				=  '''MobilenetV1/Logits/Conv2d_1c_1x1/'''
								snake_case_						:    int				=  model.classifier.weight
								snake_case_						:    Tuple				=  model.classifier.bias

				return tf_to_pt_map






def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase )		->						str:






				"""simple docstring"""


				try:
								import numpy as np
								import tensorflow as tf
				except ImportError:
								logger.error(
								    '''Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see '''
								    '''https://www.tensorflow.org/install/ for installation instructions.''' )
								raise

				# Load weights from TF model
				snake_case_						:    List[str]				=  tf.train.list_variables(_UpperCamelCase )
				snake_case_						:    Dict				=  {}
				for name, shape in init_vars:
								logger.info(f'''Loading TF weight {name} with shape {shape}''' )
								snake_case_						:    Dict				=  tf.train.load_variable(_UpperCamelCase ,				_UpperCamelCase )
								snake_case_						:    Optional[int]				=  array

				# Build TF to PyTorch weights loading map
				snake_case_						:    str				=  _build_tf_to_pytorch_map(_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase )

				for name, pointer in tf_to_pt_map.items():
								logger.info(f'''Importing {name}''' )
								if name not in tf_weights:
												logger.info(f'''{name} not in tf pre-trained weights, skipping''' )
												continue

								snake_case_						:    str				=  tf_weights[name]

								if "depthwise_weights" in name:
												logger.info('''Transposing depthwise''' )
												snake_case_						:    List[str]				=  np.transpose(_UpperCamelCase ,				(2, 3, 0, 1) )
								elif "weights" in name:
												logger.info('''Transposing''' )
												if len(pointer.shape ) == 2:  # copying into linear layer
																snake_case_						:    int				=  array.squeeze().transpose()
												else:
																snake_case_						:    Union[str, Any]				=  np.transpose(_UpperCamelCase ,				(3, 2, 0, 1) )

								if pointer.shape != array.shape:
												raise ValueError(f'''Pointer shape {pointer.shape} and array shape {array.shape} mismatched''' )

								logger.info(f'''Initialize PyTorch weight {name} {array.shape}''' )
								snake_case_						:    Any				=  torch.from_numpy(_UpperCamelCase )

								tf_weights.pop(_UpperCamelCase ,				_UpperCamelCase )
								tf_weights.pop(name + '''/RMSProp''' ,				_UpperCamelCase )
								tf_weights.pop(name + '''/RMSProp_1''' ,				_UpperCamelCase )
								tf_weights.pop(name + '''/ExponentialMovingAverage''' ,				_UpperCamelCase )

				logger.info(f'''Weights not copied to PyTorch model: {", ".join(tf_weights.keys() )}''' )
				return model






def      lowerCamelCase_						(				_UpperCamelCase ,				_UpperCamelCase )		->						torch.Tensor:






				"""simple docstring"""


				snake_case_			,							snake_case_						:    Any				=  features.shape[-2:]
				snake_case_			,							snake_case_						:    str				=  conv_layer.stride
				snake_case_			,							snake_case_						:    List[Any]				=  conv_layer.kernel_size

				if in_height % stride_height == 0:
								snake_case_						:    Tuple				=  max(kernel_height - stride_height ,				0 )
				else:
								snake_case_						:    Union[str, Any]				=  max(kernel_height - (in_height % stride_height) ,				0 )

				if in_width % stride_width == 0:
								snake_case_						:    Optional[int]				=  max(kernel_width - stride_width ,				0 )
				else:
								snake_case_						:    Dict				=  max(kernel_width - (in_width % stride_width) ,				0 )

				snake_case_						:    Union[str, Any]				=  pad_along_width // 2
				snake_case_						:    Optional[Any]				=  pad_along_width - pad_left
				snake_case_						:    Optional[Any]				=  pad_along_height // 2
				snake_case_						:    Union[str, Any]				=  pad_along_height - pad_top

				snake_case_						:    Tuple				=  (pad_left, pad_right, pad_top, pad_bottom)
				return nn.functional.pad(_UpperCamelCase ,				_UpperCamelCase ,				'''constant''' ,				0.0 )



class 					__lowerCAmelCase       (			nn.Module     ):






			def __init__(self					,							__magic_name__					,							__magic_name__					,							__magic_name__					,							__magic_name__					,							__magic_name__ = 1					,							__magic_name__ = 1					,							__magic_name__ = False					,							__magic_name__ = True					,							__magic_name__ = True					,							)     ->   None:


							'''simple docstring'''



							super().__init__()
							snake_case_						:    Optional[Any]				=  config

							if in_channels % groups != 0:
											raise ValueError(F'''Input channels ({in_channels}) are not divisible by {groups} groups.'''    )
							if out_channels % groups != 0:
											raise ValueError(F'''Output channels ({out_channels}) are not divisible by {groups} groups.'''    )

							snake_case_						:    List[str]				=  0 if config.tf_padding else int((kernel_size - 1) / 2    )

							snake_case_						:    List[str]				=  nn.Convad(
							    in_channels=__magic_name__					,							out_channels=__magic_name__					,							kernel_size=__magic_name__					,							stride=__magic_name__					,							padding=__magic_name__					,							groups=__magic_name__					,							bias=__magic_name__					,							padding_mode='''zeros'''					,							)

							if use_normalization:
											snake_case_						:    Tuple				=  nn.BatchNormad(
											    num_features=__magic_name__					,							eps=config.layer_norm_eps					,							momentum=0.9_997					,							affine=__magic_name__					,							track_running_stats=__magic_name__					,							)
							else:
											snake_case_						:    str				=  None

							if use_activation:
											if isinstance(__magic_name__					,							__magic_name__    ):
															snake_case_						:    Optional[Any]				=  ACTaFN[use_activation]
											elif isinstance(config.hidden_act					,							__magic_name__    ):
															snake_case_						:    Any				=  ACTaFN[config.hidden_act]
											else:
															snake_case_						:    List[str]				=  config.hidden_act
							else:
											snake_case_						:    Any				=  None






			def 			lowerCamelCase  (self					,							__magic_name__    )     ->   torch.Tensor:


							'''simple docstring'''



							if self.config.tf_padding:
											snake_case_						:    Union[str, Any]				=  apply_tf_padding(__magic_name__					,							self.convolution    )
							snake_case_						:    Union[str, Any]				=  self.convolution(__magic_name__    )
							if self.normalization is not None:
											snake_case_						:    int				=  self.normalization(__magic_name__    )
							if self.activation is not None:
											snake_case_						:    List[Any]				=  self.activation(__magic_name__    )
							return features




class 					__lowerCAmelCase       (			_a     ):
			lowerCamelCase_     :     Tuple       			=      MobileNetVaConfig
			lowerCamelCase_     :     Optional[int]       			=      load_tf_weights_in_mobilenet_va
			lowerCamelCase_     :     int       			=      '''mobilenet_v1'''
			lowerCamelCase_     :     str       			=      '''pixel_values'''
			lowerCamelCase_     :     List[str]       			=      False






			def 			lowerCamelCase  (self					,							__magic_name__    )     ->   None:


							'''simple docstring'''



							if isinstance(__magic_name__					,							(nn.Linear, nn.Convad)    ):
											module.weight.data.normal_(mean=0.0					,							std=self.config.initializer_range    )
											if module.bias is not None:
															module.bias.data.zero_()
							elif isinstance(__magic_name__					,							nn.BatchNormad    ):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0    )




lowerCAmelCase_    =				r'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

lowerCAmelCase_    =				r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`MobileNetV1ImageProcessor.__call__`] for details.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''



@add_start_docstrings(
    '''The bare MobileNetV1 model outputting raw hidden-states without any specific head on top.''',     _a,     )
class 					__lowerCAmelCase       (			_a     ):






			def __init__(self					,							__magic_name__					,							__magic_name__ = True    )     ->   List[Any]:


							'''simple docstring'''



							super().__init__(__magic_name__    )
							snake_case_						:    Union[str, Any]				=  config

							snake_case_						:    str				=  32
							snake_case_						:    Optional[int]				=  max(int(depth * config.depth_multiplier    )					,							config.min_depth    )

							snake_case_						:    List[str]				=  MobileNetVaConvLayer(
							    __magic_name__					,							in_channels=config.num_channels					,							out_channels=__magic_name__					,							kernel_size=3					,							stride=2					,							)

							snake_case_						:    Any				=  [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

							snake_case_						:    Union[str, Any]				=  nn.ModuleList()
							for i in range(13    ):
											snake_case_						:    List[str]				=  out_channels

											if strides[i] == 2 or i == 0:
															depth *= 2
															snake_case_						:    Dict				=  max(int(depth * config.depth_multiplier    )					,							config.min_depth    )

											self.layer.append(
											    MobileNetVaConvLayer(
											        __magic_name__					,							in_channels=__magic_name__					,							out_channels=__magic_name__					,							kernel_size=3					,							stride=strides[i]					,							groups=__magic_name__					,							)    )

											self.layer.append(
											    MobileNetVaConvLayer(
											        __magic_name__					,							in_channels=__magic_name__					,							out_channels=__magic_name__					,							kernel_size=1					,							)    )

							snake_case_						:    List[str]				=  nn.AdaptiveAvgPoolad((1, 1)    ) if add_pooling_layer else None

							# Initialize weights and apply final processing
							self.post_init()






			def 			lowerCamelCase  (self					,							__magic_name__    )     ->   Dict:


							'''simple docstring'''



							raise NotImplementedError






			@add_start_docstrings_to_model_forward(__magic_name__    )
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC					,							output_type=__magic_name__					,							config_class=_CONFIG_FOR_DOC					,							modality='''vision'''					,							expected_output=_EXPECTED_OUTPUT_SHAPE					,							)
			def 			lowerCamelCase  (self					,							__magic_name__ = None					,							__magic_name__ = None					,							__magic_name__ = None					,							)     ->   Union[tuple, BaseModelOutputWithPoolingAndNoAttention]:


							'''simple docstring'''



							snake_case_						:    Tuple				=  (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case_						:    Dict				=  return_dict if return_dict is not None else self.config.use_return_dict

							if pixel_values is None:
											raise ValueError('''You have to specify pixel_values'''    )

							snake_case_						:    Union[str, Any]				=  self.conv_stem(__magic_name__    )

							snake_case_						:    Union[str, Any]				=  () if output_hidden_states else None

							for i, layer_module in enumerate(self.layer    ):
											snake_case_						:    List[str]				=  layer_module(__magic_name__    )

											if output_hidden_states:
															snake_case_						:    Union[str, Any]				=  all_hidden_states + (hidden_states,)

							snake_case_						:    Any				=  hidden_states

							if self.pooler is not None:
											snake_case_						:    Optional[int]				=  torch.flatten(self.pooler(__magic_name__    )					,							start_dim=1    )
							else:
											snake_case_						:    Tuple				=  None

							if not return_dict:
											return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None    )

							return BaseModelOutputWithPoolingAndNoAttention(
							    last_hidden_state=__magic_name__					,							pooler_output=__magic_name__					,							hidden_states=__magic_name__					,							)




@add_start_docstrings(
    '''
    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    ''',     _a,     )
class 					__lowerCAmelCase       (			_a     ):






			def __init__(self					,							__magic_name__    )     ->   None:


							'''simple docstring'''



							super().__init__(__magic_name__    )

							snake_case_						:    Dict				=  config.num_labels
							snake_case_						:    List[str]				=  MobileNetVaModel(__magic_name__    )

							snake_case_						:    Optional[int]				=  self.mobilenet_va.layer[-1].convolution.out_channels

							# Classifier head
							snake_case_						:    Union[str, Any]				=  nn.Dropout(config.classifier_dropout_prob					,							inplace=__magic_name__    )
							snake_case_						:    Any				=  nn.Linear(__magic_name__					,							config.num_labels    ) if config.num_labels > 0 else nn.Identity()

							# Initialize weights and apply final processing
							self.post_init()






			@add_start_docstrings_to_model_forward(__magic_name__    )
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT					,							output_type=__magic_name__					,							config_class=_CONFIG_FOR_DOC					,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,							)
			def 			lowerCamelCase  (self					,							__magic_name__ = None					,							__magic_name__ = None					,							__magic_name__ = None					,							__magic_name__ = None					,							)     ->   Union[tuple, ImageClassifierOutputWithNoAttention]:


							'''simple docstring'''



							snake_case_						:    List[Any]				=  return_dict if return_dict is not None else self.config.use_return_dict

							snake_case_						:    Any				=  self.mobilenet_va(__magic_name__					,							output_hidden_states=__magic_name__					,							return_dict=__magic_name__    )

							snake_case_						:    Optional[Any]				=  outputs.pooler_output if return_dict else outputs[1]

							snake_case_						:    List[Any]				=  self.classifier(self.dropout(__magic_name__    )    )

							snake_case_						:    Optional[Any]				=  None
							if labels is not None:
											if self.config.problem_type is None:
															if self.num_labels == 1:
																			snake_case_						:    List[str]				=  '''regression'''
															elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																			snake_case_						:    List[str]				=  '''single_label_classification'''
															else:
																			snake_case_						:    Optional[Any]				=  '''multi_label_classification'''

											if self.config.problem_type == "regression":
															snake_case_						:    Optional[Any]				=  MSELoss()
															if self.num_labels == 1:
																			snake_case_						:    str				=  loss_fct(logits.squeeze()					,							labels.squeeze()    )
															else:
																			snake_case_						:    str				=  loss_fct(__magic_name__					,							__magic_name__    )
											elif self.config.problem_type == "single_label_classification":
															snake_case_						:    List[str]				=  CrossEntropyLoss()
															snake_case_						:    Union[str, Any]				=  loss_fct(logits.view(-1					,							self.num_labels    )					,							labels.view(-1    )    )
											elif self.config.problem_type == "multi_label_classification":
															snake_case_						:    Tuple				=  BCEWithLogitsLoss()
															snake_case_						:    Optional[Any]				=  loss_fct(__magic_name__					,							__magic_name__    )

							if not return_dict:
											snake_case_						:    int				=  (logits,) + outputs[2:]
											return ((loss,) + output) if loss is not None else output

							return ImageClassifierOutputWithNoAttention(
							    loss=__magic_name__					,							logits=__magic_name__					,							hidden_states=outputs.hidden_states					,							)







import PIL.Image
import PIL.ImageOps
from packaging import version
from PIL import Image


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse('''9.1.0'''):
	lowerCAmelCase_    =				{
	    '''linear''': PIL.Image.Resampling.BILINEAR,
	    '''bilinear''': PIL.Image.Resampling.BILINEAR,
	    '''bicubic''': PIL.Image.Resampling.BICUBIC,
	    '''lanczos''': PIL.Image.Resampling.LANCZOS,
	    '''nearest''': PIL.Image.Resampling.NEAREST,
	}
else:
	lowerCAmelCase_    =				{
	    '''linear''': PIL.Image.LINEAR,
	    '''bilinear''': PIL.Image.BILINEAR,
	    '''bicubic''': PIL.Image.BICUBIC,
	    '''lanczos''': PIL.Image.LANCZOS,
	    '''nearest''': PIL.Image.NEAREST,
	}






def      lowerCamelCase_						(				_UpperCamelCase )		->						Optional[int]:






				"""simple docstring"""


				snake_case_						:    Dict				=  (images / 2 + 0.5).clamp(0 ,				1 )
				snake_case_						:    Dict				=  images.cpu().permute(0 ,				2 ,				3 ,				1 ).float().numpy()
				snake_case_						:    int				=  numpy_to_pil(_UpperCamelCase )
				return images






def      lowerCamelCase_						(				_UpperCamelCase )		->						List[Any]:






				"""simple docstring"""


				if images.ndim == 3:
								snake_case_						:    Optional[Any]				=  images[None, ...]
				snake_case_						:    Any				=  (images * 255).round().astype('''uint8''' )
				if images.shape[-1] == 1:
								# special case for grayscale (single channel) images
								snake_case_						:    str				=  [Image.fromarray(image.squeeze() ,				mode='''L''' ) for image in images]
				else:
								snake_case_						:    List[Any]				=  [Image.fromarray(_UpperCamelCase ) for image in images]

				return pil_images
