Optical Flow EPE with Masks (OmniWorld-style metric)

Problem

Compute the Endpoint Error (EPE) between a predicted optical flow field and a ground-truth flow field. Optionally apply a binary validity mask so that only valid pixels contribute to the mean EPE.

Solution

import numpy as np

def optical_flow_epe(pred_flow: np.ndarray, gt_flow: np.ndarray,
                     mask: np.ndarray = None) -> float:
    # pred_flow, gt_flow: (H, W, 2)
    diff = pred_flow - gt_flow
    epe = np.sqrt(np.sum(diff ** 2, axis=-1))  # (H, W)
    if mask is not None:
        mask = mask.astype(bool)
        if mask.sum() == 0:
            return 0.0
        return float(np.mean(epe[mask]))
    return float(np.mean(epe))

Explanation

1. Compute the per-pixel difference between predicted and ground-truth flow vectors.
2. Calculate the L2 norm (Euclidean distance) across the flow channels for each pixel, giving the endpoint error per pixel.
3. If a validity mask is provided, select only those pixels where the mask is True and average their EPE values.
4. Otherwise, average over all pixels.

Complexity

• Time: O(H * W) where H and W are the spatial dimensions of the flow field
• Space: O(H * W) for the EPE map