





'''simple docstring'''






from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve
def 				_lowerCamelCase(	UpperCamelCase__     :			np.ndarray     ,							UpperCamelCase__     :			np.ndarray     ,							UpperCamelCase__     :			np.ndarray		)		->	np.ndarray:
    A							:							List[Any]				=    np.stack((horizontal_flow, vertical_flow)     ,							2		)

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

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

    A							:							Union[str, Any]				=    image[grid[:, :, 1], grid[:, :, 0]]

    # Set pixels at invalid locations to 0
    A							:							Optional[int]				=    0

    return warped
def 				_lowerCamelCase(	UpperCamelCase__     :			np.ndarray     ,							UpperCamelCase__     :			np.ndarray     ,							UpperCamelCase__     :			SupportsIndex     ,							UpperCamelCase__     :			float | None = None     ,							)		->	tuple[np.ndarray, np.ndarray]:







    if alpha is None:
        A							:							Dict				=    0.1

    # Initialize flow
    A							:							str				=    np.zeros_like(UpperCamelCase_		)
    A							:							Dict				=    np.zeros_like(UpperCamelCase_		)

    # Prepare kernels for the calculation of the derivatives and the average velocity
    A							:							Union[str, Any]				=    np.array([[-1, 1], [-1, 1]]		) * 0.2_5
    A							:							Optional[Any]				=    np.array([[-1, -1], [1, 1]]		) * 0.2_5
    A							:							Optional[int]				=    np.array([[1, 1], [1, 1]]		) * 0.2_5
    A							:							Dict				=    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(UpperCamelCase_		):
        A							:							Union[str, Any]				=    warp(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_		)
        A							:							List[Any]				=    convolve(UpperCamelCase_     ,							UpperCamelCase_		) + convolve(UpperCamelCase_     ,							UpperCamelCase_		)
        A							:							Union[str, Any]				=    convolve(UpperCamelCase_     ,							UpperCamelCase_		) + convolve(UpperCamelCase_     ,							UpperCamelCase_		)
        A							:							Optional[int]				=    convolve(UpperCamelCase_     ,							UpperCamelCase_		) + convolve(UpperCamelCase_     ,							-kernel_t		)

        A							:							Any				=    convolve(UpperCamelCase_     ,							UpperCamelCase_		)
        A							:							str				=    convolve(UpperCamelCase_     ,							UpperCamelCase_		)

        # This updates the flow as proposed in the paper (Step 12)
        A							:							Any				=    (
            derivative_x * avg_horizontal_velocity
            + derivative_y * avg_vertical_velocity
            + derivative_t
        )
        A							:							List[str]				=    update / (alpha**2 + derivative_x**2 + derivative_y**2)

        A							:							Optional[int]				=    avg_horizontal_velocity - derivative_x * update
        A							:							Dict				=    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_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase											=     {"""configuration_pegasus""": ["""PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP""", """PegasusConfig"""]}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     ["""PegasusTokenizer"""]

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     ["""PegasusTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """PegasusForCausalLM""",
         """PegasusForConditionalGeneration""",
         """PegasusModel""",
         """PegasusPreTrainedModel""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """TFPegasusForConditionalGeneration""",
         """TFPegasusModel""",
         """TFPegasusPreTrainedModel""",
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """FlaxPegasusForConditionalGeneration""",
         """FlaxPegasusModel""",
         """FlaxPegasusPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_pegasus import PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusConfig

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_pegasus import PegasusTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_pegasus_fast import PegasusTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_pegasus import (
              PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
              PegasusForCausalLM,
              PegasusForConditionalGeneration,
              PegasusModel,
              PegasusPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel, TFPegasusPreTrainedModel

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_pegasus import (
              FlaxPegasusForConditionalGeneration,
              FlaxPegasusModel,
              FlaxPegasusPreTrainedModel,
          )

else:
     import sys

     UpperCamelCase											=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


