





"""simple docstring"""




import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union




class  _a						(      enum.Enum				):
				UpperCamelCase								=    '''timm'''
				UpperCamelCase								=    '''transformers'''





def        a_				(							_lowercase       ,						_lowercase       ,						_lowercase   ):
						if stage_names is None:
												raise ValueError('''Stage_names must be set for transformers backbones'''   )

						if out_features is not None:
												if not isinstance(_lowercase       ,						(list,)   ):
																		raise ValueError(F"""out_features must be a list {type(_lowercase   )}"""   )
												if any(feat not in stage_names for feat in out_features   ):
																		raise ValueError(F"""out_features must be a subset of stage_names: {stage_names} got {out_features}"""   )

						if out_indices is not None:
												if not isinstance(_lowercase       ,						(list, tuple)   ):
																		raise ValueError(F"""out_indices must be a list or tuple, got {type(_lowercase   )}"""   )
												if any(idx >= len(_lowercase   ) for idx in out_indices   ):
																		raise ValueError('''out_indices must be valid indices for stage_names {stage_names}, got {out_indices}'''   )

						if out_features is not None and out_indices is not None:
												if len(_lowercase   ) != len(_lowercase   ):
																		raise ValueError('''out_features and out_indices should have the same length if both are set'''   )
												if out_features != [stage_names[idx] for idx in out_indices]:
																		raise ValueError('''out_features and out_indices should correspond to the same stages if both are set'''   )





def        a_				(							_lowercase       ,						_lowercase       ,						_lowercase       ,						):
						if out_indices is None and out_features is None:
												_UpperCamelCase			:					List[Any]     =    [len(_lowercase   ) - 1]
												_UpperCamelCase			:					int     =    [stage_names[-1]]
						elif out_indices is None and out_features is not None:
												_UpperCamelCase			:					str     =    [stage_names.index(_lowercase   ) for layer in out_features]
						elif out_features is None and out_indices is not None:
												_UpperCamelCase			:					Optional[Any]     =    [stage_names[idx] for idx in out_indices]
						return out_features, out_indices





def        a_				(							_lowercase       ,						_lowercase       ,						_lowercase       ,						):
						verify_out_features_out_indices(out_features=_lowercase       ,						out_indices=_lowercase       ,						stage_names=_lowercase   )
						_UpperCamelCase   ,							_UpperCamelCase			:					Optional[int]     =    _align_output_features_output_indices(
						    out_features=_lowercase       ,						out_indices=_lowercase       ,						stage_names=_lowercase   )
						# Verify that the aligned out_features and out_indices are valid
						verify_out_features_out_indices(out_features=_lowercase       ,						out_indices=_lowercase       ,						stage_names=_lowercase   )
						return output_features, output_indices




class  _a						:
				UpperCamelCase								=    None






				def 				snake_case	( self		:    int, lowerCAmelCase__		:    Optional[Any]							)	->      None:




										'''simple docstring'''



										if getattr(self, '''_backbone''', lowerCAmelCase__							) is None:
																raise ValueError('''self._backbone must be set before calling _init_timm_backbone'''							)

										# These will diagree with the defaults for the transformers models e.g. for resnet50
										# the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
										# the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
										_UpperCamelCase			:					Dict     =    [stage['''module'''] for stage in self._backbone.feature_info.info]
										_UpperCamelCase			:					List[Any]     =    [stage['''num_chs'''] for stage in self._backbone.feature_info.info]
										_UpperCamelCase			:					Optional[int]     =    self._backbone.feature_info.out_indices
										_UpperCamelCase			:					List[Any]     =    self._backbone.feature_info.module_name()

										# We verify the out indices and out features are valid
										verify_out_features_out_indices(
										    out_features=lowerCAmelCase__, out_indices=lowerCAmelCase__, stage_names=self.stage_names							)
										_UpperCamelCase   ,							_UpperCamelCase			:					List[Any]     =    out_features, out_indices






				def 				snake_case	( self		:    Any, lowerCAmelCase__		:    Optional[int]							)	->      None:




										'''simple docstring'''



										_UpperCamelCase			:					Dict     =    getattr(lowerCAmelCase__, '''stage_names'''							)
										_UpperCamelCase			:					Optional[int]     =    getattr(lowerCAmelCase__, '''out_features''', lowerCAmelCase__							)
										_UpperCamelCase			:					Tuple     =    getattr(lowerCAmelCase__, '''out_indices''', lowerCAmelCase__							)

										_UpperCamelCase			:					str     =    stage_names
										_UpperCamelCase   ,							_UpperCamelCase			:					List[str]     =    get_aligned_output_features_output_indices(
										    out_features=lowerCAmelCase__, out_indices=lowerCAmelCase__, stage_names=lowerCAmelCase__							)
										# Number of channels for each stage. This is set in the transformer backbone model init
										_UpperCamelCase			:					Any     =    None






				def 				snake_case	( self		:    Tuple, lowerCAmelCase__		:    Optional[int]							)	->      None:




										'''simple docstring'''



										_UpperCamelCase			:					str     =    config

										_UpperCamelCase			:					Any     =    getattr(lowerCAmelCase__, '''use_timm_backbone''', lowerCAmelCase__							)
										_UpperCamelCase			:					Dict     =    BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

										if self.backbone_type == BackboneType.TIMM:
																self._init_timm_backbone(lowerCAmelCase__							)
										elif self.backbone_type == BackboneType.TRANSFORMERS:
																self._init_transformers_backbone(lowerCAmelCase__							)
										else:
																raise ValueError(f"""backbone_type {self.backbone_type} not supported."""							)






				@property
				def 				snake_case	( self		:    str							)	->      List[Any]:




										'''simple docstring'''



										return self._out_features






				@out_features.setter
				def 				snake_case	( self		:    Any, lowerCAmelCase__		:    List[str]							)	->      Dict:




										'''simple docstring'''



										_UpperCamelCase   ,							_UpperCamelCase			:					List[str]     =    get_aligned_output_features_output_indices(
										    out_features=lowerCAmelCase__, out_indices=lowerCAmelCase__, stage_names=self.stage_names							)






				@property
				def 				snake_case	( self		:    Dict							)	->      Tuple:




										'''simple docstring'''



										return self._out_indices






				@out_indices.setter
				def 				snake_case	( self		:    str, lowerCAmelCase__		:    Union[Tuple[int], List[int]]							)	->      Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase   ,							_UpperCamelCase			:					Tuple     =    get_aligned_output_features_output_indices(
										    out_features=lowerCAmelCase__, out_indices=lowerCAmelCase__, stage_names=self.stage_names							)






				@property
				def 				snake_case	( self		:    Optional[Any]							)	->      str:




										'''simple docstring'''



										return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names							)}






				@property
				def 				snake_case	( self		:    str							)	->      List[Any]:




										'''simple docstring'''



										return [self.out_feature_channels[name] for name in self.out_features]






				def 				snake_case	( self		:    Union[str, Any], *lowerCAmelCase__		:    List[str], **lowerCAmelCase__		:    Tuple							)	->      List[str]:




										'''simple docstring'''



										_UpperCamelCase			:					List[str]     =    dict(inspect.signature(self.forward							).parameters							)
										_UpperCamelCase			:					Optional[int]     =    {k: v for k, v in kwargs.items() if k in signature}
										return self(*lowerCAmelCase__, **lowerCAmelCase__							)






				def 				snake_case	( self		:    Optional[Any], lowerCAmelCase__		:    Optional[int], lowerCAmelCase__		:    Optional[bool] = None, lowerCAmelCase__		:    Optional[bool] = None, lowerCAmelCase__		:    Optional[bool] = None, )	->      Tuple:




										'''simple docstring'''



										raise NotImplementedError('''This method should be implemented by the derived class.'''							)






				def 				snake_case	( self		:    Dict							)	->      Tuple:




										'''simple docstring'''



										_UpperCamelCase			:					Tuple     =    super().to_dict()
										_UpperCamelCase			:					int     =    output.pop('''_out_features'''							)
										_UpperCamelCase			:					int     =    output.pop('''_out_indices'''							)
										return output




class  _a						:






				@property
				def 				snake_case	( self		:    List[str]							)	->      List[Any]:




										'''simple docstring'''



										return self._out_features






				@out_features.setter
				def 				snake_case	( self		:    Union[str, Any], lowerCAmelCase__		:    List[str]							)	->      str:




										'''simple docstring'''



										_UpperCamelCase   ,							_UpperCamelCase			:					Optional[Any]     =    get_aligned_output_features_output_indices(
										    out_features=lowerCAmelCase__, out_indices=lowerCAmelCase__, stage_names=self.stage_names							)






				@property
				def 				snake_case	( self		:    List[Any]							)	->      Dict:




										'''simple docstring'''



										return self._out_indices






				@out_indices.setter
				def 				snake_case	( self		:    Optional[Any], lowerCAmelCase__		:    Union[Tuple[int], List[int]]							)	->      Tuple:




										'''simple docstring'''



										_UpperCamelCase   ,							_UpperCamelCase			:					Optional[int]     =    get_aligned_output_features_output_indices(
										    out_features=lowerCAmelCase__, out_indices=lowerCAmelCase__, stage_names=self.stage_names							)






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




										'''simple docstring'''



										_UpperCamelCase			:					Tuple     =    super().to_dict()
										_UpperCamelCase			:					Union[str, Any]     =    output.pop('''_out_features'''							)
										_UpperCamelCase			:					Any     =    output.pop('''_out_indices'''							)
										return output










"""simple docstring"""


from typing import List

import PIL
import torch
from PIL import Image

from ...configuration_utils import ConfigMixin
from ...models.modeling_utils import ModelMixin
from ...utils import PIL_INTERPOLATION



class    lowercase_ (   __lowerCAmelCase ,__lowerCAmelCase							):

       def __init__(	self   ):






        '''simple docstring'''

        super().__init__()

        self.register_buffer("watermark_image"      ,    torch.zeros((62, 62, 4)   )   )
        snake_case__										= None

       def 						_SCREAMING_SNAKE_CASE(	self      ,    __SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE=None   ):






        '''simple docstring'''

        snake_case__										= images[0].height
        snake_case__										= images[0].width

        snake_case__										= sample_size or h

        snake_case__										= min(h / sample_size      ,    w / sample_size   )
        snake_case__							,      snake_case__										= (int(h / coef   ), int(w / coef   )) if coef < 1 else (h, w)

        snake_case__							,      snake_case__										= 1024**2, img_w * img_h
        snake_case__										= (Sa / Sa) ** 0.5
        snake_case__							,      snake_case__							,      snake_case__										= int(K * 62   ), img_w - int(14 * K   ), img_h - int(14 * K   )

        if self.watermark_image_as_pil is None:
         snake_case__										= self.watermark_image.to(torch.uinta   ).cpu().numpy()
         snake_case__										= Image.fromarray(__SCREAMING_SNAKE_CASE      ,    mode="RGBA"   )
         snake_case__										= watermark_image

        snake_case__										= self.watermark_image_as_pil.resize(
            (wm_size, wm_size)      ,    PIL_INTERPOLATION["bicubic"]      ,    reducing_gap=__SCREAMING_SNAKE_CASE   )

        for pil_img in images:
         pil_img.paste(__SCREAMING_SNAKE_CASE      ,    box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y)      ,    mask=wm_img.split()[-1]   )

        return images

