PyTorch Lightning DataModule for MegaFlow training and validation. Parameters ---------- args : argparse.Namespace Configuration

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

MegaFlow: Zero-Shot Large Displacement Optical Flow. Args: img_size (int): Image size in pixels. patch_size (int): Size of e

PyTorch Lightning module for MegaFlow optical flow estimation. This class implements the training and validation logic for the MegaFlow model

PyTorch Lightning module for MegaFlow point tracking estimation. This class implements the training and validation logic for the MegaFlow model

Residual convolution module.

r""" ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N

PyTorch Lightning DataModule for Megaflow training and validation. Parameters ---------- args : argparse.Namespace Configuration

Feature fusion block.

Flow head that processes input tokens and predicts downsample features for flow estimation. Behavior: In forward, it will take in a list

r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_

2D image to patch embedding: (B,C,H,W) -> (B,N,D) Args: img_size: Image size. patch_size: Patch token size. in_chans

Generates and caches 2D spatial positions for patches in a grid. This class efficiently manages the generation of spatial coordinates for patches

ResNet18, output resolution is 1/8. Each block has 2 layers.

2D Rotary Position Embedding implementation. This module applies rotary position embeddings to input tokens based on their 2D spatial positio

flow propagation with self-attention on feature query: feature0, key: feature0, value: flow

Pads images such that dimensions are divisible by 8

Dataset class for Spring optical flow dataset. For train, this dataset returns image1, image2, flow and a data tuple (framenum, scene name, l