

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


# General docstring
_lowerCAmelCase					:			Union[str, Any]  				=			'MobileNetV1Config'

# Base docstring
_lowerCAmelCase					:			Union[str, Any]  				=			'google/mobilenet_v1_1.0_224'
_lowerCAmelCase					:			str  				=			[1, 1_024, 7, 7]

# Image classification docstring
_lowerCAmelCase					:			str  				=			'google/mobilenet_v1_1.0_224'
_lowerCAmelCase					:			Optional[Any]  				=			'tabby, tabby cat'


_lowerCAmelCase					:			Tuple  				=			[
    '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        UpperCamelCase_(							_snake_case   :  List[str]       ,     _snake_case   :  Dict       ,     _snake_case   :  Optional[int]=None	):




			"""simple docstring"""






			__a        ={}

			if isinstance(__snake_case       ,     __snake_case	):
						__a        =model.mobilenet_va
			else:
						__a        =model

			__a        ='MobilenetV1/Conv2d_0/'
			__a        =backbone.conv_stem.convolution.weight
			__a        =backbone.conv_stem.normalization.bias
			__a        =backbone.conv_stem.normalization.weight
			__a        =backbone.conv_stem.normalization.running_mean
			__a        =backbone.conv_stem.normalization.running_var

			for i in range(13	):
						__a        =i + 1
						__a        =i * 2

						__a        =backbone.layer[pt_index]
						__a        =F'MobilenetV1/Conv2d_{tf_index}_depthwise/'
						__a        =pointer.convolution.weight
						__a        =pointer.normalization.bias
						__a        =pointer.normalization.weight
						__a        =pointer.normalization.running_mean
						__a        =pointer.normalization.running_var

						__a        =backbone.layer[pt_index + 1]
						__a        =F'MobilenetV1/Conv2d_{tf_index}_pointwise/'
						__a        =pointer.convolution.weight
						__a        =pointer.normalization.bias
						__a        =pointer.normalization.weight
						__a        =pointer.normalization.running_mean
						__a        =pointer.normalization.running_var

			if isinstance(__snake_case       ,     __snake_case	):
						__a        ='MobilenetV1/Logits/Conv2d_1c_1x1/'
						__a        =model.classifier.weight
						__a        =model.classifier.bias

			return tf_to_pt_map






def        UpperCamelCase_(							_snake_case   :  Optional[int]       ,     _snake_case   :  str       ,     _snake_case   :  Tuple	):




			"""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
			__a        =tf.train.list_variables(__snake_case	)
			__a        ={}
			for name, shape in init_vars:
						logger.info(F'Loading TF weight {name} with shape {shape}'	)
						__a        =tf.train.load_variable(__snake_case       ,     __snake_case	)
						__a        =array

			# Build TF to PyTorch weights loading map
			__a        =_build_tf_to_pytorch_map(__snake_case       ,     __snake_case       ,     __snake_case	)

			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

						__a        =tf_weights[name]

						if "depthwise_weights" in name:
									logger.info('Transposing depthwise'	)
									__a        =np.transpose(__snake_case       ,     (2, 3, 0, 1)	)
						elif "weights" in name:
									logger.info('Transposing'	)
									if len(pointer.shape	) == 2:  # copying into linear layer
												__a        =array.squeeze().transpose()
									else:
												__a        =np.transpose(__snake_case       ,     (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}'	)
						__a        =torch.from_numpy(__snake_case	)

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

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






def        UpperCamelCase_(							_snake_case   :  torch.Tensor       ,     _snake_case   :  nn.Convad	):




			"""simple docstring"""






			__a      ,    __a        =features.shape[-2:]
			__a      ,    __a        =conv_layer.stride
			__a      ,    __a        =conv_layer.kernel_size

			if in_height % stride_height == 0:
						__a        =max(kernel_height - stride_height       ,     0	)
			else:
						__a        =max(kernel_height - (in_height % stride_height)       ,     0	)

			if in_width % stride_width == 0:
						__a        =max(kernel_width - stride_width       ,     0	)
			else:
						__a        =max(kernel_width - (in_width % stride_width)       ,     0	)

			__a        =pad_along_width // 2
			__a        =pad_along_width - pad_left
			__a        =pad_along_height // 2
			__a        =pad_along_height - pad_top

			__a        =(pad_left, pad_right, pad_top, pad_bottom)
			return nn.functional.pad(__snake_case       ,     __snake_case       ,     'constant'       ,     0.0	)






class 	__magic_name__    (			nn.Module      ):





			def __init__(		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case = 1					,  __snake_case = 1					,  __snake_case = False					,  __snake_case = True					,  __snake_case = True					,  )   -> List[Any]:




						'''simple docstring'''
						super().__init__()
						__a        =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.'					)

						__a        =0 if config.tf_padding else int((kernel_size - 1) / 2					)

						__a        =nn.Convad(
						    in_channels=_lowerCamelCase					,  out_channels=_lowerCamelCase					,  kernel_size=_lowerCamelCase					,  stride=_lowerCamelCase					,  padding=_lowerCamelCase					,  groups=_lowerCamelCase					,  bias=_lowerCamelCase					,  padding_mode='zeros'					,  )

						if use_normalization:
									__a        =nn.BatchNormad(
									    num_features=_lowerCamelCase					,  eps=config.layer_norm_eps					,  momentum=0.9997					,  affine=_lowerCamelCase					,  track_running_stats=_lowerCamelCase					,  )
						else:
									__a        =None

						if use_activation:
									if isinstance(_lowerCamelCase					,  _lowerCamelCase					):
												__a        =ACTaFN[use_activation]
									elif isinstance(config.hidden_act					,  _lowerCamelCase					):
												__a        =ACTaFN[config.hidden_act]
									else:
												__a        =config.hidden_act
						else:
									__a        =None





			def 		__magic_name__    (		self					,  __snake_case					)   -> Any:




						'''simple docstring'''
						if self.config.tf_padding:
									__a        =apply_tf_padding(_lowerCamelCase					,  self.convolution					)
						__a        =self.convolution(_lowerCamelCase					)
						if self.normalization is not None:
									__a        =self.normalization(_lowerCamelCase					)
						if self.activation is not None:
									__a        =self.activation(_lowerCamelCase					)
						return features






class 	__magic_name__    (			a_      ):
			SCREAMING_SNAKE_CASE							   =						MobileNetVaConfig
			SCREAMING_SNAKE_CASE							   =						load_tf_weights_in_mobilenet_va
			SCREAMING_SNAKE_CASE							   =						"mobilenet_v1"
			SCREAMING_SNAKE_CASE							   =						"pixel_values"
			SCREAMING_SNAKE_CASE							   =						False





			def 		__magic_name__    (		self					,  __snake_case					)   -> List[str]:




						'''simple docstring'''
						if isinstance(_lowerCamelCase					,  (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(_lowerCamelCase					,  nn.BatchNormad					):
									module.bias.data.zero_()
									module.weight.data.fill_(1.0					)


_lowerCAmelCase					:			Dict  				=			R'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

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






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





			def __init__(		self					,  __snake_case					,  __snake_case = True					)   -> Union[str, Any]:




						'''simple docstring'''
						super().__init__(_lowerCamelCase					)
						__a        =config

						__a        =32
						__a        =max(int(depth * config.depth_multiplier					)					,  config.min_depth					)

						__a        =MobileNetVaConvLayer(
						    _lowerCamelCase					,  in_channels=config.num_channels					,  out_channels=_lowerCamelCase					,  kernel_size=3					,  stride=2					,  )

						__a        =[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

						__a        =nn.ModuleList()
						for i in range(13					):
									__a        =out_channels

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

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

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

						__a        =nn.AdaptiveAvgPoolad((1, 1)					) if add_pooling_layer else None

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





			def 		__magic_name__    (		self					,  __snake_case					)   -> List[Any]:




						'''simple docstring'''
						raise NotImplementedError





			@add_start_docstrings_to_model_forward(_lowerCamelCase					)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC					,  output_type=_lowerCamelCase					,  config_class=_CONFIG_FOR_DOC					,  modality='vision'					,  expected_output=_EXPECTED_OUTPUT_SHAPE					,  )
			def 		__magic_name__    (		self					,  __snake_case = None					,  __snake_case = None					,  __snake_case = None					,  )   -> Optional[int]:




						'''simple docstring'''
						__a        =(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						__a        =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'					)

						__a        =self.conv_stem(_lowerCamelCase					)

						__a        =() if output_hidden_states else None

						for i, layer_module in enumerate(self.layer					):
									__a        =layer_module(_lowerCamelCase					)

									if output_hidden_states:
												__a        =all_hidden_states + (hidden_states,)

						__a        =hidden_states

						if self.pooler is not None:
									__a        =torch.flatten(self.pooler(_lowerCamelCase					)					,  start_dim=1					)
						else:
									__a        =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=_lowerCamelCase					,  pooler_output=_lowerCamelCase					,  hidden_states=_lowerCamelCase					,  )






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





			def __init__(		self					,  __snake_case					)   -> Dict:




						'''simple docstring'''
						super().__init__(_lowerCamelCase					)

						__a        =config.num_labels
						__a        =MobileNetVaModel(_lowerCamelCase					)

						__a        =self.mobilenet_va.layer[-1].convolution.out_channels

						# Classifier head
						__a        =nn.Dropout(config.classifier_dropout_prob					,  inplace=_lowerCamelCase					)
						__a        =nn.Linear(_lowerCamelCase					,  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(_lowerCamelCase					)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT					,  output_type=_lowerCamelCase					,  config_class=_CONFIG_FOR_DOC					,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,  )
			def 		__magic_name__    (		self					,  __snake_case = None					,  __snake_case = None					,  __snake_case = None					,  __snake_case = None					,  )   -> List[Any]:




						'''simple docstring'''
						__a        =return_dict if return_dict is not None else self.config.use_return_dict

						__a        =self.mobilenet_va(_lowerCamelCase					,  output_hidden_states=_lowerCamelCase					,  return_dict=_lowerCamelCase					)

						__a        =outputs.pooler_output if return_dict else outputs[1]

						__a        =self.classifier(self.dropout(_lowerCamelCase					)					)

						__a        =None
						if labels is not None:
									if self.config.problem_type is None:
												if self.num_labels == 1:
															__a        ='regression'
												elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
															__a        ='single_label_classification'
												else:
															__a        ='multi_label_classification'

									if self.config.problem_type == "regression":
												__a        =MSELoss()
												if self.num_labels == 1:
															__a        =loss_fct(logits.squeeze()					,  labels.squeeze()					)
												else:
															__a        =loss_fct(_lowerCamelCase					,  _lowerCamelCase					)
									elif self.config.problem_type == "single_label_classification":
												__a        =CrossEntropyLoss()
												__a        =loss_fct(logits.view(-1					,  self.num_labels					)					,  labels.view(-1					)					)
									elif self.config.problem_type == "multi_label_classification":
												__a        =BCEWithLogitsLoss()
												__a        =loss_fct(_lowerCamelCase					,  _lowerCamelCase					)

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

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














"""simple docstring"""



from dataclasses import dataclass
from typing import Tuple

import numpy as np
import torch

@dataclass
class        a     :
    UpperCAmelCase_ :     torch.Tensor  # [batch_size x 3]
    UpperCAmelCase_ :     torch.Tensor  # [batch_size x 3]
    UpperCAmelCase_ :     torch.Tensor  # [batch_size x 3]
    UpperCAmelCase_ :     torch.Tensor  # [batch_size x 3]
    UpperCAmelCase_ :     int
    UpperCAmelCase_ :     int
    UpperCAmelCase_ :     float
    UpperCAmelCase_ :     float
    UpperCAmelCase_ :     Tuple[int]







    def 			UpperCamelCase_      (			self       ):
      assert self.x.shape[0] == self.y.shape[0] == self.z.shape[0] == self.origin.shape[0]
      assert self.x.shape[1] == self.y.shape[1] == self.z.shape[1] == self.origin.shape[1] == 3
      assert len(self.x.shape       ) == len(self.y.shape       ) == len(self.z.shape       ) == len(self.origin.shape       ) == 2







    def 			UpperCamelCase_      (			self       ):
      return torch.from_numpy(np.array([self.width, self.height]							,		dtype=np.floataa       )       )







    def 			UpperCamelCase_      (			self       ):
      return torch.from_numpy(np.array([self.x_fov, self.y_fov]							,		dtype=np.floataa       )       )







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       torch.arange(self.height * self.width       )
      lowercase							    =       torch.stack(
          [
              pixel_indices % self.width,
              torch.div(_lowerCamelCase							,		self.width							,		rounding_mode='trunc'       ),
          ]							,		axis=1							,		)
      return coords







    @property
    def 			UpperCamelCase_      (			self       ):
      lowercase    ,     *lowercase							    =       self.shape
      lowercase							    =       int(np.prod(_lowerCamelCase       )       )

      lowercase							    =       self.get_image_coords()
      lowercase							    =       torch.broadcast_to(coords.unsqueeze(0       )							,		[batch_size * inner_batch_size, *coords.shape]       )
      lowercase							    =       self.get_camera_rays(_lowerCamelCase       )

      lowercase							    =       rays.view(_lowerCamelCase							,		inner_batch_size * self.height * self.width							,		2							,		3       )

      return rays







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase       ):
      lowercase    ,     *lowercase    ,     lowercase							    =       coords.shape
      assert n_coords == 2
      assert batch_size == self.origin.shape[0]

      lowercase							    =       coords.view(_lowerCamelCase							,		-1							,		2       )

      lowercase							    =       self.resolution()
      lowercase							    =       self.fov()

      lowercase							    =       (flat.float() / (res - 1)) * 2 - 1
      lowercase							    =       fracs * torch.tan(fov / 2       )

      lowercase							    =       fracs.view(_lowerCamelCase							,		-1							,		2       )
      lowercase							    =       (
          self.z.view(_lowerCamelCase							,		1							,		3       )
          + self.x.view(_lowerCamelCase							,		1							,		3       ) * fracs[:, :, :1]
          + self.y.view(_lowerCamelCase							,		1							,		3       ) * fracs[:, :, 1:]
      )
      lowercase							    =       directions / directions.norm(dim=-1							,		keepdim=_lowerCamelCase       )
      lowercase							    =       torch.stack(
          [
              torch.broadcast_to(self.origin.view(_lowerCamelCase							,		1							,		3       )							,		[batch_size, directions.shape[1], 3]       ),
              directions,
          ]							,		dim=2							,		)
      return rays.view(_lowerCamelCase							,		*_lowerCamelCase							,		2							,		3       )







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase							,		_lowerCamelCase       ):
      assert width * self.height == height * self.width, "The aspect ratio should not change."
      return DifferentiableProjectiveCamera(
          origin=self.origin							,		x=self.x							,		y=self.y							,		z=self.z							,		width=_lowerCamelCase							,		height=_lowerCamelCase							,		x_fov=self.x_fov							,		y_fov=self.y_fov							,		)


def 						_SCREAMING_SNAKE_CASE (    __snake_case			:  int    ):





  '''simple docstring'''





  lowercase							    =       []
  lowercase							    =       []
  lowercase							    =       []
  lowercase							    =       []
  for theta in np.linspace(0		,					2 * np.pi		,					num=20    ):
    lowercase							    =       np.array([np.sin(__snake_case    ), np.cos(__snake_case    ), -0.5]    )
    z /= np.sqrt(np.sum(z**2    )    )
    lowercase							    =       -z * 4
    lowercase							    =       np.array([np.cos(__snake_case    ), -np.sin(__snake_case    ), 0.0]    )
    lowercase							    =       np.cross(__snake_case		,					__snake_case    )
    origins.append(__snake_case    )
    xs.append(__snake_case    )
    ys.append(__snake_case    )
    zs.append(__snake_case    )
  return DifferentiableProjectiveCamera(
      origin=torch.from_numpy(np.stack(__snake_case		,					axis=0    )    ).float()		,					x=torch.from_numpy(np.stack(__snake_case		,					axis=0    )    ).float()		,					y=torch.from_numpy(np.stack(__snake_case		,					axis=0    )    ).float()		,					z=torch.from_numpy(np.stack(__snake_case		,					axis=0    )    ).float()		,					width=__snake_case		,					height=__snake_case		,					x_fov=0.7		,					y_fov=0.7		,					shape=(1, len(__snake_case    ))		,					)


