# Official implementation of GOCor


This is the official implementation of our paper : 

[**GOCor: Bringing Globally Optimized Correspondence Volumes into Your Neural Network.**](https://arxiv.org/abs/2009.07823)
<br />
Authors: [Prune Truong](https://prunetruong.com/) *, [Martin Danelljan](https://martin-danelljan.github.io/) *, Luc Van Gool, Radu Timofte<br />

\[[Paper](https://arxiv.org/abs/2009.07823)\]\[[Website](https://prunetruong.com/research/gocor)\]\[[Video](https://www.youtube.com/watch?v=V22MyFChBCs)\]



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. 


![alt text](/images/corr_diff_iteration.jpg)



Also check out our related work [**GLU-Net**](https://arxiv.org/abs/1912.05524) and the code
[here](https://github.com/PruneTruong/GLU-Net) !

<br />

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](https://github.com/PruneTruong/PDCNet).

<br />

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 :
```bash
@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 <a name="Installation"></a>

Inference runs for torch >= 1.0. 


* Create and activate conda environment with Python 3.x

```bash
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:
```bash
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. 

```bash
pip install cupy-cuda90 --no-cache-dir 
```


* **Download an archive with pre-trained models [click](https://drive.google.com/file/d/1DKoIhE80xW7y3phbLDPfJiYYFo3Vf8am/view?usp=sharing) and extract it to the project folder**                                                


# 2. Models

Pre-trained weights can be downloaded from [here](https://drive.google.com/file/d/1DKoIhE80xW7y3phbLDPfJiYYFo3Vf8am/view?usp=sharing).
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](https://github.com/PruneTruong/GLU-Net)
* **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](https://github.com/NVlabs/PWC-Net/tree/master/PyTorch), 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:

```bash
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/)
```
![alt text](/images/eth3d_warped_query_image_GLUNet_GOCor_dynamic.png)

For baseline GLU-Net, the output is instead:

```bash
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/

```
![alt text](/images/eth3d_warped_query_image_GLUNet_dynamic.png)



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


```bash
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/
```

![alt text](/images/kitti2015_warped_query_image_PWCNet_GOCor_chairs_things.png)


```bash
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/
```
![alt text](/images/kitti2015_warped_query_image_PWCNet_chairs_things.png)

<br />


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 <a name="Acknowledgement"></a>

We borrow code from public projects, such as [pytracking](https://github.com/visionml/pytracking), [GLU-Net](https://github.com/PruneTruong/GLU-Net), 
[DGC-Net](https://github.com/AaltoVision/DGC-Net), [PWC-Net](https://github.com/NVlabs/PWC-Net), 
[NC-Net](https://github.com/ignacio-rocco/ncnet), [Flow-Net-Pytorch](https://github.com/ClementPinard/FlowNetPytorch), 
[RAFT](https://github.com/princeton-vl/RAFT) ...

