the hornschunck method estimates the optical flow for every single pixel of a sequence of images it works by assuming brightness constancy between two consecutive frames and smoothness in the optical flow useful resources wikipedia https en wikipedia orgwikihorne28093schunckmethod paper http image diku dkimagecanonmaterialhornschunckopticalflow pdf warps the pixels of an image into a new image using the horizontal and vertical flows pixels that are warped from an invalid location are set to 0 parameters image grayscale image horizontalflow horizontal flow verticalflow vertical flow returns warped image warpnp array0 1 2 0 3 0 2 2 2 np array0 1 1 1 0 0 1 1 1 np array0 0 0 0 1 0 0 0 1 array0 0 0 3 1 0 0 2 3 create a grid of all pixel coordinates and subtract the flow to get the target pixels coordinates find the locations outside of the original image set pixels at invalid locations to 0 this function performs the hornschunck algorithm and returns the estimated optical flow it is assumed that the input images are grayscale and normalized to be in 0 1 parameters image0 first image of the sequence image1 second image of the sequence alpha regularization constant numiter number of iterations performed returns estimated horizontal vertical flow np roundhornschuncknp array0 0 2 0 0 2 np array0 2 0 0 2 0 alpha0 1 numiter110 astypenp int32 array 0 1 1 0 1 1 blankline 0 0 0 0 0 0 dtypeint32 initialize flow prepare kernels for the calculation of the derivatives and the average velocity iteratively refine the flow this updates the flow as proposed in the paper step 12 warps the pixels of an image into a new image using the horizontal and vertical flows pixels that are warped from an invalid location are set to 0 parameters image grayscale image horizontal_flow horizontal flow vertical_flow vertical flow returns warped image warp np array 0 1 2 0 3 0 2 2 2 np array 0 1 1 1 0 0 1 1 1 np array 0 0 0 0 1 0 0 0 1 array 0 0 0 3 1 0 0 2 3 create a grid of all pixel coordinates and subtract the flow to get the target pixels coordinates find the locations outside of the original image set pixels at invalid locations to 0 this function performs the horn schunck algorithm and returns the estimated optical flow it is assumed that the input images are grayscale and normalized to be in 0 1 parameters image0 first image of the sequence image1 second image of the sequence alpha regularization constant num_iter number of iterations performed returns estimated horizontal vertical flow np round horn_schunck np array 0 0 2 0 0 2 np array 0 2 0 0 2 0 alpha 0 1 num_iter 110 astype np int32 array 0 1 1 0 1 1 blankline 0 0 0 0 0 0 dtype int32 initialize flow prepare kernels for the calculation of the derivatives and the average velocity iteratively refine the flow this updates the flow as proposed in the paper step 12
from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve


def warp(
    image: np.ndarray, horizontal_flow: np.ndarray, vertical_flow: np.ndarray
) -> np.ndarray:
    
    flow = np.stack((horizontal_flow, vertical_flow), 2)

    
    
    grid = np.stack(
        np.meshgrid(np.arange(0, image.shape[1]), np.arange(0, image.shape[0])), 2
    )
    grid = np.round(grid - flow).astype(np.int32)

    
    invalid = (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]))
    grid[invalid] = 0

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

    
    warped[invalid[:, :, 0] | invalid[:, :, 1]] = 0

    return warped


def horn_schunck(
    image0: np.ndarray,
    image1: np.ndarray,
    num_iter: SupportsIndex,
    alpha: float | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    
    if alpha is None:
        alpha = 0.1

    
    horizontal_flow = np.zeros_like(image0)
    vertical_flow = np.zeros_like(image0)

    
    kernel_x = np.array([[-1, 1], [-1, 1]]) * 0.25
    kernel_y = np.array([[-1, -1], [1, 1]]) * 0.25
    kernel_t = np.array([[1, 1], [1, 1]]) * 0.25
    kernel_laplacian = np.array(
        [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]
    )

    
    for _ in range(num_iter):
        warped_image = warp(image0, horizontal_flow, vertical_flow)
        derivative_x = convolve(warped_image, kernel_x) + convolve(image1, kernel_x)
        derivative_y = convolve(warped_image, kernel_y) + convolve(image1, kernel_y)
        derivative_t = convolve(warped_image, kernel_t) + convolve(image1, -kernel_t)

        avg_horizontal_velocity = convolve(horizontal_flow, kernel_laplacian)
        avg_vertical_velocity = convolve(vertical_flow, kernel_laplacian)

        
        update = (
            derivative_x * avg_horizontal_velocity
            + derivative_y * avg_vertical_velocity
            + derivative_t
        )
        update = update / (alpha**2 + derivative_x**2 + derivative_y**2)

        horizontal_flow = avg_horizontal_velocity - derivative_x * update
        vertical_flow = avg_vertical_velocity - derivative_y * update

    return horizontal_flow, vertical_flow


if __name__ == "__main__":
    import doctest

    doctest.testmod()