Official implementation of GOCor

This is the official implementation of our paper :

GOCor: Bringing Globally Optimized Correspondence Volumes into Your Neural Network.
Authors: Prune Truong *, Martin Danelljan *, Luc Van Gool, Radu Timofte

[Paper][Website][Video]

The feature correlation layer serves as a key neural network module in numerous computer vision problems that involve dense correspondences between image pairs. It predicts a correspondence volume by evaluating dense scalar products between feature vectors extracted from pairs of locations in two images. However, this point-to-point feature comparison is insufficient when disambiguating multiple similar regions in an image, severely affecting the performance of the end task. This work proposes GOCor, a fully differentiable dense matching module, acting as a direct replacement to the feature correlation layer. The correspondence volume generated by our module is the result of an internal optimization procedure that explicitly accounts for similar regions in the scene. Moreover, our approach is capable of effectively learning spatial matching priors to resolve further matching ambiguities.

Also check out our related work GLU-Net and the code here !

In this repo, we only provide code to test on image pairs as well as the pre-trained weights of the networks evaluated in GOCor paper. We will not release the training code. However, since GOCor module is a plug-in replacement for the feature correlation layer, it can be integrated into any architecture and trained using the original training code. We will release general training and evaluation code in a general dense correspondence repo, coming soon here.

For any questions, issues or recommendations, please contact Prune at prune.truong@vision.ee.ethz.ch

Citation

If our project is helpful for your research, please consider citing :

@inproceedings{GOCor_Truong_2020,
      title = {{GOCor}: Bringing Globally Optimized Correspondence Volumes into Your Neural Network},
      author    = {Prune Truong 
                   and Martin Danelljan 
                   and Luc Van Gool 
                   and Radu Timofte},
      year = {2020},
      booktitle = {Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information
                   Processing Systems 2020, {NeurIPS} 2020}
}

1. Installation

Inference runs for torch >= 1.0.

Create and activate conda environment with Python 3.x

conda create -n GOCor_env python=3.7
conda activate GOCor_env

Install all dependencies (except for cupy, see below) by running the following command:

pip install torch torchvision imageio matplotlib numpy opencv-python pandas

Note: CUDA is required to run the code. Indeed, the correlation layer is implemented in CUDA using CuPy, which is why CuPy is a required dependency. It can be installed using pip install cupy or alternatively using one of the provided binary packages as outlined in the CuPy repository. The code was developed using Python 3.7 & PyTorch 1.0 & CUDA 9.0, which is why I installed cupy for cuda90. For another CUDA version, change accordingly.

pip install cupy-cuda90==7.8.0 --no-cache-dir

There are some issues with latest versions of cupy. So for all cuda, install cupy version 7.8.0. For example, on cuda10,

pip install cupy-cuda100==7.8.0 --no-cache-dir

Download an archive with pre-trained models click and extract it to the project folder

2. Models

Pre-trained weights can be downloaded from here. We provide the pre-trained weights of:

GLU-Net trained on the static data, these are given for reference, they correspond to the weights 'GLUNet_DPED_CityScape_ADE.pth' that we provided here
GLU-Net-GOCor trained on the static data, corresponds to network in the GOCor paper
GLU-Net trained on the dynamic data
GLU-Net-GOCor trained on the dynamic data, corresponds to network in the GOCor paper
PWC-Net finetuned on chairs-things (by us), they are given for reference
PWC-Net-GOCor finetuned on chair-things, corresponds to network in the GOCor paper
PWC-Net further finetuned on sintel (by us), for reference
PWC-Net-GOCor further finetuned on sintel, corresponds to network in the GOCor paper

For reference, you can also use the weights from the original PWC-Net repo, where the networks are trained on chairs-things and further finetuned on sintel. As explained in the paper, for training our PWC-Net-based models, we initialize the network parameters with the pre-trained weights trained on chairs-things.

All networks are created in 'model_selection.py'

3. Test on your own images

You can test the networks on a pair of images using test_models.py and the provided trained model weights. You must first choose the model and pre-trained weights to use. The inputs are the paths to the query and reference images. The images are then passed to the network which outputs the corresponding flow field relating the reference to the query image. The query is then warped according to the estimated flow, and a figure is saved.

For this pair of images (provided to check that the code is working properly) and using GLU-Net-GOCor trained on the dynamic dataset, the output is:

python test_models.py --model GLUNet_GOCor --pre_trained_model dynamic --path_query_image images/eth3d_query.png --path_reference_image images/eth3d_reference.png --write_dir evaluation/

additional optional arguments:
--pre_trained_models_dir (default is pre_trained_models/)

For baseline GLU-Net, the output is instead:

python test_models.py --model GLUNet --pre_trained_model dynamic --path_query_image images/eth3d_query.png --path_reference_image images/eth3d_reference.png --write_dir evaluation/

And for PWC-Net-GOCor and baseline PWC-Net:

python test_models.py --model PWCNet_GOCor --pre_trained_model chairs_things --path_query_image images/kitti2015_query.png --path_reference_image images/kitti2015_reference.png --write_dir evaluation/

python test_models.py --model PWCNet --pre_trained_model chairs_things --path_query_image images/kitti2015_query.png --path_reference_image images/kitti2015_reference.png --write_dir evaluation/

Possible model choices are : GLUNet, GLUNet_GOCor, PWCNet, PWCNet_GOCor

Possible pre-trained model choices are: static, dynamic, chairs_things, chairs_things_ft_sintel

4. Acknowledgement