AADFlowNet

Adaptive Angle-aware Dense Optical Flow Network

简体中文 · Documentation

AADFlowNet is a compact PyTorch network for dense two-dimensional optical flow estimation from a current BGR image and a reference BGR image. It is designed for vision-based tactile sensing, while keeping its model contract independent of a particular camera or device SDK.

The repository contains the model, dataset loader, four-phase training entry point, dataset and stream inference, and ONNX export. Pretrained weights and datasets are not distributed in this source release.

Inference preview

The animation below shows qualitative model inference with the tactile image and the predicted dense flow visualization.

Highlights

• Three-stage estimation: magnitude, angle-aware reconstruction, and residual refinement.
• Four-phase training with explicit parameter freezing for each phase.
• A compact model with 132,119 parameters in the default configuration.
• Strict input, output, dataset, and checkpoint validation.
• Dataset evaluation, camera/video-stream inference, and static-shape ONNX export.
• No dependency on a proprietary tactile-camera SDK.

Architecture

The model estimates flow in three stages:

1. Magnitude estimation predicts the displacement length.
2. Angle-aware reconstruction estimates direction and combines it with the magnitude to produce a dense base flow.
3. Residual refinement corrects local errors and returns the final flow.

Training follows four consecutive phases: magnitude training, angle-aware base flow training, residual refinement, and joint fine-tuning. The final phase keeps the magnitude branch fixed while fine-tuning the direction and refinement branches.

Repository layout

AADFlowNet/
├── model/
│   ├── aadflownet.py       # network definition
│   ├── dataset.py          # dataset validation and data loaders
│   └── runtime.py          # preprocessing and checkpoint loading
├── srcs/
│   ├── train.py            # four-phase training
│   ├── infer.py            # dataset and stream inference
│   └── export_onnx.py      # ONNX export and validation
├── docs/
│   ├── README.md
│   ├── README_zh-CN.md
│   ├── DATASET.md
│   ├── DATASET_zh-CN.md
│   ├── CHECKPOINTS.md
│   ├── CHECKPOINTS_zh-CN.md
│   └── media/
├── requirements.txt
└── LICENSE

Run all commands below from the repository root. The entry points are executed as Python modules so imports resolve consistently across platforms.

Requirements and installation

• Python 3.10 or newer
• PyTorch 2.6 or newer
• NumPy and OpenCV
• ONNX and ONNX Script for export

python -m pip install --upgrade pip
python -m pip install -r requirements.txt

requirements.txt installs the standard PyPI PyTorch build. If you need a specific CUDA or accelerator build, install PyTorch with the command generated by the official PyTorch selector, then install the remaining requirements.

Quick model check

The following code constructs the network without a checkpoint and verifies the public tensor contract:

import torch

from model import create_model

model = create_model().eval()
image_pair = torch.zeros(1, 6, 320, 320, dtype=torch.float32)

with torch.inference_mode():
    flow = model(image_pair)

print(flow.shape)  # torch.Size([1, 2, 320, 320])

This is a structural smoke test. Meaningful flow estimation requires trained AADFlowNet weights.

Dataset

The dataset root may contain one or more training and validation sequences. Training directories are named train or train__*; validation directories are named val or val__*.

dataset/
├── train/
│   ├── initial/frame.png
│   └── current/
│       ├── frames/000001.png
│       ├── flow_x/000001.npy
│       └── flow_y/000001.npy
└── val/
    ├── initial/frame.png
    └── current/
        ├── frames/000001.png
        ├── flow_x/000001.npy
        └── flow_y/000001.npy

Each flow file is a two-dimensional float32 NumPy array. Its stem must match the corresponding frame stem. See Dataset format for the complete validation, coordinate, and multi-sequence rules.

Training

The training entry point runs the four phases in order and writes checkpoints to the selected output directory:

python -m srcs.train \
  --data /path/to/dataset \
  --output runs/aadflownet

Inspect every training option with:

python -m srcs.train --help

The default schedule uses 20 magnitude epochs, 20 direction epochs, 30 refinement epochs, and 10 fine-tuning epochs. Tune these values for your data and hardware rather than treating them as benchmark settings.

Inference

Evaluate one labelled validation sequence and report mean endpoint error:

python -m srcs.infer dataset \
  --checkpoint /path/to/aadflownet.pth \
  --data /path/to/dataset/val

Run a camera or video stream. If --reference is omitted, the first frame is used as the reference:

python -m srcs.infer stream \
  --checkpoint /path/to/aadflownet.pth \
  --source 0 \
  --show

Use python -m srcs.infer --help and the subcommand help for batch size, workers, device selection, input size, output video, visualization density, and frame limits.

ONNX export

python -m srcs.export_onnx \
  --checkpoint /path/to/aadflownet.pth \
  --output aadflownet.onnx

The exporter uses the torch.export-based ONNX path, defaults to ONNX opset 18, and validates the saved model with onnx.checker. The exported model has a static batch size of one; the default tensor shape is 1 x 6 x 320 x 320.

Model contract

Property	Contract
Input	float32, shape B x 6 x H x W
Channel order	current BGR, then reference BGR
Image range	[0, 1]
Output	float32, shape B x 2 x H x W
Output channels	(flow_x, flow_y)
Flow unit	pixels at the model output resolution

Current and reference images must have the same size. Stream utilities resize both images consistently before constructing the six-channel input.

Checkpoints and safety

Training checkpoints contain a strict state_dict plus model and input metadata. The loader also accepts a raw state dictionary. Architecture or tensor-contract mismatches fail instead of silently skipping parameters.

PyTorch checkpoints must still be treated as untrusted files. Only load files from a source you trust and verify published checksums before use. See Checkpoint format and safety for the full contract and integrity-check commands.

Documentation

Topic	English	简体中文
Project guide	README	项目说明
Dataset format	DATASET.md	DATASET_zh-CN.md
Checkpoints	CHECKPOINTS.md	CHECKPOINTS_zh-CN.md

Scope and limitations

• This repository does not include pretrained weights or training data.
• The provided commands define the execution workflow, not a published accuracy or latency benchmark.
• Dataset quality, label convention, image scale, and sensor configuration directly affect the learned flow field.
• Validate exported models in the target inference runtime before deployment.

License

AADFlowNet is released under the MIT License. Pretrained weights, training data, and third-party device binaries are not distributed with this repository.