initial commit

Author: az-ja

.gitignore

@@ -0,0 +1,18 @@
+correlation.egg-info
+checkpoints/*
+data/
+images/*
+runs/*
+.vscode/*
+models/*
+*.png
+*.csv
+*.pyc
+
+*.sh
+*.zip
+job_*.txt
+eval_all.log
+alt_cuda_corr/build/
+alt_cuda_corr/dist/
+*.flo

LICENSE

@@ -0,0 +1,29 @@
+Copyright (c) 2023 Azin Jahedi. All Rights Reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+ 1. Redistribution of source code must retain the above copyright notice, this
+    list of conditions and the following disclaimer.
+
+ 2. Redistribution in binary form must reproduce the above copyright notice,
+    this list of conditions and the following disclaimer in the documentation
+    and/or other materials provided with the distribution.
+
+ 3. Neither the name of the copyright holder nor the names of its contributors
+    may be used to endorse or promote products derived from this software without
+    specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
+INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
+OF THE POSSIBILITY OF SUCH DAMAGE.
+
+You acknowledge that this software is not designed, licensed or intended for use
+in the design, construction, operation or maintenance of any military facility.

README.md

@@ -0,0 +1,72 @@
+# CCMR
+This is the inference code of our `CCMR` optical flow estimation method.
+
+ **[CCMR: High Resolution Optical Flow Estimation via Coarse-to-Fine Context-Guided Motion Reasoning](https://openaccess.thecvf.com/content/WACV2024/papers/Jahedi_CCMR_High_Resolution_Optical_Flow_Estimation_via_Coarse-To-Fine_Context-Guided_Motion_WACV_2024_paper.pdf)**<br/>
+> _WACV 2024_ <br/>
+> Azin Jahedi, Maximilian Luz, Marc Rivinius and Andrés Bruhn
+
+## Requirements
+
+The code has been tested with PyTorch 1.10.2+cu113.
+Install the required dependencies via
+```
+pip install -r requirements.txt
+```
+
+Alternatively you can also manually install the following packages in your virtual environment:
+- `torch`, `torchvision`, and `torchaudio` (e.g., with `--extra-index-url https://download.pytorch.org/whl/cu113` for CUDA 11.3)
+- `matplotlib`
+- `scipy`
+- `tensorboard`
+- `opencv-python`
+- `tqdm`
+- `parse`
+- `timm`
+- `flowpy`
+
+
+## Pre-Trained Checkpoints
+
+You can find our trained models in the release assets.
+
+
+## Datasets
+
+Datasets are expected to be located under `./data` in the following layout:
+```
+./data
+  ├── kitti15                   # KITTI 2015
+  │  └── dataset
+  │     ├── testing/...
+  │     └── training/...
+  └── sintel                    # Sintel
+     ├── test/...
+     └── training/...
+  
+```
+
+## Running CCMR(+)
+
+For running `CCMR+` on MPI Sintel images you need about 4.5 GB of GPU VRAM. `CCMR` (the 3-scale version) needs about 3 GBs of VRAM, using the following Cuda module.
+ 
+To compile the CUDA correlation module run the following once:
+```Shell
+cd alt_cuda_corr && python setup.py install && cd ..
+```
+
+And to reproduce our benchmark results after finetuning run:
+```Shell
+python evaluate.py --model_type "CCMR+" --model models/CCMR+_sintel.pth --dataset sintel_test
+python evaluate.py --model_type "CCMR+" --model models/CCMR+_kitti.pth --dataset kitti_test
+
+python evaluate.py --model_type "CCMR" --model models/CCMR_sintel.pth --dataset sintel_test
+python evaluate.py --model_type "CCMR" --model models/CCMR_kitti.pth --dataset kitti_test
+```
+
+## License
+- Our code is licensed under the BSD 3-Clause **No Military** License. See [LICENSE](LICENSE).
+- The provided checkpoints are under the [CC BY-NC-SA 3.0](https://creativecommons.org/licenses/by-nc-sa/3.0/) license.
+
+## Acknowledgement
+Parts of this repository are adapted from [RAFT](https://github.com/princeton-vl/RAFT) ([license](licenses/RAFT/LICENSE)), [MS-RAFT+](https://github.com/cv-stuttgart/MS_RAFT_plus) ([license](https://github.com/cv-stuttgart/MS_RAFT_plus/blob/main/LICENSE)), and [XCiT](https://github.com/facebookresearch/xcit/tree/main) ([license](https://github.com/facebookresearch/xcit/blob/main/LICENSE)).
+We thank the authors.

alt_cuda_corr/correlation.cpp

@@ -0,0 +1,59 @@
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <vector>
+
+// CUDA forward declarations
+std::vector<torch::Tensor> corr_cuda_forward(
+    torch::Tensor fmap1,
+    torch::Tensor fmap2,
+    torch::Tensor coords,
+    int radius);
+
+std::vector<torch::Tensor> corr_cuda_backward(
+  torch::Tensor fmap1,
+  torch::Tensor fmap2,
+  torch::Tensor coords,
+  torch::Tensor corr_grad,
+  int radius);
+
+// C++ interface
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> corr_forward(
+    torch::Tensor fmap1,
+    torch::Tensor fmap2,
+    torch::Tensor coords,
+    int radius) {
+  CHECK_INPUT(fmap1);
+  CHECK_INPUT(fmap2);
+  CHECK_INPUT(coords);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(coords));
+
+  return corr_cuda_forward(fmap1, fmap2, coords, radius);
+}
+
+
+std::vector<torch::Tensor> corr_backward(
+    torch::Tensor fmap1,
+    torch::Tensor fmap2,
+    torch::Tensor coords,
+    torch::Tensor corr_grad,
+    int radius) {
+  CHECK_INPUT(fmap1);
+  CHECK_INPUT(fmap2);
+  CHECK_INPUT(coords);
+  CHECK_INPUT(corr_grad);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(coords));
+
+  return corr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius);
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &corr_forward, "CORR forward");
+  m.def("backward", &corr_backward, "CORR backward");
+}