GMFlow

PyTorch官方实现，对应论文：

GMFlow: 通过全局匹配学习光流，CVPR 2022，口头报告

作者：Haofei Xu、Jing Zhang、Jianfei Cai、Hamid Rezatofighi、Dacheng Tao

2022年11月15日更新：请查看我们的新工作：统一光流、立体视觉和深度估计及代码：unimatch，用于将GMFlow扩展到立体视觉和深度任务。我们还发布了更多具有不同速度-精度权衡的预训练GMFlow模型，详情请见：MODEL_ZOO.md。快来体验我们的Colab和HuggingFace演示，在浏览器中玩转GMFlow吧！

GMFlow的视频介绍（中文）已在bilibili上线！

https://user-images.githubusercontent.com/19343475/174446408-520b8a6c-9714-4ff3-978c-98e23ab29c1f.mp4

我们通过将光流问题重新表述为全局匹配问题，简化了光流估计的整个流程。

亮点

• 灵活且模块化的设计

  我们将端到端的光流框架分解为五个组件：特征提取、特征增强、特征匹配、光流传播和光流精炼。用户可以通过组合不同的组件轻松构建自定义的光流模型。

• 高精度

  仅需一次精炼，GMFlow在极具挑战性的Sintel基准测试上就超越了需要31次精炼的RAFT。

• 高效性

  基础GMFlow模型（无精炼）处理Sintel数据（436x1024）时，运行时间仅为57毫秒（V100）或26毫秒（A100）。由于GMFlow不需要大量顺序计算，因此在高端GPU（如A100）上比RAFT获得更高的加速效果。此外，GMFlow还简化了反向光流的计算，无需两次前向传播网络即可得到双向光流。利用正反向一致性检查，双向光流还可用于遮挡检测。

  

安装

我们的代码基于PyTorch 1.9.0、CUDA 10.2和Python 3.8。更高版本的PyTorch也应该可以正常运行。

我们建议使用conda进行安装：

conda env create -f environment.yml
conda activate gmflow

演示

所有预训练模型均可从Google Drive下载。

您可以在一系列图像上运行训练好的模型并可视化结果：

CUDA_VISIBLE_DEVICES=0 python main.py \
--inference_dir demo/sintel_market_1 \
--output_path output/gmflow-norefine-sintel_market_1 \
--resume pretrained/gmflow_sintel-0c07dcb3.pth

您还可以通过启用--pred_bidir_flow来预测双向光流，并使用--fwd_bwd_consistency_check进行正反向一致性检查。更多示例可在scripts/demo.sh中找到。

数据集

用于训练和评估GMFlow的数据集如下：

• FlyingChairs
• FlyingThings3D
• Sintel
• KITTI
• HD1K

默认情况下，数据加载器datasets.py假设数据集位于datasets文件夹中，并按以下方式组织：

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

建议将您的数据集根目录符号链接到datasets：

ln -s $YOUR_DATASET_ROOT datasets

否则，您可能需要修改datasets.py中的相应路径。

评估

您可以通过以下命令评估训练好的GMFlow模型：

CUDA_VISIBLE_DEVICES=0 python main.py --eval --val_dataset things sintel --resume pretrained/gmflow_things-e9887eda.pth 

更多评估脚本可在scripts/evaluate.sh中找到。

对于提交至Sintel和KITTI在线测试集的任务，您可以运行scripts/submission.sh。

训练

关于FlyingChairs、FlyingThings3D、Sintel和KITTI数据集的所有训练脚本，均收录于scripts/train_gmflow.sh和scripts/train_gmflow_with_refine.sh中。请注意，基础GMFlow模型（无精炼）可在4张16GB显存的V100 GPU上训练。而若要训练带有精炼的GMFlow，则默认需要8张16GB显存的V100 GPU、4张32GB显存的V100 GPU，或4张40GB显存的A100 GPU。您可能需要根据硬件情况调整批量大小和训练轮数。

我们支持使用TensorBoard监控和可视化训练过程。您可以先启动一个TensorBoard会话：

tensorboard --logdir checkpoints

然后在浏览器中访问http://localhost:6006。

引用

如果您在研究中使用了我们的工作，请考虑引用我们的论文：

@inproceedings{xu2022gmflow,
  title={GMFlow: Learning Optical Flow via Global Matching},
  author={Xu, Haofei and Zhang, Jing and Cai, Jianfei and Rezatofighi, Hamid and Tao, Dacheng},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={8121-8130},
  year={2022}
}

致谢

没有以下几个优秀的开源项目，本项目将无法实现：RAFT、LoFTR、DETR、Swin、mmdetection以及Detectron2。我们感谢原作者们的杰出贡献。

GMFlow 快速上手指南

GMFlow 是一个基于全局匹配（Global Matching）范式的光流估计模型，由 CVPR 2022 Oral 论文提出。该模型具有高精度、高效率的特点，无需多次迭代细化即可在 Sintel 等基准测试中超越传统方法。

环境准备

本项目基于 PyTorch 1.9.0、CUDA 10.2 和 Python 3.8 开发，更高版本的 PyTorch 通常也能兼容。

前置依赖：

• NVIDIA GPU（推荐 V100 或 A100 以获得最佳性能）
• Conda 包管理器（推荐）
• CUDA Toolkit（版本需与 PyTorch 匹配）

安装步骤

推荐使用 Conda 创建隔离环境并安装依赖：

# 克隆仓库（如果尚未克隆）
git clone https://github.com/autonomousvision/gmflow.git
cd gmflow

# 创建并激活虚拟环境
conda env create -f environment.yml
conda activate gmflow

提示：如果 environment.yml 下载缓慢，可手动创建环境后，使用国内镜像源（如清华源）安装核心依赖： pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 下载预训练模型

所有预训练模型可通过 Google Drive 下载。将下载的权重文件（例如 gmflow_sintel-0c07dcb3.pth）放置在 pretrained/ 目录下。

2. 运行推理示例

以下命令演示了如何对一组图像进行光流估计并可视化结果：

CUDA_VISIBLE_DEVICES=0 python main.py \
--inference_dir demo/sintel_market_1 \
--output_path output/gmflow-norefine-sintel_market_1 \
--resume pretrained/gmflow_sintel-0c07dcb3.pth

3. 双向光流与遮挡检测（可选）

GMFlow 支持高效计算双向光流并进行前后向一致性检查，用于遮挡区域检测：

CUDA_VISIBLE_DEVICES=0 python main.py \
--inference_dir demo/sintel_market_1 \
--output_path output/gmflow-bidir-sintel_market_1 \
--resume pretrained/gmflow_sintel-0c07dcb3.pth \
--pred_bidir_flow \
--fwd_bwd_consistency_check

运行完成后，结果图像将保存至指定的 output_path 目录中。