Skip to content

Commit ac59957

Browse files
egorchistovegorchistov
committed
Initial release
1 parent 25bb18a commit ac59957

42 files changed

Lines changed: 5671 additions & 0 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.gitignore‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
.idea/
2+
.vscode/
3+
.gradio/
4+
__pycache__/
5+
assets/
6+
ckpts/
7+
datasets/
8+
submissions/
9+
wandb/

‎LICENSE‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
BSD 3-Clause License
2+
3+
Copyright (c) 2024, Princeton Vision & Learning Lab
4+
Copyright (c) 2025, MSU Graphics & Media Lab
5+
6+
Redistribution and use in source and binary forms, with or without
7+
modification, are permitted provided that the following conditions are met:
8+
9+
1. Redistributions of source code must retain the above copyright notice, this
10+
list of conditions and the following disclaimer.
11+
12+
2. Redistributions in binary form must reproduce the above copyright notice,
13+
this list of conditions and the following disclaimer in the documentation
14+
and/or other materials provided with the distribution.
15+
16+
3. Neither the name of the copyright holder nor the names of its
17+
contributors may be used to endorse or promote products derived from
18+
this software without specific prior written permission.
19+
20+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

‎README.md‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
<p align="center">
2+
<h1 align="center">MEMFOF: High-Resolution Training for Memory-Efficient Multi-Frame Optical Flow Estimation</h1>
3+
<p align="center">
4+
<a href="https://github.com/VladBargin">Vladislav Bargatin</a>
5+
·
6+
<a href="http://github.com/egorchistov">Egor Chistov</a>
7+
·
8+
<a href="https://github.com/AlexanderYakovenko1">Alexander Yakovenko</a>
9+
·
10+
<a href="https://linkedin.com/in/dmitriyvatolin">Dmitriy Vatolin</a>
11+
</p>
12+
<h3 align="center">ICCV 2025</h3>
13+
<h3 align="center"><a href="https://arxiv.org/abs/2506.23151">📄 Paper</a> | <a href="https://msu-video-group.github.io/memfof">🌐 Project Page</a> | <a href="https://colab.research.google.com/github/msu-video-group/memfof/blob/dev/demo.ipynb">🚀 Colab</a> | <a href="https://huggingface.co/spaces/egorchistov/MEMFOF">🤗 Demo</a></h3>
14+
</p>
15+
16+
## 🛠️ Installation
17+
18+
Our code is developed with pytorch >= 2.5.0, CUDA >= 12.6 and python >= 3.10.
19+
20+
```shell
21+
git clone https://github.com/msu-video-group/memfof.git
22+
cd memfof
23+
pip3 install -r requirements.txt
24+
```
25+
26+
## 🚀 Demo
27+
28+
Given a video sequence, our code supports generating prediction results of optical flow.
29+
30+
> 🏞️ Prefer MEMFOF-Tartan-T-TSKH model for real-world videos — it is trained with higher diversity and robustness in mind.
31+
32+
Refer to [demo.ipynb](https://colab.research.google.com/github/msu-video-group/memfof/blob/dev/demo.ipynb) for examle usage or run the following command to host a [demo page](https://huggingface.co/spaces/egorchistov/MEMFOF).
33+
34+
```shell
35+
python3 demo.py
36+
```
37+
38+
## 📦 Models
39+
40+
- [`MEMFOF-Tartan`](https://huggingface.co/egorchistov/MEMFOF-Tartan)
41+
- [`MEMFOF-Tartan-T`](https://huggingface.co/egorchistov/MEMFOF-Tartan-T)
42+
- [`MEMFOF-Tartan-T-TSKH`](https://huggingface.co/egorchistov/MEMFOF-Tartan-T-TSKH) (✅ Recommended for real-world videos)
43+
- [`MEMFOF-Tartan-T-TSKH-kitti`](https://huggingface.co/egorchistov/MEMFOF-Tartan-T-TSKH-kitti)
44+
- [`MEMFOF-Tartan-T-TSKH-sintel`](https://huggingface.co/egorchistov/MEMFOF-Tartan-T-TSKH-sintel)
45+
- [`MEMFOF-Tartan-T-TSKH-spring`](https://huggingface.co/egorchistov/MEMFOF-Tartan-T-TSKH-spring)
46+
47+
## 🗂️ Datasets
48+
49+
To train MEMFOF, you will need to download the required datasets: [FlyingThings3D](https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html), [Sintel](http://sintel.is.tue.mpg.de/), [KITTI](http://www.cvlibs.net/datasets/kitti/eval_scene_flow.php?benchmark=flow), [HD1K](http://hci-benchmark.iwr.uni-heidelberg.de/), [TartanAir](https://theairlab.org/tartanair-dataset/), and [Spring](https://spring-benchmark.org/).
50+
51+
By default `datasets.py` will search for the datasets in these locations. You can create symbolic links to wherever the datasets were downloaded in the `datasets` folder.
52+
53+
```shell
54+
├── datasets
55+
├── Sintel
56+
├── KITTI
57+
├── FlyingThings3D
58+
├── HD1K
59+
├── Spring
60+
├── test
61+
├── train
62+
├── TartanAir
63+
```
64+
65+
## 📊 Evaluation and Submission
66+
67+
Please refer to [eval.sh](eval.sh) and [submission.sh](submission.sh) for more details.
68+
69+
## 🏋️ Training
70+
71+
Our training setup is configured for **4 nodes with 8 GPUs each**, using a fixed effective batch size.
72+
If you run the script on fewer resources, the per-device batch size may become too large and lead to **out-of-memory (OOM)** errors.
73+
74+
In such cases, you’ll need to manually lower the `effective_batch_size` in the config — **note that this will affect the final results**, as training dynamics and convergence may change.
75+
76+
Our training script is optimized for use with the slurm workload manager. A typical submission script looks like this:
77+
78+
```shell
79+
# (submit.sh)
80+
#!/bin/bash
81+
82+
#SBATCH --nodes=4
83+
#SBATCH --gres=gpu:8
84+
#SBATCH --ntasks-per-node=8
85+
#SBATCH --cpus-per-task=16
86+
87+
srun bash train.sh
88+
```
89+
90+
Alternatively, multi-node training is also supported via other launch methods, such as torchrun:
91+
92+
```shell
93+
OMP_NUM_THREADS=16 torchrun \
94+
--nproc_per_node=8 \
95+
--nnodes=4 \
96+
--node_rank <NODE_RANK> \
97+
--master_addr <MASTER_ADDR> \
98+
--master_port <MASTER_PORT> \
99+
--no-python bash train.sh
100+
```
101+
102+
For more details, refer to the [PyTorch Lightning documentation](https://lightning.ai/docs/pytorch/2.5.1/clouds/cluster.html).
103+
104+
We use Weights & Biases (WandB) for experiment tracking by default. To disable logging, set the environment variable:
105+
106+
```shell
107+
export WANDB_MODE=disabled
108+
```
109+
110+
## ❓ Need Help?
111+
112+
Feel free to open an issue if you have any questions.
113+
114+
## 📚 Citation
115+
116+
```
117+
@article{bargatin2025memfof,
118+
title={MEMFOF: High-Resolution Training for Memory-Efficient Multi-Frame Optical Flow Estimation},
119+
author={Bargatin, Vladislav and Chistov, Egor and Yakovenko, Alexander and Vatolin, Dmitriy},
120+
journal={arXiv preprint arXiv:2506.23151},
121+
year={2025}
122+
}
123+
```
124+
125+
## 🙏 Acknowledgements
126+
127+
This project relies on code from existing repositories: [SEA-RAFT](https://github.com/princeton-vl/SEA-RAFT), [VideoFlow](https://github.com/XiaoyuShi97/VideoFlow), and [GMA](https://github.com/zacjiang/GMA). We thank the original authors for their excellent work.

‎config/eval/kitti.json‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
{
2+
"name": "kitti",
3+
"dataset": "kitti",
4+
"val_datasets": ["kitti"],
5+
"monitor": null,
6+
7+
"use_var": true,
8+
"var_min": 0,
9+
"var_max": 10,
10+
"pretrain": "resnet34",
11+
"radius": 4,
12+
"dim": 512,
13+
"num_blocks": 2,
14+
"iters": 8,
15+
16+
"image_size": [750, 1920],
17+
"scale": 1,
18+
"effective_batch_size": 32,
19+
"num_workers": 16,
20+
"epsilon": 1e-8,
21+
"lr": 3e-5,
22+
"wdecay": 1e-5,
23+
"dropout": 0,
24+
"clip": 1.0,
25+
"gamma": 0.85,
26+
"num_steps": 2500,
27+
"val_steps": 250,
28+
29+
"restore_ckpt": "ckpts/Tartan-T-TSKH-kitti.ckpt",
30+
"coarse_config": null
31+
}

‎config/eval/sintel.json‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
{
2+
"name": "sintel",
3+
"dataset": "sintel",
4+
"val_datasets": ["sintel"],
5+
"monitor": null,
6+
7+
"use_var": true,
8+
"var_min": 0,
9+
"var_max": 10,
10+
"pretrain": "resnet34",
11+
"radius": 4,
12+
"dim": 512,
13+
"num_blocks": 2,
14+
"iters": 8,
15+
16+
"image_size": [872, 1920],
17+
"scale": 1,
18+
"effective_batch_size": 32,
19+
"num_workers": 16,
20+
"epsilon": 1e-8,
21+
"lr": 3e-5,
22+
"wdecay": 1e-5,
23+
"dropout": 0,
24+
"clip": 1.0,
25+
"gamma": 0.85,
26+
"num_steps": 12500,
27+
"val_steps": 1250,
28+
29+
"restore_ckpt": "ckpts/Tartan-T-TSKH-sintel.ckpt",
30+
"coarse_config": null
31+
}

‎config/eval/spring.json‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
{
2+
"name": "spring",
3+
"dataset": "spring",
4+
"val_datasets": ["spring"],
5+
"monitor": null,
6+
7+
"use_var": true,
8+
"var_min": 0,
9+
"var_max": 10,
10+
"pretrain": "resnet34",
11+
"radius": 4,
12+
"dim": 512,
13+
"num_blocks": 2,
14+
"iters": 8,
15+
16+
"image_size": [1080, 1920],
17+
"scale": 0,
18+
"effective_batch_size": 32,
19+
"num_workers": 16,
20+
"epsilon": 1e-8,
21+
"lr": 4.8e-5,
22+
"wdecay": 1e-5,
23+
"dropout": 0,
24+
"clip": 1.0,
25+
"gamma": 0.85,
26+
"num_steps": 60000,
27+
"val_steps": 3000,
28+
29+
"restore_ckpt": "ckpts/Tartan-T-TSKH-spring.ckpt",
30+
"coarse_config": null
31+
}

‎config/parser.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import json
2+
import argparse
3+
4+
5+
def json_to_args(json_path: str) -> argparse.Namespace:
6+
"""Convert JSON configuration file to argparse.Namespace object.
7+
8+
Parameters
9+
----------
10+
json_path : str
11+
Path to the JSON configuration file.
12+
13+
Returns
14+
-------
15+
argparse.Namespace
16+
Namespace object containing configuration parameters.
17+
"""
18+
args = argparse.Namespace()
19+
20+
with open(json_path, "r") as f:
21+
data = json.load(f)
22+
args.__dict__.update(data)
23+
return args
24+
25+
26+
def parse_args(parser: argparse.ArgumentParser) -> argparse.Namespace:
27+
"""Parse command line arguments and merge them with JSON configuration.
28+
29+
This function first parses command line arguments, then loads configuration
30+
from a JSON file specified by the 'cfg' argument, and finally merges both
31+
configurations with command line arguments taking precedence.
32+
33+
Parameters
34+
----------
35+
parser : argparse.ArgumentParser
36+
Argument parser instance with defined arguments.
37+
38+
Returns
39+
-------
40+
argparse.Namespace
41+
Merged configuration from both command line and JSON file.
42+
"""
43+
entry = parser.parse_args()
44+
args = json_to_args(entry.cfg)
45+
args.__dict__.update(vars(entry))
46+
return args

0 commit comments

Comments
 (0)