




"""simple docstring"""



from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


SCREAMING_SNAKE_CASE    =	[
    """openmmlab/upernet-convnext-tiny""",
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
SCREAMING_SNAKE_CASE    =	"""UperNetConfig"""




class      UpperCAmelCase_  (						nn.Module							):




    def __init__(   self						:	int     ,			snake_case_						:	int     ,			snake_case_						:	int     ,			snake_case_						:	Union[int, Tuple[int, int]]     ,			snake_case_						:	Union[int, Tuple[int, int], str] = 0     ,			snake_case_						:	bool = False     ,			snake_case_						:	Union[int, Tuple[int, int]] = 1     ,			)				->							None:




     '''simple docstring'''

     super().__init__()
     A__	=    nn.Convad(
         in_channels=snake_case_     ,			out_channels=snake_case_     ,			kernel_size=snake_case_     ,			padding=snake_case_     ,			bias=snake_case_     ,			dilation=snake_case_     ,			)
     A__	=    nn.BatchNormad(snake_case_				)
     A__	=    nn.ReLU()




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




     '''simple docstring'''

     A__	=    self.conv(snake_case_				)
     A__	=    self.batch_norm(snake_case_				)
     A__	=    self.activation(snake_case_				)

     return output






class      UpperCAmelCase_  (						nn.Module							):




    def __init__(   self						:	Optional[Any]     ,			snake_case_						:	int     ,			snake_case_						:	int     ,			snake_case_						:	int				)				->							None:




     '''simple docstring'''

     super().__init__()
     A__	=    [
         nn.AdaptiveAvgPoolad(snake_case_				),
         UperNetConvModule(snake_case_     ,			snake_case_     ,			kernel_size=1				),
     ]
     for i, layer in enumerate(self.layers				):
      self.add_module(str(snake_case_				)     ,			snake_case_				)




    def 							__magic_name__							(   self						:	Union[str, Any]     ,			snake_case_						:	torch.Tensor				)				->							torch.Tensor:




     '''simple docstring'''

     A__	=    input
     for layer in self.layers:
      A__	=    layer(snake_case_				)
     return hidden_state






class      UpperCAmelCase_  (						nn.Module							):




    def __init__(   self						:	Tuple     ,			snake_case_						:	Tuple[int, ...]     ,			snake_case_						:	int     ,			snake_case_						:	int     ,			snake_case_						:	bool				)				->							None:




     '''simple docstring'''

     super().__init__()
     A__	=    pool_scales
     A__	=    align_corners
     A__	=    in_channels
     A__	=    channels
     A__	=    []
     for i, pool_scale in enumerate(snake_case_				):
      A__	=    UperNetPyramidPoolingBlock(pool_scale=snake_case_     ,			in_channels=snake_case_     ,			channels=snake_case_				)
      self.blocks.append(snake_case_				)
      self.add_module(str(snake_case_				)     ,			snake_case_				)




    def 							__magic_name__							(   self						:	Optional[int]     ,			snake_case_						:	torch.Tensor				)				->							List[torch.Tensor]:




     '''simple docstring'''

     A__	=    []
     for ppm in self.blocks:
      A__	=    ppm(snake_case_				)
      A__	=    nn.functional.interpolate(
          snake_case_     ,			size=x.size()[2:]     ,			mode="bilinear"     ,			align_corners=self.align_corners				)
      ppm_outs.append(snake_case_				)
     return ppm_outs






class      UpperCAmelCase_  (						nn.Module							):




    def __init__(   self						:	str     ,			snake_case_						:	int     ,			snake_case_						:	Any				)				->							int:




     '''simple docstring'''

     super().__init__()

     A__	=    config
     A__	=    config.pool_scales  # e.g. (1, 2, 3, 6)
     A__	=    in_channels
     A__	=    config.hidden_size
     A__	=    False
     A__	=    nn.Convad(self.channels     ,			config.num_labels     ,			kernel_size=1				)

     # PSP Module
     A__	=    UperNetPyramidPoolingModule(
         self.pool_scales     ,			self.in_channels[-1]     ,			self.channels     ,			align_corners=self.align_corners     ,			)
     A__	=    UperNetConvModule(
         self.in_channels[-1] + len(self.pool_scales				) * self.channels     ,			self.channels     ,			kernel_size=3     ,			padding=1     ,			)
     # FPN Module
     A__	=    nn.ModuleList()
     A__	=    nn.ModuleList()
     for in_channels in self.in_channels[:-1]:  # skip the top layer
      A__	=    UperNetConvModule(snake_case_     ,			self.channels     ,			kernel_size=1				)
      A__	=    UperNetConvModule(self.channels     ,			self.channels     ,			kernel_size=3     ,			padding=1				)
      self.lateral_convs.append(snake_case_				)
      self.fpn_convs.append(snake_case_				)

     A__	=    UperNetConvModule(
         len(self.in_channels				) * self.channels     ,			self.channels     ,			kernel_size=3     ,			padding=1     ,			)




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




     '''simple docstring'''

     self.apply(self._init_weights				)




    def 							__magic_name__							(   self						:	Tuple     ,			snake_case_						:	Dict				)				->							str:




     '''simple docstring'''

     if isinstance(snake_case_     ,			nn.Convad				):
      module.weight.data.normal_(mean=0.0     ,			std=self.config.initializer_range				)
      if module.bias is not None:
       module.bias.data.zero_()




    def 							__magic_name__							(   self						:	Optional[int]     ,			snake_case_						:	Optional[int]				)				->							Optional[int]:




     '''simple docstring'''

     A__	=    inputs[-1]
     A__	=    [x]
     psp_outs.extend(self.psp_modules(snake_case_				)				)
     A__	=    torch.cat(snake_case_     ,			dim=1				)
     A__	=    self.bottleneck(snake_case_				)

     return output




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




     '''simple docstring'''

     A__	=    [lateral_conv(encoder_hidden_states[i]				) for i, lateral_conv in enumerate(self.lateral_convs				)]

     laterals.append(self.psp_forward(snake_case_				)				)

     # build top-down path
     A__	=    len(snake_case_				)
     for i in range(used_backbone_levels - 1     ,			0     ,			-1				):
      A__	=    laterals[i - 1].shape[2:]
      A__	=    laterals[i - 1] + nn.functional.interpolate(
          laterals[i]     ,			size=snake_case_     ,			mode="bilinear"     ,			align_corners=self.align_corners				)

     # build outputs
     A__	=    [self.fpn_convs[i](laterals[i]				) for i in range(used_backbone_levels - 1				)]
     # append psp feature
     fpn_outs.append(laterals[-1]				)

     for i in range(used_backbone_levels - 1     ,			0     ,			-1				):
      A__	=    nn.functional.interpolate(
          fpn_outs[i]     ,			size=fpn_outs[0].shape[2:]     ,			mode="bilinear"     ,			align_corners=self.align_corners				)
     A__	=    torch.cat(snake_case_     ,			dim=1				)
     A__	=    self.fpn_bottleneck(snake_case_				)
     A__	=    self.classifier(snake_case_				)

     return output






class      UpperCAmelCase_  (						nn.Module							):




    def __init__(   self						:	Optional[int]     ,			snake_case_						:	Optional[int]     ,			snake_case_						:	int = 2     ,			snake_case_						:	int = 3     ,			snake_case_						:	Union[int, Tuple[int, int]] = 1				)				->							None:




     '''simple docstring'''

     super().__init__()

     A__	=    config
     A__	=    config.auxiliary_in_channels
     A__	=    config.auxiliary_channels
     A__	=    config.auxiliary_num_convs
     A__	=    config.auxiliary_concat_input
     A__	=    in_index

     A__	=    (kernel_size // 2) * dilation
     A__	=    []
     convs.append(
         UperNetConvModule(
             self.in_channels     ,			self.channels     ,			kernel_size=snake_case_     ,			padding=snake_case_     ,			dilation=snake_case_				)				)
     for i in range(self.num_convs - 1				):
      convs.append(
          UperNetConvModule(
              self.channels     ,			self.channels     ,			kernel_size=snake_case_     ,			padding=snake_case_     ,			dilation=snake_case_				)				)
     if self.num_convs == 0:
      A__	=    nn.Identity()
     else:
      A__	=    nn.Sequential(*snake_case_				)
     if self.concat_input:
      A__	=    UperNetConvModule(
          self.in_channels + self.channels     ,			self.channels     ,			kernel_size=snake_case_     ,			padding=kernel_size // 2				)

     A__	=    nn.Convad(self.channels     ,			config.num_labels     ,			kernel_size=1				)




    def 							__magic_name__							(   self						:	Tuple				)				->							Tuple:




     '''simple docstring'''

     self.apply(self._init_weights				)




    def 							__magic_name__							(   self						:	Union[str, Any]     ,			snake_case_						:	List[Any]				)				->							List[str]:




     '''simple docstring'''

     if isinstance(snake_case_     ,			nn.Convad				):
      module.weight.data.normal_(mean=0.0     ,			std=self.config.initializer_range				)
      if module.bias is not None:
       module.bias.data.zero_()




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




     '''simple docstring'''

     A__	=    encoder_hidden_states[self.in_index]
     A__	=    self.convs(snake_case_				)
     if self.concat_input:
      A__	=    self.conv_cat(torch.cat([hidden_states, output]     ,			dim=1				)				)
     A__	=    self.classifier(snake_case_				)
     return output






class      UpperCAmelCase_  (						__UpperCamelCase							):
    lowercase__						 =				UperNetConfig
    lowercase__						 =				'''pixel_values'''
    lowercase__						 =				True




    def 							__magic_name__							(   self						:	Optional[Any]     ,			snake_case_						:	Union[str, Any]				)				->							Optional[int]:




     '''simple docstring'''

     if isinstance(snake_case_     ,			snake_case_				):
      module.backbone.init_weights()
      module.decode_head.init_weights()
      module.auxiliary_head.init_weights()




    def 							__magic_name__							(   self						:	str				)				->							Tuple:




     '''simple docstring'''

     self.backbone.init_weights()
     self.decode_head.init_weights()
     self.auxiliary_head.init_weights()




    def 							__magic_name__							(   self						:	Any     ,			snake_case_						:	Dict     ,			snake_case_						:	Optional[int]=False				)				->							str:




     '''simple docstring'''

     if isinstance(snake_case_     ,			snake_case_				):
      A__	=    value






SCREAMING_SNAKE_CASE    =	R"""
    Parameters:
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.
        config ([`UperNetConfig`]): 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.
"""

SCREAMING_SNAKE_CASE    =	R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See
            `attentions` under returned tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers of the backbone. 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(
    '''UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes.''',       __UpperCamelCase,       )
class      UpperCAmelCase_  (						__UpperCamelCase							):




    def __init__(   self						:	Optional[int]     ,			snake_case_						:	Union[str, Any]				)				->							Dict:




     '''simple docstring'''

     super().__init__(snake_case_				)

     A__	=    AutoBackbone.from_config(config.backbone_config				)

     # Semantic segmentation head(s)
     A__	=    UperNetHead(snake_case_     ,			in_channels=self.backbone.channels				)
     A__	=    UperNetFCNHead(snake_case_				) if config.use_auxiliary_head else None

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




    @add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"				)				)
    @replace_return_docstrings(output_type=snake_case_     ,			config_class=_CONFIG_FOR_DOC				)
    def 							__magic_name__							(   self						:	Union[str, Any]     ,			snake_case_						:	Optional[torch.Tensor] = None     ,			snake_case_						:	Optional[bool] = None     ,			snake_case_						:	Optional[bool] = None     ,			snake_case_						:	Optional[torch.Tensor] = None     ,			snake_case_						:	Optional[bool] = None     ,			)				->							Union[tuple, SemanticSegmenterOutput]:




     '''simple docstring'''

     A__	=    return_dict if return_dict is not None else self.config.use_return_dict
     A__	=    (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     A__	=    output_attentions if output_attentions is not None else self.config.output_attentions

     A__	=    self.backbone.forward_with_filtered_kwargs(
         snake_case_     ,			output_hidden_states=snake_case_     ,			output_attentions=snake_case_				)
     A__	=    outputs.feature_maps

     A__	=    self.decode_head(snake_case_				)
     A__	=    nn.functional.interpolate(snake_case_     ,			size=pixel_values.shape[2:]     ,			mode="bilinear"     ,			align_corners=snake_case_				)

     A__	=    None
     if self.auxiliary_head is not None:
      A__	=    self.auxiliary_head(snake_case_				)
      A__	=    nn.functional.interpolate(
          snake_case_     ,			size=pixel_values.shape[2:]     ,			mode="bilinear"     ,			align_corners=snake_case_				)

     A__	=    None
     if labels is not None:
      if self.config.num_labels == 1:
       raise ValueError("The number of labels should be greater than one"				)
      else:
       # compute weighted loss
       A__	=    CrossEntropyLoss(ignore_index=self.config.loss_ignore_index				)
       A__	=    loss_fct(snake_case_     ,			snake_case_				)
       A__	=    loss_fct(snake_case_     ,			snake_case_				)
       A__	=    main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

     return SemanticSegmenterOutput(
         loss=snake_case_     ,			logits=snake_case_     ,			hidden_states=outputs.hidden_states     ,			attentions=outputs.attentions     ,			)







"""simple docstring"""



import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 return x + 2






class      UpperCAmelCase_  (						unittest.TestCase							):




    def 							__magic_name__							(   self						:	Any				)				->							Any:




     '''simple docstring'''

     A__	=    "x = 3"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result == 3
     self.assertDictEqual(snake_case_     ,			{"x": 3}				)

     A__	=    "x = y"
     A__	=    {"y": 5}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 5, "y": 5}				)




    def 							__magic_name__							(   self						:	Optional[Any]				)				->							Any:




     '''simple docstring'''

     A__	=    "y = add_two(x)"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 5}				)

     # Won't work without the tool
     with CaptureStdout() as out:
      A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result is None
     assert "tried to execute add_two" in out.out




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




     '''simple docstring'''

     A__	=    "x = 3"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result == 3
     self.assertDictEqual(snake_case_     ,			{"x": 3}				)




    def 							__magic_name__							(   self						:	Dict				)				->							Optional[Any]:




     '''simple docstring'''

     A__	=    "test_dict = {'x': x, 'y': add_two(x)}"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 5}				)
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_dict": {"x": 3, "y": 5}}				)




    def 							__magic_name__							(   self						:	int				)				->							str:




     '''simple docstring'''

     A__	=    "x = 3\ny = 5"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 5}				)




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




     '''simple docstring'''

     A__	=    "text = f'This is x: {x}.'"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == "This is x: 3."
     self.assertDictEqual(snake_case_     ,			{"x": 3, "text": "This is x: 3."}				)




    def 							__magic_name__							(   self						:	str				)				->							Optional[int]:




     '''simple docstring'''

     A__	=    "if x <= 3:\n    y = 2\nelse:\n    y = 5"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 2
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 2}				)

     A__	=    {"x": 8}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 8, "y": 5}				)




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




     '''simple docstring'''

     A__	=    "test_list = [x, add_two(x)]"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     self.assertListEqual(snake_case_     ,			[3, 5]				)
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_list": [3, 5]}				)




    def 							__magic_name__							(   self						:	Tuple				)				->							int:




     '''simple docstring'''

     A__	=    "y = x"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result == 3
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 3}				)




    def 							__magic_name__							(   self						:	str				)				->							Dict:




     '''simple docstring'''

     A__	=    "test_list = [x, add_two(x)]\ntest_list[1]"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_list": [3, 5]}				)

     A__	=    "test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_dict": {"x": 3, "y": 5}}				)




    def 							__magic_name__							(   self						:	Optional[int]				)				->							Tuple:




     '''simple docstring'''

     A__	=    "x = 0\nfor i in range(3):\n    x = i"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{"range": range}     ,			state=snake_case_				)
     assert result == 2
     self.assertDictEqual(snake_case_     ,			{"x": 2, "i": 2}				)





