Awesome
NeuFlow_v2
Official PyTorch implementation of paper:
NeuFlow v2: Push High-Efficiency Optical Flow To the Limit
Authors: Zhiyong Zhang, Aniket Gupta, Huaizu Jiang, Hanumant Singh
Installation (PyTorch >= 2.0 is required)
conda create --name neuflow python==3.8
conda activate neuflow
conda install pytorch==2.0.1 torchvision==0.15.2 pytorch-cuda=11.7 -c pytorch -c nvidia
pip install numpy opencv-python
Inference with HuggingFace š¤
Install huggingface-hub
pip install huggingface-hub
Inference uses our pretrained model, trained with multiple datasets, neuflow_mixed.pth:
python infer_hf.py
<img src="example_result.jpg" width="400" >
Inference
Inference uses our pretrained model, trained with multiple datasets, neuflow_mixed.pth:
python infer.py
Datasets
The datasets used to train and evaluate NeuFlow are as follows:
By default the dataloader assumes the datasets are located in folder datasets
and are organized as follows:
datasets
āāā FlyingChairs_release
āĀ Ā āāā data
āāā FlyingThings3D
āĀ Ā āāā frames_cleanpass
āĀ Ā āāā frames_finalpass
āĀ Ā āāā optical_flow
āāā HD1K
āĀ Ā āāā hd1k_challenge
āĀ Ā āāā hd1k_flow_gt
āĀ Ā āāā hd1k_flow_uncertainty
āĀ Ā āāā hd1k_input
āāā KITTI
āĀ Ā āāā testing
āĀ Ā āāā training
āāā Sintel
āĀ Ā āāā test
āĀ Ā āāā training
Symlink your dataset root to datasets
:
ln -s $YOUR_DATASET_ROOT datasets
Convert all your images and flows to .npy format to speed up data loading. This script provides an example of converting FlyingThings cleanpass data.
python images_flows_to_npy.py
Training
Simple training script:
python train.py \
--checkpoint_dir $YOUR_CHECKPOINT_DIR \
--stage things \
--val_dataset things sintel kitti \
--batch_size 32 \
--num_workers 4 \
--lr 1e-4 \
--val_freq 1000 \
--resume neuflow_things.pth \
--strict_resume
We trained on the FlyingThings dataset using 8x A5000 GPUs with the following command:
python -m torch.distributed.launch --nproc_per_node=8 --master_port=29501 train.py \
--checkpoint_dir $YOUR_CHECKPOINT_DIR \
--stage things \
--val_dataset things sintel kitti \
--batch_size 256 \
--num_workers 8 \
--lr 8e-4 \
--val_freq 500 \
--distributed
Evaluation
python eval.py \
--resume neuflow_things.pth