


'''simple docstring'''







import random
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  import PIL

  from transformers import FlavaImageProcessor
  from transformers.image_utils import PILImageResampling
  from transformers.models.flava.image_processing_flava import (
      FLAVA_CODEBOOK_MEAN,
      FLAVA_CODEBOOK_STD,
      FLAVA_IMAGE_MEAN,
      FLAVA_IMAGE_STD,
  )
else:
  A__:   Dict    =	None







class A__	(						unittest.TestCase				):






      def __init__(			self      :Dict					,						SCREAMING_SNAKE_CASE      :Optional[int]					,						SCREAMING_SNAKE_CASE      :Any=7					,						SCREAMING_SNAKE_CASE      :Tuple=3					,						SCREAMING_SNAKE_CASE      :int=3_0					,						SCREAMING_SNAKE_CASE      :str=4_0_0					,						SCREAMING_SNAKE_CASE      :int=True					,						SCREAMING_SNAKE_CASE      :List[str]=None					,						SCREAMING_SNAKE_CASE      :Optional[Any]=True					,						SCREAMING_SNAKE_CASE      :int=None					,						SCREAMING_SNAKE_CASE      :Any=None					,						SCREAMING_SNAKE_CASE      :Optional[int]=True					,						SCREAMING_SNAKE_CASE      :int=1 / 2_5_5					,						SCREAMING_SNAKE_CASE      :Optional[Any]=True					,						SCREAMING_SNAKE_CASE      :Dict=FLAVA_IMAGE_MEAN					,						SCREAMING_SNAKE_CASE      :int=FLAVA_IMAGE_STD					,						SCREAMING_SNAKE_CASE      :Any=1_4					,						SCREAMING_SNAKE_CASE      :Union[str, Any]=7_5					,						SCREAMING_SNAKE_CASE      :Optional[int]=None					,						SCREAMING_SNAKE_CASE      :List[Any]=1_6					,						SCREAMING_SNAKE_CASE      :Tuple=0.3					,						SCREAMING_SNAKE_CASE      :Union[str, Any]=None					,						SCREAMING_SNAKE_CASE      :Optional[Any]=True					,						SCREAMING_SNAKE_CASE      :List[Any]=None					,						SCREAMING_SNAKE_CASE      :Optional[int]=None					,						SCREAMING_SNAKE_CASE      :Dict=True					,						SCREAMING_SNAKE_CASE      :int=None					,						SCREAMING_SNAKE_CASE      :Dict=True					,						SCREAMING_SNAKE_CASE      :int=True					,						SCREAMING_SNAKE_CASE      :Optional[int]=FLAVA_CODEBOOK_MEAN					,						SCREAMING_SNAKE_CASE      :Union[str, Any]=FLAVA_CODEBOOK_STD					,						)	-> Any:




             '''simple docstring'''
             _a   :			Union[str, Any]					 =size if size is not None else {"""height""": 2_2_4, """width""": 2_2_4}
             _a   :			int					 =crop_size if crop_size is not None else {"""height""": 2_2_4, """width""": 2_2_4}
             _a   :			int					 =codebook_size if codebook_size is not None else {"""height""": 1_1_2, """width""": 1_1_2}
             _a   :			Dict					 =codebook_crop_size if codebook_crop_size is not None else {"""height""": 1_1_2, """width""": 1_1_2}

             _a   :			int					 =parent
             _a   :			Dict					 =batch_size
             _a   :			Dict					 =num_channels
             _a   :			str					 =do_resize
             _a   :			Dict					 =do_rescale
             _a   :			Tuple					 =rescale_factor
             _a   :			List[Any]					 =min_resolution
             _a   :			Optional[int]					 =max_resolution
             _a   :			int					 =size
             _a   :			str					 =resample if resample is not None else PILImageResampling.BICUBIC
             _a   :			str					 =do_normalize
             _a   :			List[str]					 =image_mean
             _a   :			List[str]					 =image_std
             _a   :			int					 =do_center_crop
             _a   :			str					 =crop_size

             _a   :			Any					 =input_size_patches
             _a   :			List[Any]					 =total_mask_patches
             _a   :			Tuple					 =mask_group_max_patches
             _a   :			List[Any]					 =mask_group_min_patches
             _a   :			int					 =mask_group_min_aspect_ratio
             _a   :			Tuple					 =mask_group_max_aspect_ratio

             _a   :			List[str]					 =codebook_do_resize
             _a   :			Optional[Any]					 =codebook_size
             _a   :			int					 =codebook_resample if codebook_resample is not None else PILImageResampling.LANCZOS
             _a   :			str					 =codebook_do_center_crop
             _a   :			Any					 =codebook_crop_size
             _a   :			Optional[int]					 =codebook_do_map_pixels
             _a   :			Optional[Any]					 =codebook_do_normalize
             _a   :			Union[str, Any]					 =codebook_image_mean
             _a   :			Any					 =codebook_image_std






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




             '''simple docstring'''
             return {
                 "image_mean": self.image_mean,
                 "image_std": self.image_std,
                 "do_normalize": self.do_normalize,
                 "do_resize": self.do_resize,
                 "size": self.size,
                 "resample": self.resample,
                 "do_rescale": self.do_rescale,
                 "rescale_factor": self.rescale_factor,
                 "do_center_crop": self.do_center_crop,
                 "crop_size": self.crop_size,
                 "input_size_patches": self.input_size_patches,
                 "total_mask_patches": self.total_mask_patches,
                 "mask_group_max_patches": self.mask_group_max_patches,
                 "mask_group_min_patches": self.mask_group_min_patches,
                 "mask_group_min_aspect_ratio": self.mask_group_min_aspect_ratio,
                 "mask_group_max_aspect_ratio": self.mask_group_min_aspect_ratio,
                 "codebook_do_resize": self.codebook_do_resize,
                 "codebook_size": self.codebook_size,
                 "codebook_resample": self.codebook_resample,
                 "codebook_do_center_crop": self.codebook_do_center_crop,
                 "codebook_crop_size": self.codebook_crop_size,
                 "codebook_do_map_pixels": self.codebook_do_map_pixels,
                 "codebook_do_normalize": self.codebook_do_normalize,
                 "codebook_image_mean": self.codebook_image_mean,
                 "codebook_image_std": self.codebook_image_std,
             }






      def __UpperCAmelCase	(			self      :Tuple     )	-> Union[str, Any]:




             '''simple docstring'''
             return (self.size["height"], self.size["width"])






      def __UpperCAmelCase	(			self      :str     )	-> str:




             '''simple docstring'''
             return (
                 (self.input_size_patches, self.input_size_patches)
                 if not isinstance(self.input_size_patches					,						_lowerCamelCase     )
                 else self.input_size_patches
             )






      def __UpperCAmelCase	(			self      :int     )	-> Union[str, Any]:




             '''simple docstring'''
             return (self.codebook_size["height"], self.codebook_size["width"])







@require_torch
@require_vision
class A__	(						_lowerCAmelCase			,					unittest.TestCase				):
      __UpperCamelCase	:  int     	=							FlavaImageProcessor if is_vision_available() else None
      __UpperCamelCase	:  Dict     	=							None






      def __UpperCAmelCase	(			self      :int     )	-> Optional[Any]:




             '''simple docstring'''
             _a   :			Optional[int]					 =FlavaImageProcessingTester(self     )






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




             '''simple docstring'''
             return self.image_processor_tester.prepare_image_processor_dict()






      def __UpperCAmelCase	(			self      :Tuple     )	-> int:




             '''simple docstring'''
             _a   :			Dict					 =self.image_processing_class(**self.image_processor_dict     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""image_mean"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""image_std"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""do_normalize"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""do_resize"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""resample"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""crop_size"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""do_center_crop"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""do_rescale"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""rescale_factor"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""masking_generator"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_do_resize"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_size"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_resample"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_do_center_crop"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_crop_size"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_do_map_pixels"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_do_normalize"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_image_mean"""     )     )
             self.assertTrue(hasattr(_lowerCamelCase					,						"""codebook_image_std"""     )     )






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




             '''simple docstring'''
             _a   :			str					 =self.image_processing_class.from_dict(self.image_processor_dict     )
             self.assertEqual(image_processor.size					,						{"""height""": 2_2_4, """width""": 2_2_4}     )
             self.assertEqual(image_processor.crop_size					,						{"""height""": 2_2_4, """width""": 2_2_4}     )
             self.assertEqual(image_processor.codebook_size					,						{"""height""": 1_1_2, """width""": 1_1_2}     )
             self.assertEqual(image_processor.codebook_crop_size					,						{"""height""": 1_1_2, """width""": 1_1_2}     )

             _a   :			Any					 =self.image_processing_class.from_dict(
                 self.image_processor_dict					,						size=4_2					,						crop_size=8_4					,						codebook_size=3_3					,						codebook_crop_size=6_6     )
             self.assertEqual(image_processor.size					,						{"""height""": 4_2, """width""": 4_2}     )
             self.assertEqual(image_processor.crop_size					,						{"""height""": 8_4, """width""": 8_4}     )
             self.assertEqual(image_processor.codebook_size					,						{"""height""": 3_3, """width""": 3_3}     )
             self.assertEqual(image_processor.codebook_crop_size					,						{"""height""": 6_6, """width""": 6_6}     )






      def __UpperCAmelCase	(			self      :Optional[int]     )	-> Dict:




             '''simple docstring'''
             pass






      def __UpperCAmelCase	(			self      :str     )	-> List[str]:




             '''simple docstring'''
             _a   :			Any					 =self.image_processing_class(**self.image_processor_dict     )
             # create random PIL images
             _a   :			Tuple					 =prepare_image_inputs(self.image_processor_tester					,						equal_resolution=_lowerCamelCase     )
             for image in image_inputs:
                    self.assertIsInstance(_lowerCamelCase					,						PIL.Image.Image     )

             # Test not batched input
             _a   :			Union[str, Any]					 =image_processing(image_inputs[0]					,						return_tensors="""pt"""     )

             # Test no bool masked pos
             self.assertFalse("""bool_masked_pos""" in encoded_images     )

             _a ,   _a   :			int					 =self.image_processor_tester.get_expected_image_size()

             self.assertEqual(
                 encoded_images.pixel_values.shape					,						(1, self.image_processor_tester.num_channels, expected_height, expected_width)					,						)

             # Test batched
             _a   :			str					 =image_processing(_lowerCamelCase					,						return_tensors="""pt"""     )
             _a ,   _a   :			Union[str, Any]					 =self.image_processor_tester.get_expected_image_size()

             # Test no bool masked pos
             self.assertFalse("""bool_masked_pos""" in encoded_images     )

             self.assertEqual(
                 encoded_images.pixel_values.shape					,						(
                     self.image_processor_tester.batch_size,
                     self.image_processor_tester.num_channels,
                     expected_height,
                     expected_width,
                 )					,						)






      def __UpperCAmelCase	(			self      :Tuple					,						SCREAMING_SNAKE_CASE      :Union[str, Any]					,						SCREAMING_SNAKE_CASE      :Tuple     )	-> List[Any]:




             '''simple docstring'''
             _a   :			int					 =self.image_processing_class(**self.image_processor_dict     )
             # create random tensors
             _a   :			int					 =prepare_image_inputs(self.image_processor_tester					,						equal_resolution=_lowerCamelCase					,						**_lowerCamelCase     )
             for image in image_inputs:
                    self.assertIsInstance(_lowerCamelCase					,						_lowerCamelCase     )

             # Test not batched input
             _a   :			Any					 =image_processing(image_inputs[0]					,						return_tensors="""pt"""     )

             _a ,   _a   :			Optional[Any]					 =self.image_processor_tester.get_expected_image_size()
             self.assertEqual(
                 encoded_images.pixel_values.shape					,						(1, self.image_processor_tester.num_channels, expected_height, expected_width)					,						)

             _a   :			str					 =image_processing(_lowerCamelCase					,						return_image_mask=_lowerCamelCase					,						return_tensors="""pt"""     )

             _a ,   _a   :			Dict					 =self.image_processor_tester.get_expected_image_size()
             self.assertEqual(
                 encoded_images.pixel_values.shape					,						(
                     self.image_processor_tester.batch_size,
                     self.image_processor_tester.num_channels,
                     expected_height,
                     expected_width,
                 )					,						)

             _a ,   _a   :			Optional[int]					 =self.image_processor_tester.get_expected_mask_size()
             self.assertEqual(
                 encoded_images.bool_masked_pos.shape					,						(
                     self.image_processor_tester.batch_size,
                     expected_height,
                     expected_width,
                 )					,						)

             # Test batched
             _a   :			Optional[int]					 =image_processing(_lowerCamelCase					,						return_tensors="""pt"""     ).pixel_values

             _a ,   _a   :			int					 =self.image_processor_tester.get_expected_image_size()
             self.assertEqual(
                 encoded_images.shape					,						(
                     self.image_processor_tester.batch_size,
                     self.image_processor_tester.num_channels,
                     expected_height,
                     expected_width,
                 )					,						)

             # Test masking
             _a   :			Optional[Any]					 =image_processing(_lowerCamelCase					,						return_image_mask=_lowerCamelCase					,						return_tensors="""pt"""     )

             _a ,   _a   :			int					 =self.image_processor_tester.get_expected_image_size()
             self.assertEqual(
                 encoded_images.pixel_values.shape					,						(
                     self.image_processor_tester.batch_size,
                     self.image_processor_tester.num_channels,
                     expected_height,
                     expected_width,
                 )					,						)

             _a ,   _a   :			Tuple					 =self.image_processor_tester.get_expected_mask_size()
             self.assertEqual(
                 encoded_images.bool_masked_pos.shape					,						(
                     self.image_processor_tester.batch_size,
                     expected_height,
                     expected_width,
                 )					,						)






      def __UpperCAmelCase	(			self      :List[Any]     )	-> Optional[int]:




             '''simple docstring'''
             self._test_call_framework(np.ndarray					,						prepare_kwargs={"""numpify""": True}     )






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




             '''simple docstring'''
             self._test_call_framework(torch.Tensor					,						prepare_kwargs={"""torchify""": True}     )






      def __UpperCAmelCase	(			self      :Dict     )	-> Dict:




             '''simple docstring'''
             random.seed(1_2_3_4     )
             _a   :			Union[str, Any]					 =self.image_processing_class(**self.image_processor_dict     )
             _a   :			List[str]					 =prepare_image_inputs(self.image_processor_tester					,						equal_resolution=_lowerCamelCase					,						torchify=_lowerCamelCase     )

             # Test not batched input
             _a   :			Union[str, Any]					 =image_processing(image_inputs[0]					,						return_image_mask=_lowerCamelCase					,						return_tensors="""pt"""     )
             self.assertEqual(encoded_images.bool_masked_pos.sum().item()					,						7_5     )






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




             '''simple docstring'''
             _a   :			Dict					 =self.image_processing_class(**self.image_processor_dict     )
             # create random PIL images
             _a   :			Optional[int]					 =prepare_image_inputs(self.image_processor_tester					,						equal_resolution=_lowerCamelCase     )
             for image in image_inputs:
                    self.assertIsInstance(_lowerCamelCase					,						PIL.Image.Image     )

             # Test not batched input
             _a   :			int					 =image_processing(image_inputs[0]					,						return_codebook_pixels=_lowerCamelCase					,						return_tensors="""pt"""     )
             _a ,   _a   :			List[Any]					 =self.image_processor_tester.get_expected_codebook_image_size()
             self.assertEqual(
                 encoded_images.codebook_pixel_values.shape					,						(1, self.image_processor_tester.num_channels, expected_height, expected_width)					,						)

             # Test batched
             _a   :			Optional[int]					 =image_processing(_lowerCamelCase					,						return_codebook_pixels=_lowerCamelCase					,						return_tensors="""pt"""     )
             _a ,   _a   :			Optional[int]					 =self.image_processor_tester.get_expected_codebook_image_size()
             self.assertEqual(
                 encoded_images.codebook_pixel_values.shape					,						(
                     self.image_processor_tester.batch_size,
                     self.image_processor_tester.num_channels,
                     expected_height,
                     expected_width,
                 )					,						)






'''simple docstring'''







from typing import Dict, Optional, Tuple

import torch


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor  )  ->      torch.Tensor:
       _a   :			Tuple					 =boundaries[1] - boundaries[0]
       _a   :			Dict					 =boundaries + step / 2
       _a   :			List[Any]					 =torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1  )]   ,dim=0  )
       return bin_centers


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor   ,_UpperCAmelCase     :     torch.Tensor   ,)  ->      Tuple[torch.Tensor, torch.Tensor]:
       _a   :			Any					 =_calculate_bin_centers(_UpperCAmelCase  )
       return (
           torch.sum(aligned_distance_error_probs * bin_centers   ,dim=-1  ),
           bin_centers[-1],
       )


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor   ,_UpperCAmelCase     :     int = 31   ,_UpperCAmelCase     :     int = 64   ,**_UpperCAmelCase     :     Dict   ,)  ->      Dict[str, torch.Tensor]:
       _a   :			Dict					 =torch.linspace(0   ,_UpperCAmelCase   ,steps=(no_bins - 1)   ,device=logits.device  )

       _a   :			int					 =torch.nn.functional.softmax(_UpperCAmelCase   ,dim=-1  )
       _a ,   _a   :			List[Any]					 =_calculate_expected_aligned_error(
           alignment_confidence_breaks=_UpperCAmelCase   ,aligned_distance_error_probs=_UpperCAmelCase   ,)

       return {
           "aligned_confidence_probs": aligned_confidence_probs,
           "predicted_aligned_error": predicted_aligned_error,
           "max_predicted_aligned_error": max_predicted_aligned_error,
       }


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor   ,_UpperCAmelCase     :     Optional[torch.Tensor] = None   ,_UpperCAmelCase     :     int = 31   ,_UpperCAmelCase     :     int = 64   ,_UpperCAmelCase     :     float = 1e-8   ,**_UpperCAmelCase     :     Optional[Any]   ,)  ->      torch.Tensor:
       if residue_weights is None:
              _a   :			Dict					 =logits.new_ones(logits.shape[-2]  )

       _a   :			str					 =torch.linspace(0   ,_UpperCAmelCase   ,steps=(no_bins - 1)   ,device=logits.device  )

       _a   :			Optional[int]					 =_calculate_bin_centers(_UpperCAmelCase  )
       torch.sum(_UpperCAmelCase  )
       _a   :			List[str]					 =logits.shape[-2]
       _a   :			List[Any]					 =max(_UpperCAmelCase   ,19  )

       _a   :			Union[str, Any]					 =1.2_4 * (clipped_n - 15) ** (1.0 / 3) - 1.8

       _a   :			Tuple					 =torch.nn.functional.softmax(_UpperCAmelCase   ,dim=-1  )

       _a   :			Dict					 =1.0 / (1 + (bin_centers**2) / (da**2))
       _a   :			int					 =torch.sum(probs * tm_per_bin   ,dim=-1  )

       _a   :			List[Any]					 =residue_weights / (eps + residue_weights.sum())
       _a   :			List[str]					 =torch.sum(predicted_tm_term * normed_residue_mask   ,dim=-1  )

       _a   :			int					 =per_alignment * residue_weights

       _a   :			List[Any]					 =(weighted == torch.max(_UpperCAmelCase  )).nonzero()[0]
       return per_alignment[tuple(_UpperCAmelCase  )]


