torch-nvidia-of-sdk 5.0.0__py3-none-manylinux_2_34_x86_64.whl

of/__init__.py

@@ -0,0 +1,3 @@
+from . import io, visualization, metrics, datasets # noqa
+
+from .methods import TorchNVOpticalFlow # noqa

of/datasets.py

@@ -0,0 +1,229 @@
+"""
+Dataset loaders for optical flow benchmarks.
+"""
+
+import numpy as np
+from pathlib import Path
+from typing import List, Dict, Optional, Union
+from dataclasses import dataclass
+import imageio.v3 as imageio
+
+from .io import read_flo
+
+
+@dataclass
+class FlowSample:
+    """A single optical flow sample with metadata."""
+
+    current_frame: Path
+    reference_frame: Path
+    fw_flow: Optional[Path] = None
+    bw_flow: Optional[Path] = None
+    fw_valid_mask: Optional[Path] = None
+    bw_valid_mask: Optional[Path] = None
+    name: str = ""
+    metadata: Dict = None
+
+    def __post_init__(self):
+        if self.metadata is None:
+            self.metadata = {}
+
+
+class OpticalFlowDataset:
+    """Base class for optical flow datasets."""
+
+    def __init__(self, root_dir: Union[str, Path]):
+        self.root_dir = Path(root_dir)
+        self.samples: List[FlowSample] = []
+        self.setup()
+
+    def setup(self):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> FlowSample:
+        raise NotImplementedError("Subclasses should implement this method.")
+
+
+class CSEMDataset(OpticalFlowDataset):
+    """
+    project_root/
+    ├── images/
+    │   ├── frame_0001.png
+    │   ├── frame_0002.png
+    |   ├── ...
+    │   └── frame_N.png
+    ├── flow/
+    │   ├── forward/          # Flow from t to t+1
+    │   │   ├── flow_0001.flo
+    │   │   ├── flow_0002.flo
+    │   │   ├── ...
+    │   │   └── flow_<N-1>.flo
+    │   └── backward/         # Flow from t to t-1 (for occlusion checks)
+    │       ├── flow_0002.flo
+    │       ├── flow_0003.flo
+    │       ├── ...
+    │       └── flow_<N>.flo
+    """
+
+    def setup(self):
+        frame_paths = sorted((self.root_dir / "images").glob("*.png"))
+        for i in range(len(frame_paths) - 1):
+            current_frame = frame_paths[i]
+            reference_frame = frame_paths[i + 1]
+            name = current_frame.stem.split("_")[-1]
+            frame_number = int(name)
+            fw_flow_path = self.root_dir / "flow" / "forward" / f"flow_{name}.flo"
+            bw_flow_path = self.root_dir / "flow" / "backward" / f"flow_{name}.flo"
+            sample = FlowSample(
+                current_frame=current_frame,
+                reference_frame=reference_frame,
+                fw_flow=fw_flow_path if fw_flow_path.exists() else None,
+                bw_flow=bw_flow_path if bw_flow_path.exists() else None,
+                name=name,
+                metadata={"frame_number": frame_number},
+            )
+            self.samples.append(sample)
+
+    def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
+        sample = self.samples[idx]
+        data = {
+            "current_frame": imageio.imread(sample.current_frame),
+            "reference_frame": imageio.imread(sample.reference_frame),
+            "name": sample.name,
+            "metadata": sample.metadata,
+        }
+        if sample.fw_flow:
+            data["fw_flow"] = read_flo(sample.fw_flow)
+        if sample.bw_flow:
+            data["bw_flow"] = read_flo(sample.bw_flow)
+        return data
+
+
+class MPISintelDataset(OpticalFlowDataset):
+    """
+    MPI-Sintel optical flow dataset.
+
+    The dataset has two passes: 'clean' and 'final'.
+    Directory structure:
+        root/split/pass_name/sequence/frame_xxxx.png
+        root/split/flow/sequence/frame_xxxx.flo
+        root/split/invalid/sequence/frame_xxxx.png
+    """
+
+    def __init__(
+        self,
+        root_dir: Union[str, Path],
+        split: str = "training",
+        pass_name: str = "clean",
+    ):
+        """
+        Initialize MPI-Sintel dataset.
+
+        Args:
+            root_dir: Path to MPI-Sintel root directory
+            split: Either 'training' or 'test'
+            pass_name: Either 'clean' or 'final'
+        """
+        # Set attributes before super().__init__ because it calls setup()
+        self.split = split
+        self.pass_name = pass_name
+
+        super().__init__(root_dir)
+
+    def setup(self):
+        """Build list of samples from directory structure."""
+        images_dir = self.root_dir / self.split / self.pass_name
+
+        if not images_dir.exists():
+            raise ValueError(f"Images directory not found: {images_dir}")
+
+        flow_dir = self.root_dir / self.split / "flow"
+        # Sintel provides 'invalid' masks (occlusions + out of bounds)
+        # We map this to fw_valid_mask (loading logic should handle the inversion if needed)
+        invalid_dir = self.root_dir / self.split / "invalid"
+
+        # Iterate through sequences
+        for seq_dir in sorted(images_dir.iterdir()):
+            if not seq_dir.is_dir():
+                continue
+
+            seq_name = seq_dir.name
+            # Get all image pairs in the sequence
+            image_files = sorted(seq_dir.glob("*.png"))
+
+            for i in range(len(image_files) - 1):
+                img1 = image_files[i]
+                img2 = image_files[i + 1]
+
+                fw_flow_path = None
+                fw_valid_mask_path = None
+
+                # Training split has ground truth flow and masks
+                if self.split == "training":
+                    # Flow filename matches image filename
+                    cand_flow = flow_dir / seq_name / img1.name.replace(".png", ".flo")
+                    if cand_flow.exists():
+                        fw_flow_path = cand_flow
+
+                    cand_mask = invalid_dir / seq_name / img1.name
+                    if cand_mask.exists():
+                        fw_valid_mask_path = cand_mask
+
+                sample = FlowSample(
+                    current_frame=img1,
+                    reference_frame=img2,
+                    fw_flow=fw_flow_path,
+                    bw_flow=None,  # Sintel standard structure doesn't provide backward flow
+                    fw_valid_mask=fw_valid_mask_path,
+                    bw_valid_mask=None,
+                    name=f"{seq_name}_{img1.stem}",
+                    metadata={
+                        "sequence": seq_name,
+                        "pass": self.pass_name,
+                        "frame_index": i,
+                    },
+                )
+
+                self.samples.append(sample)
+
+    def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
+        sample = self.samples[idx]
+        data = {
+            "current_frame": imageio.imread(sample.current_frame),
+            "reference_frame": imageio.imread(sample.reference_frame),
+            "name": sample.name,
+            "metadata": sample.metadata,
+        }
+
+        if sample.fw_flow:
+            data["fw_flow"] = read_flo(sample.fw_flow)
+
+        if sample.fw_valid_mask:
+            # Sintel masks are PNGs where logic 1 usually means invalid.
+            # Reading as is; downstream transforms can invert/threshold.
+            data["fw_valid_mask"] = imageio.imread(sample.fw_valid_mask)
+
+        return data
+
+
+def get_dataset(name: str, root_dir: Union[str, Path], **kwargs) -> OpticalFlowDataset:
+    """
+    Factory function to get a dataset by name.
+
+    Args:
+        name: Dataset name ('mpi-sintel')
+        root_dir: Root directory of the dataset
+        **kwargs: Additional arguments for the dataset
+
+    Returns:
+        OpticalFlowDataset instance
+    """
+    name = name.lower()
+
+    if name == "mpi-sintel":
+        return MPISintelDataset(root_dir, **kwargs)
+    else:
+        raise ValueError(f"Unknown dataset: {name}. Only 'mpi-sintel' is supported.")

of/io.py

@@ -0,0 +1,80 @@
+"""
+Input/Output utilities for optical flow data.
+Supports reading and writing optical flow in various formats.
+"""
+
+import struct
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+from torch import Tensor
+
+FlowType = Union[NDArray, Tensor]
+
+
+def read_flo(filepath: Union[str, Path]) -> NDArray:  # noqa
+    """
+    Read optical flow from .flo file (Middlebury format).
+
+    Args:
+        filepath: Path to the .flo file
+
+    Returns:
+        Optical flow as numpy array of shape (H, W, 2) with dtype float32
+
+    Raises:
+        ValueError: If file format is invalid
+    """
+    filepath = Path(filepath)
+
+    with open(filepath, "rb") as f:
+        # Check magic number
+        magic = struct.unpack("f", f.read(4))[0]
+        if magic != 202021.25:
+            raise ValueError(f"Invalid .flo file format. Magic number: {magic}")
+
+        # Read dimensions
+        width = struct.unpack("i", f.read(4))[0]
+        height = struct.unpack("i", f.read(4))[0]
+
+        # Read flow data
+        data = np.fromfile(f, np.float32)
+
+    # Reshape to (height, width, 2)
+    flow = data.reshape((height, width, 2))
+    return flow
+
+
+def write_flo(filepath: Union[str, Path], flow: FlowType) -> None:
+    """
+    Write optical flow to .flo file (Middlebury format).
+
+    Args:
+        filepath: Path to save the .flo file
+        flow: Optical flow as numpy array of shape (H, W, 2)
+
+    Raises:
+        ValueError: If flow array shape is invalid
+    """
+    filepath = Path(filepath)
+    if isinstance(flow, torch.Tensor):
+        flow = flow.detach().cpu().numpy()
+
+    if flow.ndim != 3 or flow.shape[2] != 2:
+        raise ValueError(f"Flow must have shape (H, W, 2), got {flow.shape}")
+
+    height, width, _ = flow.shape
+
+    with open(filepath, "wb") as f:
+        # Write magic number
+        f.write(struct.pack("f", 202021.25))
+
+        # Write dimensions
+        f.write(struct.pack("i", width))
+        f.write(struct.pack("i", height))
+
+        # Write flow data
+        flow.astype(np.float32).tofile(f)

of/methods/__init__.py

@@ -0,0 +1,3 @@
+from .nvidia_sdk import TorchNVOpticalFlow
+
+__all__ = ["TorchNVOpticalFlow"]

of/methods/nvidia_sdk.py

@@ -0,0 +1,69 @@
+import torch
+import numpy as np
+from typing import Literal
+from .nvof_torch import TorchNVOpticalFlow as _C_TorchNVOpticalFlow
+
+
+class TorchNVOpticalFlow(_C_TorchNVOpticalFlow):
+    @classmethod
+    def from_tensor(
+        cls,
+        input: torch.Tensor,
+        preset: Literal["slow", "medium", "fast"] = "fast",
+        grid_size: Literal[1, 2, 4] = 1,
+    ):
+        """
+        Factory method to create engine from a tensor.
+
+        Args:
+            input: Tensor of shape (H, W, C) to infer dimensions from
+            preset: Performance preset ("slow", "medium", "fast")
+            grid_size: Grid size for optical flow (1, 2, or 4)
+            bidirectional: Whether to compute bidirectional flow # TODO
+
+        Returns:
+            TorchNVOpticalFlow instance
+        """
+        h = input.size(0)
+        w = input.size(1)
+        gpu = input.get_device()
+        return cls(w, h, gpu, preset, grid_size, False)
+
+    @torch.no_grad()
+    def compute_flow(
+        self,
+        input: torch.Tensor,
+        reference: torch.Tensor,
+        upsample: bool = True,
+    ) -> torch.Tensor:
+        """
+        Compute optical flow between input and reference frames.
+
+        Args:
+            input: Input tensor of shape (H, W, 3) or (H, W, 4), uint8, CUDA
+            reference: Reference tensor of shape (H, W, 3) or (H, W, 4), uint8, CUDA
+            upsample: Whether to upsample the output to full resolution
+
+        Returns:
+            Flow tensor of shape (H, W, 2), int16
+        """
+        # NVOpticalFlow requires ABGR uint8
+        if isinstance(input, np.ndarray):
+            input = torch.from_numpy(input)
+        if isinstance(reference, np.ndarray):
+            reference = torch.from_numpy(reference)
+
+        if input.shape[-1] == 3:
+            alpha_input = input.sum(dim=-1, keepdim=True).div(3).clamp(0, 255).byte()
+            input = torch.cat([alpha_input, input], dim=-1).to(f"cuda:{self.gpu_id()}")
+
+        if reference.shape[-1] == 3:
+            alpha_reference = (
+                reference.sum(dim=-1, keepdim=True).div(3).clamp(0, 255).byte()
+            )
+            reference = torch.cat([alpha_reference, reference], dim=-1).to(
+                f"cuda:{self.gpu_id()}"
+            )
+
+        # Checks performed in C++
+        return super().compute_flow(input, reference, upsample=upsample).float() / 32.0

of/metrics.py

@@ -0,0 +1,174 @@
+"""
+Error metrics for optical flow evaluation.
+"""
+
+import numpy as np
+from typing import Optional, Dict, Tuple
+
+
+def endpoint_error(flow_pred: np.ndarray, flow_gt: np.ndarray, 
+                   valid_mask: Optional[np.ndarray] = None) -> np.ndarray:
+    """
+    Compute End-Point Error (EPE) between predicted and ground truth flow.
+    
+    EPE = sqrt((u_pred - u_gt)^2 + (v_pred - v_gt)^2)
+    
+    Args:
+        flow_pred: Predicted flow of shape (H, W, 2)
+        flow_gt: Ground truth flow of shape (H, W, 2)
+        valid_mask: Optional boolean mask of valid pixels (H, W)
+        
+    Returns:
+        EPE map of shape (H, W)
+    """
+    diff = flow_pred - flow_gt
+    epe = np.sqrt(np.sum(diff**2, axis=2))
+    
+    if valid_mask is not None:
+        epe = epe * valid_mask
+    
+    return epe
+
+
+def average_endpoint_error(flow_pred: np.ndarray, flow_gt: np.ndarray,
+                           valid_mask: Optional[np.ndarray] = None) -> float:
+    """
+    Compute Average End-Point Error (AEE).
+    
+    Args:
+        flow_pred: Predicted flow of shape (H, W, 2)
+        flow_gt: Ground truth flow of shape (H, W, 2)
+        valid_mask: Optional boolean mask of valid pixels (H, W)
+        
+    Returns:
+        Average EPE as a scalar float
+    """
+    epe = endpoint_error(flow_pred, flow_gt, valid_mask)
+    
+    if valid_mask is not None:
+        return epe[valid_mask].mean()
+    else:
+        return epe.mean()
+
+
+def angular_error(flow_pred: np.ndarray, flow_gt: np.ndarray,
+                  valid_mask: Optional[np.ndarray] = None) -> np.ndarray:
+    """
+    Compute Angular Error (AE) between predicted and ground truth flow.
+    
+    The angular error is computed as the angle between the 3D vectors
+    (u, v, 1) for predicted and ground truth flows.
+    
+    Args:
+        flow_pred: Predicted flow of shape (H, W, 2)
+        flow_gt: Ground truth flow of shape (H, W, 2)
+        valid_mask: Optional boolean mask of valid pixels (H, W)
+        
+    Returns:
+        Angular error map in degrees of shape (H, W)
+    """
+    # Convert to 3D vectors
+    u_pred, v_pred = flow_pred[:, :, 0], flow_pred[:, :, 1]
+    u_gt, v_gt = flow_gt[:, :, 0], flow_gt[:, :, 1]
+    
+    # Compute dot product of normalized vectors
+    num = 1 + u_pred * u_gt + v_pred * v_gt
+    denom = np.sqrt(1 + u_pred**2 + v_pred**2) * np.sqrt(1 + u_gt**2 + v_gt**2)
+    
+    # Avoid division by zero
+    denom = np.maximum(denom, 1e-10)
+    
+    # Compute angle
+    cos_angle = np.clip(num / denom, -1.0, 1.0)
+    ae = np.arccos(cos_angle) * 180.0 / np.pi
+    
+    if valid_mask is not None:
+        ae = ae * valid_mask
+    
+    return ae
+
+
+def average_angular_error(flow_pred: np.ndarray, flow_gt: np.ndarray,
+                         valid_mask: Optional[np.ndarray] = None) -> float:
+    """
+    Compute Average Angular Error (AAE).
+    
+    Args:
+        flow_pred: Predicted flow of shape (H, W, 2)
+        flow_gt: Ground truth flow of shape (H, W, 2)
+        valid_mask: Optional boolean mask of valid pixels (H, W)
+        
+    Returns:
+        Average angular error in degrees as a scalar float
+    """
+    ae = angular_error(flow_pred, flow_gt, valid_mask)
+    
+    if valid_mask is not None:
+        return ae[valid_mask].mean()
+    else:
+        return ae.mean()
+
+
+def fl_all(flow_pred: np.ndarray, flow_gt: np.ndarray,
+           valid_mask: Optional[np.ndarray] = None,
+           threshold: float = 3.0) -> float:
+    """
+    Compute FL-all metric (percentage of outliers).
+    
+    An outlier is defined as EPE > threshold OR angular error > threshold degrees.
+    
+    Args:
+        flow_pred: Predicted flow of shape (H, W, 2)
+        flow_gt: Ground truth flow of shape (H, W, 2)
+        valid_mask: Optional boolean mask of valid pixels (H, W)
+        threshold: Threshold for outlier detection (default: 3.0 pixels)
+        
+    Returns:
+        Percentage of outliers (0-100)
+    """
+    epe = endpoint_error(flow_pred, flow_gt, valid_mask)
+    ae = angular_error(flow_pred, flow_gt, valid_mask)
+    
+    outliers = (epe > threshold) | (ae > threshold)
+    
+    if valid_mask is not None:
+        return 100.0 * outliers[valid_mask].sum() / valid_mask.sum()
+    else:
+        return 100.0 * outliers.sum() / outliers.size
+
+
+def compute_all_metrics(flow_pred: np.ndarray, flow_gt: np.ndarray,
+                       valid_mask: Optional[np.ndarray] = None,
+                       thresholds: Tuple[float, ...] = (1.0, 3.0, 5.0)) -> Dict[str, float]:
+    """
+    Compute all standard optical flow metrics.
+    
+    Args:
+        flow_pred: Predicted flow of shape (H, W, 2)
+        flow_gt: Ground truth flow of shape (H, W, 2)
+        valid_mask: Optional boolean mask of valid pixels (H, W)
+        thresholds: Thresholds for outlier metrics
+        
+    Returns:
+        Dictionary containing all metrics
+    """
+    metrics = {}
+    
+    # End-point error
+    metrics['EPE'] = average_endpoint_error(flow_pred, flow_gt, valid_mask)
+    
+    # Angular error
+    metrics['AE'] = average_angular_error(flow_pred, flow_gt, valid_mask)
+    
+    # Outlier percentages at different thresholds
+    for threshold in thresholds:
+        metrics[f'FL-{threshold}px'] = fl_all(flow_pred, flow_gt, valid_mask, threshold)
+    
+    # Root mean square error
+    epe = endpoint_error(flow_pred, flow_gt, valid_mask)
+    if valid_mask is not None:
+        metrics['RMSE'] = np.sqrt((epe[valid_mask]**2).mean())
+    else:
+        metrics['RMSE'] = np.sqrt((epe**2).mean())
+    
+    return metrics

of/visualization.py

@@ -0,0 +1,378 @@
+"""
+Visualization utilities for optical flow.
+
+INTERPRETATION GUIDE:
+---------------------
+To correctly read these visualizations, you must distinguish between the "Color Wheel"
+logic and "Masking" logic.
+
+1. PIXEL COLORS (Valid Flow):
+   - WHITE:        Zero Motion (Stationary).
+   - COLOR:        Motion Direction (Hue) and Speed (Saturation).
+                   (Red=Right, Green=Up, Cyan=Left, Blue=Down).
+   - BRIGHT/VIVID: Fast motion.
+   - PALE/PASTEL:  Slow motion.
+
+2. BLACK PIXELS (Invalid/Occluded):
+   - BLACK:        Invalid data, NaN, Occluded, or Unknown.
+
+   *Note on KITTI*: In KITTI ground truth, the sky and upper image are invalid.
+   They appear BLACK. The stationary road appears WHITE. If you see a visualization
+   that is Black where the road should be, the data is likely NaN/Invalid, not stationary.
+
+3. CONVENTIONS (Normalization):
+   - 'middlebury': Adapts to the image's max speed. Good for seeing details in
+                   slow scenes.
+   - 'kitti':      Uses fixed scaling (usually saturates at 3px or similar).
+                   Good for comparing speed between different clips (fast scenes
+                   look vivid, slow scenes look white/pale).
+"""
+
+from typing import List, Optional, Tuple, Union, Dict, Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+
+
+def _get_color_wheel() -> np.ndarray:
+    """
+    Generates the standard optical flow color wheel (Middlebury/KITTI standard).
+
+    Returns:
+        np.ndarray: Color wheel palette of shape (55, 3)
+    """
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros((ncols, 3))
+
+    col = 0
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY, 1) / RY)
+    col += RY
+
+    # YG
+    colorwheel[col : col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG, 1) / YG)
+    colorwheel[col : col + YG, 1] = 255
+    col += YG
+
+    # GC
+    colorwheel[col : col + GC, 1] = 255
+    colorwheel[col : col + GC, 2] = np.floor(255 * np.arange(0, GC, 1) / GC)
+    col += GC
+
+    # CB
+    colorwheel[col : col + CB, 1] = 255 - np.floor(255 * np.arange(0, CB, 1) / CB)
+    colorwheel[col : col + CB, 2] = 255
+    col += CB
+
+    # BM
+    colorwheel[col : col + BM, 2] = 255
+    colorwheel[col : col + BM, 0] = np.floor(255 * np.arange(0, BM, 1) / BM)
+    col += BM
+
+    # MR
+    colorwheel[col : col + MR, 2] = 255 - np.floor(255 * np.arange(0, MR, 1) / MR)
+    colorwheel[col : col + MR, 0] = 255
+
+    return colorwheel / 255.0
+
+
+def flow_to_color(
+    flow: np.ndarray, max_flow: Optional[float] = None, convention: str = "middlebury"
+) -> np.ndarray:
+    """
+    Convert optical flow to an RGB image using the standard color wheel.
+
+    Args:
+        flow: Optical flow array of shape (H, W, 2).
+
+        max_flow: Normalization factor.
+                  - If None: Defaults to the max magnitude in the flow array.
+                  - If float: Any flow larger than this is clamped/desaturated.
+
+        convention: Controls the default normalization behavior if max_flow is None.
+                   - 'middlebury': Always normalizes to the current image max.
+                   - 'kitti': Defaults to a fixed scale (e.g. 3.0) if max_flow is None.
+
+    Returns:
+        np.ndarray: RGB image of shape (H, W, 3).
+                    Valid flow is colored (0 motion = White).
+                    Invalid/NaN flow is Black.
+    """
+
+    assert flow.ndim == 3 and flow.shape[2] == 2
+
+    u = flow[:, :, 0]
+    v = flow[:, :, 1]
+
+    # 1. Detect Invalid Flow (NaN, Inf, or > 1e9)
+    # Middlebury uses 1e9 as a magic number for "unknown"
+    idx_unknown = (np.abs(u) > 1e9) | (np.abs(v) > 1e9) | np.isnan(u) | np.isnan(v)
+
+    # Temporarily clean data for calculation (avoid warnings)
+    u = np.where(idx_unknown, 0, u)
+    v = np.where(idx_unknown, 0, v)
+
+    mag = np.sqrt(u**2 + v**2)
+    angle = np.arctan2(-v, -u) / np.pi
+
+    fk = (angle + 1) / 2 * (55 - 1)
+    k0 = np.floor(fk).astype(int)
+    k1 = k0 + 1
+    k1[k1 == 55] = 0
+    f = fk - k0
+
+    colorwheel = _get_color_wheel()
+    col0 = colorwheel[k0]
+    col1 = colorwheel[k1]
+
+    col = (1 - f)[:, :, None] * col0 + f[:, :, None] * col1
+
+    if max_flow is None:
+        if convention.lower() == "kitti":
+            max_flow = 3.0
+        else:
+            max_flow = np.max(mag)
+            if max_flow == 0:
+                max_flow = 1.0
+
+    col *= mag[:, :, None] / max_flow
+
+    # Handle saturation
+    idx_sat = mag > max_flow
+    if np.any(idx_sat):
+        col[idx_sat] = col[idx_sat] * (max_flow / mag[idx_sat])[:, None]
+        col[idx_sat] = col[idx_sat] * 0.75 + 0.25
+
+    col = np.clip(col, 0, 1)
+
+    # 2. Apply Black for Unknown/Invalid pixels
+    if np.any(idx_unknown):
+        col[idx_unknown] = np.array([0, 0, 0])
+
+    return (col * 255).astype(np.uint8)
+
+
+def visualize_flow(
+    flow: np.ndarray,
+    ax: Optional[Axes] = None,
+    title: str = "Optical Flow",
+    max_flow: Optional[float] = None,
+    convention: str = "middlebury",
+    figsize: Tuple[int, int] = (8, 6),
+) -> Tuple[Figure, Axes]:
+    """
+    Visualize dense optical flow.
+    """
+    if ax is None:
+        fig, ax = plt.subplots(figsize=figsize)
+    else:
+        fig = ax.get_figure()
+
+    rgb = flow_to_color(flow, max_flow=max_flow, convention=convention)
+
+    ax.imshow(rgb)
+    ax.set_title(title)
+    ax.axis("off")
+
+    return fig, ax
+
+
+def visualize_flow_arrows(
+    flow: np.ndarray,
+    ax: Optional[Axes] = None,
+    image: Optional[np.ndarray] = None,
+    step: int = 16,
+    scale: float = 1.0,
+    color: str = "r",
+    title: str = "Flow Vectors",
+    figsize: Tuple[int, int] = (8, 6),
+) -> Tuple[Figure, Axes]:
+    """
+    Visualize flow as sparse arrows (quiver plot).
+    """
+    if ax is None:
+        fig, ax = plt.subplots(figsize=figsize)
+    else:
+        fig = ax.get_figure()
+
+    h, w = flow.shape[:2]
+
+    if image is not None:
+        ax.imshow(image, cmap="gray" if image.ndim == 2 else None)
+    else:
+        ax.imshow(np.zeros((h, w, 3), dtype=np.uint8))
+
+    y, x = np.mgrid[step // 2 : h : step, step // 2 : w : step]
+    u = flow[y, x, 0]
+    v = flow[y, x, 1]
+
+    ax.quiver(
+        x,
+        y,
+        u,
+        -v,
+        color=color,
+        angles="xy",
+        scale_units="xy",
+        scale=1 / scale,
+        width=0.003,
+    )
+
+    ax.set_title(title)
+    ax.set_xlim(0, w)
+    ax.set_ylim(h, 0)
+    ax.axis("off")
+
+    return fig, ax
+
+
+def compare_flows(
+    flow_pred: np.ndarray,
+    flow_gt: np.ndarray,
+    axes: Optional[Union[np.ndarray, List[Axes]]] = None,
+    max_flow: Optional[float] = None,
+    convention: str = "middlebury",
+    titles: Tuple[str, str, str] = ("Prediction", "Ground Truth", "EPE Error"),
+    figsize: Tuple[int, int] = (16, 5),
+) -> Tuple[Figure, np.ndarray]:
+    """
+    Compare prediction vs ground truth and visualize error.
+    """
+    if axes is None:
+        fig, axes = plt.subplots(1, 3, figsize=figsize)
+    else:
+        fig = axes[0].get_figure()
+        if len(axes) != 3:
+            raise ValueError("compare_flows requires 3 axes")
+
+    # 1. Visualize Prediction
+    visualize_flow(
+        flow_pred, ax=axes[0], title=titles[0], max_flow=max_flow, convention=convention
+    )
+
+    # 2. Visualize GT
+    visualize_flow(
+        flow_gt, ax=axes[1], title=titles[1], max_flow=max_flow, convention=convention
+    )
+
+    # 3. Visualize Endpoint Error (EPE)
+    epe = np.linalg.norm(flow_pred - flow_gt, axis=2)
+
+    # We use a valid mask to compute mean error only on valid pixels
+    # Otherwise NaN errors will break the visualization
+    mask = (np.abs(flow_gt[:, :, 0]) < 1e9) & (np.abs(flow_gt[:, :, 1]) < 1e9)
+    if mask.sum() > 0:
+        mean_epe = epe[mask].mean()
+    else:
+        mean_epe = 0.0
+
+    im_err = axes[2].imshow(epe, cmap="magma")
+    axes[2].set_title(f"{titles[2]} (Mean: {mean_epe:.2f})")
+    axes[2].axis("off")
+
+    plt.colorbar(im_err, ax=axes[2], fraction=0.046, pad=0.04)
+
+    plt.tight_layout()
+    return fig, axes
+
+
+def concatenate_flows(
+    arrays: List[np.ndarray],
+    positions: List[Tuple[int, int]],
+    kwargs_list: Optional[List[Dict[str, Any]]] = None,
+) -> np.ndarray:
+    """
+    Concatenates strictly typed images (RGB) and flow fields (Float) into a grid.
+
+    The function handles color space conversions internally so you can use standard
+    RGB colors for text and input images, while leveraging OpenCV for text rendering.
+
+    Args:
+        arrays: List of numpy arrays.
+            - 3 Channels: MUST be RGB uint8 (Standard Image).
+            - 2 Channels: MUST be Float Optical Flow.
+        positions: List of (col, row) tuples determining grid placement.
+        kwargs_list: Optional list of dictionaries (one per array) for customization.
+            Supported keys:
+
+            [Visual Params]
+            - 'caption': str         -> Text to write on the image.
+            - 'font_scale': float    -> Size of text (default: 1.0).
+            - 'font_thickness': int  -> Thickness of text (default: 2).
+            - 'font_color': tuple    -> RGB tuple (default: (255, 255, 255)).
+            - 'text_pos': tuple      -> (x, y) bottom-left position (default: (20, 40)).
+
+            [Flow Params (Only used if array has 2 channels)]
+            - 'max_flow': float      -> Normalization max magnitude.
+            - 'convention': str      -> 'middlebury' or 'kitti'.
+
+    Returns:
+        np.ndarray: RGB image (H_grid, W_grid, 3) ready for Matplotlib/TensorBoard.
+    """
+    import cv2
+
+    H, W = arrays[0].shape[:2]
+    N = len(arrays)
+
+    if kwargs_list is None:
+        kwargs_list = [{}] * N
+
+    # Canvas Size calculation
+    max_u = max(p[0] for p in positions)
+    max_v = max(p[1] for p in positions)
+
+    # Create internal BGR canvas for OpenCV operations
+    frame_bgr = np.zeros(((max_v + 1) * H, (max_u + 1) * W, 3), dtype=np.uint8)
+
+    for arr, (u, v), kw in zip(arrays, positions, kwargs_list):
+        # --- 1. Process Content (Strict Types -> BGR) ---
+        if arr.shape[2] == 2:
+            # TYPE: FLOW (Float) -> RGB -> BGR
+            flow_args = {k: v for k, v in kw.items() if k in ["max_flow", "convention"]}
+            rgb_flow = flow_to_color(arr, **flow_args)
+            img_chunk_bgr = cv2.cvtColor(rgb_flow, cv2.COLOR_RGB2BGR)
+
+        elif arr.shape[2] == 3:
+            # TYPE: IMAGE (RGB uint8) -> BGR
+            img_chunk_bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
+
+        else:
+            raise ValueError(
+                f"Array has {arr.shape[2]} channels. Expected 2 (Flow) or 3 (RGB Image)."
+            )
+
+        # --- 2. Add Text (OpenCV draws on BGR) ---
+        caption = kw.get("caption", None)
+        if caption:
+            # User specifies color in RGB (e.g., Red=(255,0,0))
+            # We flip to BGR because we are drawing on a BGR image
+            font_color_rgb = kw.get("font_color", (255, 255, 255))
+            font_color_bgr = font_color_rgb[::-1]
+
+            cv2.putText(
+                img_chunk_bgr,
+                str(caption),
+                kw.get("text_pos", (20, 40)),
+                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=kw.get("font_scale", 1.0),
+                color=font_color_bgr,
+                thickness=kw.get("font_thickness", 2),
+                lineType=cv2.LINE_AA,
+            )
+
+        # --- 3. Place in Grid ---
+        y, x = v * H, u * W
+        frame_bgr[y : y + H, x : x + W] = img_chunk_bgr
+
+    # Convert final result back to RGB for return
+    return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

torch_nvidia_of_sdk-5.0.0.dist-info/METADATA

@@ -0,0 +1,218 @@
+Metadata-Version: 2.2
+Name: torch-nvidia-of-sdk
+Version: 5.0.0
+Summary: PyTorch bindings for NVIDIA Optical Flow SDK, providing hardware-accelerated optical flow computation with PyTorch end-to-end integration in Nvidia and Python.
+Author-Email: Juan Montesinos <jfmontgar@gmail.com>
+Requires-Python: >=3.13
+Requires-Dist: imageio[ffmpeg]>=2.37.2
+Requires-Dist: matplotlib>=3.10.8
+Requires-Dist: numpy
+Requires-Dist: torch>=2.10.0
+Requires-Dist: tqdm>=4.67.2
+Provides-Extra: full
+Requires-Dist: opencv-contrib-python-headless>=4.13.0.90; extra == "full"
+Description-Content-Type: text/markdown
+
+# Torch Optical Flow
+
+PyTorch bindings for NVIDIA Optical Flow SDK, providing hardware-accelerated optical flow computation with PyTorch
+end-to-end integration in Nvidia and Python.
+
+Please read more about the NVIDIA Optical Flow SDK here: [https://developer.nvidia.com/optical-flow-sdk](https://developer.nvidia.com/optical-flow-sdk)
+
+# What's this repo about?
+- Hardware-accelerated optical flow using a special processor in Nvidia GPUs. No gradients are computed, this is for inference only.
+- Frame interpolation and ROI, or other additional content in the SDK is not supported.  
+- Configurable speed (slow, medium, fast) vs and grid size (1, 2, 4)
+- Support for various ABGR8 format, namely, RGB images.
+- End-to-end GPU processing with PyTorch.
+- Biderectional optical flow computation (forward and backward) in a single call. This is supported by the SDK, but not exposed in the wrappers they provide.
+
+The package comes with basic functionality for optical flow:
+- .flo reader and writer
+- Optical Flow common metrics
+- Visualization utilities
+
+
+# Requirements
+### System Requirements
+- NVIDIA GPU with Optical Flow SDK support (Turing, Ampere or Ada)
+- Tested on linux (Ubuntu), the SDK is compatible with windows too. Read optical_flow_sdk>Read_Me.pdf for windows instructions.
+
+### Software Requirements
+- CUDA toolkit >=10.2
+- Linux drivers "nvidia-smi" >=528.85
+- GCC >= 5.1
+- CMake >= 3.14
+- When you pip install torch, it comes with its own CUDA binaries. Get the same or higher CUDA toolkit version as your PyTorch installation.
+
+# Installation
+## pip / uv
+
+## Build from Source
+This repository uses uv. 
+A oneshot comand to build, install and test the package would be:
+
+```bash
+rm -rf build _skbuild .venv  && CC=gcc CXX=g++ uv sync --extra full --reinstall-package torch-nvidia-of-sdk && uv run examples/minimal_example.py
+```
+`--reinstall-package` forces `uv` to re-compile the package. Clearing caches is not really needed but I'm paranoid.
+`--extra full` is analogous to pip extras `pip install torch-nvidia-of-sdk[full]`. It just adds headless opencv for visualization
+## Compiling your own wheel
+`CC=gcc CXX=g++ uv build --wheel --package torch-nvidia-of-sdk` will build a wheel in `dist/` that you can install with pip.
+
+# Quick Start
+
+Try the minimal example to get started quickly:
+
+```bash
+# Run the minimal example (uses sample frames from assets/)
+uv run  examples/minimal_example.py
+```
+
+This will:
+1. Load two sample frames from the `assets/` directory
+2. Compute optical flow using NVOF
+3. Generate visualizations and save results to `output/`
+
+See [`examples/README.md`](examples/README.md) for more examples and tutorials.
+
+# Basic Usage
+
+```python
+import torch
+import numpy as np
+from of import TorchNVOpticalFlow
+from of.io import read_flo, write_flo
+from of.visualization import flow_to_color
+
+# Load your images (RGB format, uint8)
+img1 = torch.from_numpy(np.array(...)).cuda()  # Shape: (H, W, 3)
+img2 = torch.from_numpy(np.array(...)).cuda()
+
+# Initialize optical flow engine
+flow_engine = TorchNVOpticalFlow(
+    width=img1.shape[1],
+    height=img1.shape[0],
+    gpu_id=0,
+    preset="medium",  # "slow", "medium", or "fast"
+    grid_size=1,      # 1, 2, or 4
+)
+
+# Compute optical flow
+flow = flow_engine.compute_flow(img1, img2, upsample=True)
+
+# Flow is a (H, W, 2) tensor where flow[..., 0] is x-displacement, flow[..., 1] is y-displacement
+print(f"Flow shape: {flow.shape}")
+
+# Visualize flow as RGB image
+flow_rgb = flow_to_color(flow.cpu().numpy())
+
+# Save flow to .flo file
+write_flo("output_flow.flo", flow)
+```
+
+# API Reference
+
+## Core Class: `TorchNVOpticalFlow`
+
+### Constructor
+
+```python
+TorchNVOpticalFlow(
+    width: int,
+    height: int,
+    gpu_id: int = 0,
+    preset: str = "medium",
+    grid_size: int = 1,
+    bidirectional: bool = False
+)
+```
+
+**Parameters:**
+- `width`: Width of input images in pixels
+- `height`: Height of input images in pixels
+- `gpu_id`: CUDA device ID (default: 0)
+- `preset`: Speed/quality preset. Options:
+  - `"slow"`: Highest quality, slowest
+  - `"medium"`: Balanced (recommended)
+  - `"fast"`: Fastest, lower quality
+- `grid_size`: Output grid size. Options: 1, 2, or 4
+  - 1: Full resolution output (default)
+  - 2/4: Downsampled output (faster, use with `upsample=True` to restore resolution)
+- `bidirectional`: Enable bidirectional flow computation (forward and backward)
+
+### Methods
+
+#### `compute_flow(input, reference, upsample=True)`
+
+Compute forward optical flow between two frames.
+
+**Parameters:**
+- `input`: First frame as CUDA tensor of shape `(H, W, 4)`, dtype `uint8`, RGBA format
+- `reference`: Second frame as CUDA tensor of shape `(H, W, 4)`, dtype `uint8`, RGBA format
+- `upsample`: If True and grid_size > 1, upsample flow to full resolution (default: True)
+
+**Returns:**
+- `torch.Tensor`: Optical flow of shape `(H, W, 2)`, dtype `float32`
+  - `flow[..., 0]`: Horizontal displacement (x)
+  - `flow[..., 1]`: Vertical displacement (y)
+
+**Example:**
+```python
+flow = flow_engine.compute_flow(img1_rgba, img2_rgba, upsample=True)
+```
+
+#### `compute_flow_bidirectional(input, reference, upsample=True)`
+
+Compute both forward and backward optical flow.
+
+**Parameters:**
+- `input`: First frame as CUDA tensor of shape `(H, W, 4)`, dtype `uint8`, RGBA format
+- `reference`: Second frame as CUDA tensor of shape `(H, W, 4)`, dtype `uint8`, RGBA format
+- `upsample`: If True and grid_size > 1, upsample flows to full resolution (default: True)
+
+**Returns:**
+- `Tuple[torch.Tensor, torch.Tensor]`: Forward and backward flows, each of shape `(H, W, 2)`
+
+**Example:**
+```python
+forward_flow, backward_flow = flow_engine.compute_flow_bidirectional(
+    img1_rgba, img2_rgba, upsample=True
+)
+```
+
+#### `output_shape()`
+
+Get the output shape for the current configuration.
+
+**Returns:**
+- `List[int]`: Output shape as `[height, width, 2]`
+
+---
+
+## I/O Utilities (`of.io`)
+
+### `read_flo(filepath)`
+
+Read optical flow from `.flo` file (Middlebury format).
+
+**Parameters:**
+- `filepath`: Path to `.flo` file (str or Path)
+
+**Returns:**
+- `np.ndarray`: Flow array of shape `(H, W, 2)`, dtype `float32`
+
+### `write_flo(filepath, flow)`
+
+Write optical flow to `.flo` file (Middlebury format).
+
+**Parameters:**
+- `filepath`: Output file path (str or Path)
+- `flow`: Flow array of shape `(H, W, 2)` (numpy array or torch tensor)
+  
+# Examples
+
+This repository includes several examples in the `examples/` directory:
+
+See [`examples/README.md`](examples/README.md) for detailed documentation and usage instructions.

torch_nvidia_of_sdk-5.0.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+of/__init__.py,sha256=4UKSViAFdZWuiGfxIbactVbQkQeA7hnoGZd2_PpdveM,105
+of/datasets.py,sha256=Yf2l0vNEkzkEVo0Yn3dkV7RSPoTxos_WOIPlOHap5Nc,7629
+of/io.py,sha256=AXMlM2zh6zAt5l4hbwTsHq3lpXXGA_8kX3mtTWZR3hs,2075
+of/metrics.py,sha256=djeP2XkYt5FcGkQsdKhtynhSbBOLpemm9iz6I3Td3mg,5497
+of/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+of/visualization.py,sha256=Wc-d7-QvTzvKIBSGC0eE62qEI9R95E1MbispRM_2uvI,11621
+of/methods/__init__.py,sha256=wabPZD0BntWyXUcrAvbjY0hfMg1ArtHtD-bMTyvmgS8,77
+of/methods/nvidia_sdk.py,sha256=0nZgsHCZrJXiFsPn6RhVdhrPWoHKvrQr2Pru0uLyRW4,2363
+of/methods/nvof_torch.cpython-313-x86_64-linux-gnu.so,sha256=54urcZT1reYkrnfYxHK2nROZOa98otjsHk2A0R0Xuqw,313808
+torch_nvidia_of_sdk.libs/libcuda-df918354.so.580.105.08,sha256=b022zSCMvDgYt5zAe0-ijrqSxG2JFnJMfFzmm063juc,96295392
+torch_nvidia_of_sdk.libs/libcudart-381c0faa.so.12.9.37,sha256=aSNYYYNLW27faxYhvRRkFbGQDgOiSvsmnkYOVnP7oyc,753656
+torch_nvidia_of_sdk-5.0.0.dist-info/METADATA,sha256=I6OgY1-CAzrcjVdF1tRAVzxrhmVfeeo77f2qbTPxVPE,7104
+torch_nvidia_of_sdk-5.0.0.dist-info/WHEEL,sha256=JVXwG5HcbWr_fVpHgfrVgL0hYKvnp2fH_zuuklvZsD4,115
+torch_nvidia_of_sdk-5.0.0.dist-info/RECORD,,
+torch_nvidia_of_sdk-5.0.0.dist-info/sboms/auditwheel.cdx.json,sha256=3Gc-fQ9TGDEcEAu3DXWD50ff7jSXkLzkbsTPjFwfImY,1956

torch_nvidia_of_sdk-5.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: scikit-build-core 0.11.6
+Root-Is-Purelib: false
+Tag: py3-none-manylinux_2_34_x86_64
+

torch_nvidia_of_sdk-5.0.0.dist-info/sboms/auditwheel.cdx.json

@@ -0,0 +1 @@
+{"bomFormat": "CycloneDX", "specVersion": "1.4", "version": 1, "metadata": {"component": {"type": "library", "bom-ref": "pkg:pypi/torch_nvidia_of_sdk@5.0.0?file_name=torch_nvidia_of_sdk-5.0.0-py3-none-manylinux_2_34_x86_64.whl", "name": "torch_nvidia_of_sdk", "version": "5.0.0", "purl": "pkg:pypi/torch_nvidia_of_sdk@5.0.0?file_name=torch_nvidia_of_sdk-5.0.0-py3-none-manylinux_2_34_x86_64.whl"}, "tools": [{"name": "auditwheel", "version": "6.6.0"}]}, "components": [{"type": "library", "bom-ref": "pkg:pypi/torch_nvidia_of_sdk@5.0.0?file_name=torch_nvidia_of_sdk-5.0.0-py3-none-manylinux_2_34_x86_64.whl", "name": "torch_nvidia_of_sdk", "version": "5.0.0", "purl": "pkg:pypi/torch_nvidia_of_sdk@5.0.0?file_name=torch_nvidia_of_sdk-5.0.0-py3-none-manylinux_2_34_x86_64.whl"}, {"type": "library", "bom-ref": "pkg:deb/ubuntu/libnvidia-compute-580@580.105.08-0ubuntu1#4bd77591991c90dfe7081d8439e3bbfc934181610ce8d8a0ff77e83466d8866b", "name": "libnvidia-compute-580", "version": "580.105.08-0ubuntu1", "purl": "pkg:deb/ubuntu/libnvidia-compute-580@580.105.08-0ubuntu1"}, {"type": "library", "bom-ref": "pkg:deb/ubuntu/cuda-cudart-12-9@12.9.37-1#52366d11334f0bdf1ec18cb5b3389421d9afd26e46a0ea03841c7a6067e9ea15", "name": "cuda-cudart-12-9", "version": "12.9.37-1", "purl": "pkg:deb/ubuntu/cuda-cudart-12-9@12.9.37-1"}], "dependencies": [{"ref": "pkg:pypi/torch_nvidia_of_sdk@5.0.0?file_name=torch_nvidia_of_sdk-5.0.0-py3-none-manylinux_2_34_x86_64.whl", "dependsOn": ["pkg:deb/ubuntu/libnvidia-compute-580@580.105.08-0ubuntu1#4bd77591991c90dfe7081d8439e3bbfc934181610ce8d8a0ff77e83466d8866b", "pkg:deb/ubuntu/cuda-cudart-12-9@12.9.37-1#52366d11334f0bdf1ec18cb5b3389421d9afd26e46a0ea03841c7a6067e9ea15"]}, {"ref": "pkg:deb/ubuntu/libnvidia-compute-580@580.105.08-0ubuntu1#4bd77591991c90dfe7081d8439e3bbfc934181610ce8d8a0ff77e83466d8866b"}, {"ref": "pkg:deb/ubuntu/cuda-cudart-12-9@12.9.37-1#52366d11334f0bdf1ec18cb5b3389421d9afd26e46a0ea03841c7a6067e9ea15"}]}