






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     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = 0      ,		_SCREAMING_SNAKE_CASE = False      ,		_SCREAMING_SNAKE_CASE = 1      ,		)  ->				None:



        """simple docstring"""






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


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






        UpperCamelCase      = self.conv(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = self.batch_norm(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = self.activation(_SCREAMING_SNAKE_CASE				)

        return output






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				None:



        """simple docstring"""






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


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






        UpperCamelCase      = input
        for layer in self.layers:
              UpperCamelCase      = layer(_SCREAMING_SNAKE_CASE				)
        return hidden_state






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				None:



        """simple docstring"""






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


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				List[torch.Tensor]:



        """simple docstring"""






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






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				Optional[int]:



        """simple docstring"""






        super().__init__()

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

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

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

  def   A__						(				self				)  ->				Tuple:



        """simple docstring"""






        self.apply(self._init_weights				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				Dict:



        """simple docstring"""






        if isinstance(_SCREAMING_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   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				int:



        """simple docstring"""






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

        return output


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






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

        laterals.append(self.psp_forward(_SCREAMING_SNAKE_CASE				)				)

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

        # build outputs
        UpperCamelCase      = [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				):
              UpperCamelCase      = nn.functional.interpolate(
                  fpn_outs[i]      ,		size=fpn_outs[0].shape[2:]      ,		mode="""bilinear"""      ,		align_corners=self.align_corners				)
        UpperCamelCase      = torch.cat(_SCREAMING_SNAKE_CASE      ,		dim=1				)
        UpperCamelCase      = self.fpn_bottleneck(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = self.classifier(_SCREAMING_SNAKE_CASE				)

        return output






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = 2      ,		_SCREAMING_SNAKE_CASE = 3      ,		_SCREAMING_SNAKE_CASE = 1				)  ->				None:



        """simple docstring"""






        super().__init__()

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

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

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

  def   A__						(				self				)  ->				Dict:



        """simple docstring"""






        self.apply(self._init_weights				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				Optional[int]:



        """simple docstring"""






        if isinstance(_SCREAMING_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   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






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






class     a_						(      lowerCamelCase    ):
  lowercase							=				UperNetConfig
  lowercase							=				"""pixel_values"""
  lowercase							=				True

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				str:



        """simple docstring"""






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

  def   A__						(				self				)  ->				Dict:



        """simple docstring"""






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


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE=False				)  ->				Union[str, Any]:



        """simple docstring"""






        if isinstance(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = value


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

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

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				List[Any]:



        """simple docstring"""






        super().__init__(_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = AutoBackbone.from_config(config.backbone_config				)

        # Semantic segmentation head(s)
        UpperCamelCase      = UperNetHead(_SCREAMING_SNAKE_CASE      ,		in_channels=self.backbone.channels				)
        UpperCamelCase      = UperNetFCNHead(_SCREAMING_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=_SCREAMING_SNAKE_CASE      ,		config_class=_CONFIG_FOR_DOC				)
  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		)  ->				Union[tuple, SemanticSegmenterOutput]:



        """simple docstring"""






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

        UpperCamelCase      = self.backbone.forward_with_filtered_kwargs(
            _SCREAMING_SNAKE_CASE      ,		output_hidden_states=_SCREAMING_SNAKE_CASE      ,		output_attentions=_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = outputs.feature_maps

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

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

        UpperCamelCase      = 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
                    UpperCamelCase      = CrossEntropyLoss(ignore_index=self.config.loss_ignore_index				)
                    UpperCamelCase      = loss_fct(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
                    UpperCamelCase      = loss_fct(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
                    UpperCamelCase      = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

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








'''simple docstring'''

import importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Union

import torch

from ..utils import BaseOutput


SCREAMING_SNAKE_CASE__     =       'scheduler_config.json'






class     a_						(      lowerCamelCase    ):
  lowercase							=				1
  lowercase							=				2
  lowercase							=				3
  lowercase							=				4
  lowercase							=				5
  lowercase							=				6
  lowercase							=				7
  lowercase							=				8
  lowercase							=				9
  lowercase							=				10
  lowercase							=				11
  lowercase							=				12
  lowercase							=				13
  lowercase							=				14






@dataclass
class     a_						(      lowerCamelCase    ):
  lowercase							=				42






class     a_						:
  lowercase							=				SCHEDULER_CONFIG_NAME
  lowercase							=				[]
  lowercase							=				True

  @classmethod
  def   A__						(				cls      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE=False      ,		**_SCREAMING_SNAKE_CASE      ,		)  ->				Optional[Any]:



        """simple docstring"""






        UpperCamelCase  ,UpperCamelCase  ,UpperCamelCase      = cls.load_config(
            pretrained_model_name_or_path=_SCREAMING_SNAKE_CASE      ,		subfolder=_SCREAMING_SNAKE_CASE      ,		return_unused_kwargs=_SCREAMING_SNAKE_CASE      ,		return_commit_hash=_SCREAMING_SNAKE_CASE      ,		**_SCREAMING_SNAKE_CASE      ,		)
        return cls.from_config(_SCREAMING_SNAKE_CASE      ,		return_unused_kwargs=_SCREAMING_SNAKE_CASE      ,		**_SCREAMING_SNAKE_CASE				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = False      ,		**_SCREAMING_SNAKE_CASE				)  ->				Dict:



        """simple docstring"""






        self.save_config(save_directory=_SCREAMING_SNAKE_CASE      ,		push_to_hub=_SCREAMING_SNAKE_CASE      ,		**_SCREAMING_SNAKE_CASE				)

  @property
  def   A__						(				self				)  ->				Tuple:



        """simple docstring"""






        return self._get_compatibles()


  @classmethod
  def   A__						(				cls				)  ->				List[Any]:



        """simple docstring"""






        UpperCamelCase      = list(set([cls.__name__] + cls._compatibles				)				)
        UpperCamelCase      = importlib.import_module(__name__.split("""."""				)[0]				)
        UpperCamelCase      = [
            getattr(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				) for c in compatible_classes_str if hasattr(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
        ]
        return compatible_classes
