






from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve



def      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__	):


		'''simple docstring'''

		__A    =							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    =							np.stack(
		    np.meshgrid(np.arange(0					,     image.shape[1]	)					,     np.arange(0					,     image.shape[0]	)	)					,     2	)
		__A    =							np.round(grid - flow	).astype(np.intaa	)

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

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

		# Set pixels at invalid locations to 0
		__A    =							0

		return warped



def      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__ = None					,     ):


		'''simple docstring'''

		if alpha is None:
				__A    =							0.1

		# Initialize flow
		__A    =							np.zeros_like(lowerCAmelCase__	)
		__A    =							np.zeros_like(lowerCAmelCase__	)

		# Prepare kernels for the calculation of the derivatives and the average velocity
		__A    =							np.array([[-1, 1], [-1, 1]]	) * 0.25
		__A    =							np.array([[-1, -1], [1, 1]]	) * 0.25
		__A    =							np.array([[1, 1], [1, 1]]	) * 0.25
		__A    =							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(lowerCAmelCase__	):
				__A    =							warp(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__	)
				__A    =							convolve(lowerCAmelCase__					,     lowerCAmelCase__	) + convolve(lowerCAmelCase__					,     lowerCAmelCase__	)
				__A    =							convolve(lowerCAmelCase__					,     lowerCAmelCase__	) + convolve(lowerCAmelCase__					,     lowerCAmelCase__	)
				__A    =							convolve(lowerCAmelCase__					,     lowerCAmelCase__	) + convolve(lowerCAmelCase__					,     -kernel_t	)

				__A    =							convolve(lowerCAmelCase__					,     lowerCAmelCase__	)
				__A    =							convolve(lowerCAmelCase__					,     lowerCAmelCase__	)

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

				__A    =							avg_horizontal_velocity - derivative_x * update
				__A    =							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_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case_  :							Any            =	{
    "configuration_dpr": ["DPR_PRETRAINED_CONFIG_ARCHIVE_MAP", "DPRConfig"],
    "tokenization_dpr": [
        "DPRContextEncoderTokenizer",
        "DPRQuestionEncoderTokenizer",
        "DPRReaderOutput",
        "DPRReaderTokenizer",
    ],
}


try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							str            =	[
					    "DPRContextEncoderTokenizerFast",
					    "DPRQuestionEncoderTokenizerFast",
					    "DPRReaderTokenizerFast",
					]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							Optional[int]            =	[
					    "DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "DPRContextEncoder",
					    "DPRPretrainedContextEncoder",
					    "DPRPreTrainedModel",
					    "DPRPretrainedQuestionEncoder",
					    "DPRPretrainedReader",
					    "DPRQuestionEncoder",
					    "DPRReader",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							Dict            =	[
					    "TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "TFDPRContextEncoder",
					    "TFDPRPretrainedContextEncoder",
					    "TFDPRPretrainedQuestionEncoder",
					    "TFDPRPretrainedReader",
					    "TFDPRQuestionEncoder",
					    "TFDPRReader",
					]


if TYPE_CHECKING:
					from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
					from .tokenization_dpr import (
					    DPRContextEncoderTokenizer,
					    DPRQuestionEncoderTokenizer,
					    DPRReaderOutput,
					    DPRReaderTokenizer,
					)

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_dpr_fast import (
										    DPRContextEncoderTokenizerFast,
										    DPRQuestionEncoderTokenizerFast,
										    DPRReaderTokenizerFast,
										)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_dpr import (
										    DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    DPRContextEncoder,
										    DPRPretrainedContextEncoder,
										    DPRPreTrainedModel,
										    DPRPretrainedQuestionEncoder,
										    DPRPretrainedReader,
										    DPRQuestionEncoder,
										    DPRReader,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_dpr import (
										    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFDPRContextEncoder,
										    TFDPRPretrainedContextEncoder,
										    TFDPRPretrainedQuestionEncoder,
										    TFDPRPretrainedReader,
										    TFDPRQuestionEncoder,
										    TFDPRReader,
										)

else:
					import sys

					snake_case_  :							List[str]            =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





