

'''simple docstring'''

from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve

def 			__UpperCamelCase							(		lowercase__  :					np.ndarray,   lowercase__  :					np.ndarray,   lowercase__  :					np.ndarray	):






     '''simple docstring'''





     __lowercase     						=np.stack((horizontal_flow, vertical_flow),   2	)

     # Create a grid of all pixel coordinates and subtract the flow to get the
     # target pixels coordinates
     __lowercase     						=np.stack(
         np.meshgrid(np.arange(0,   image.shape[1]	),   np.arange(0,   image.shape[0]	)	),   2	)
     __lowercase     						=np.round(grid - flow	).astype(np.intaa	)

     # Find the locations outside of the original image
     __lowercase     						=(grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]	))
     __lowercase     						=0

     __lowercase     						=image[grid[:, :, 1], grid[:, :, 0]]

     # Set pixels at invalid locations to 0
     __lowercase     						=0

     return warped

def 			__UpperCamelCase							(		lowercase__  :					np.ndarray,   lowercase__  :					np.ndarray,   lowercase__  :					SupportsIndex,   lowercase__  :					float | None = None,   ):






     '''simple docstring'''





     if alpha is None:
          __lowercase     						=0.1

     # Initialize flow
     __lowercase     						=np.zeros_like(lowercase__	)
     __lowercase     						=np.zeros_like(lowercase__	)

     # Prepare kernels for the calculation of the derivatives and the average velocity
     __lowercase     						=np.array([[-1, 1], [-1, 1]]	) * 0.25
     __lowercase     						=np.array([[-1, -1], [1, 1]]	) * 0.25
     __lowercase     						=np.array([[1, 1], [1, 1]]	) * 0.25
     __lowercase     						=np.array(
         [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]	)

     # Iteratively refine the flow
     for _ in range(lowercase__	):
          __lowercase     						=warp(lowercase__,   lowercase__,   lowercase__	)
          __lowercase     						=convolve(lowercase__,   lowercase__	) + convolve(lowercase__,   lowercase__	)
          __lowercase     						=convolve(lowercase__,   lowercase__	) + convolve(lowercase__,   lowercase__	)
          __lowercase     						=convolve(lowercase__,   lowercase__	) + convolve(lowercase__,   -kernel_t	)

          __lowercase     						=convolve(lowercase__,   lowercase__	)
          __lowercase     						=convolve(lowercase__,   lowercase__	)

          # This updates the flow as proposed in the paper (Step 12)
          __lowercase     						=(
              derivative_x * avg_horizontal_velocity
              + derivative_y * avg_vertical_velocity
              + derivative_t
          )
          __lowercase     						=update / (alpha**2 + derivative_x**2 + derivative_y**2)

          __lowercase     						=avg_horizontal_velocity - derivative_x * update
          __lowercase     						=avg_vertical_velocity - derivative_y * update

     return horizontal_flow, vertical_flow


if __name__ == "__main__":
       import doctest

       doctest.testmod()





'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCAmelCase							=					{
    '''configuration_videomae''': ['''VIDEOMAE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VideoMAEConfig'''],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase							=					[
           '''VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''VideoMAEForPreTraining''',
           '''VideoMAEModel''',
           '''VideoMAEPreTrainedModel''',
           '''VideoMAEForVideoClassification''',
       ]

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase							=					['''VideoMAEFeatureExtractor''']
       UpperCAmelCase							=					['''VideoMAEImageProcessor''']

if TYPE_CHECKING:
       from .configuration_videomae import VIDEOMAE_PRETRAINED_CONFIG_ARCHIVE_MAP, VideoMAEConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_videomae import (
                  VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST,
                  VideoMAEForPreTraining,
                  VideoMAEForVideoClassification,
                  VideoMAEModel,
                  VideoMAEPreTrainedModel,
              )

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_videomae import VideoMAEFeatureExtractor
              from .image_processing_videomae import VideoMAEImageProcessor

else:
       import sys

       UpperCAmelCase							=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


