├── README.md ├── _config.yml ├── figure ├── qualitative.png ├── quantitative_flow.png └── quantitative_stereo.png ├── models ├── MonodepthModel.py ├── PWC_net.py ├── __init__.py ├── __pycache__ │ ├── MonodepthModel.cpython-36.pyc │ ├── MonodepthModel.cpython-37.pyc │ ├── PWC_net.cpython-36.pyc │ ├── PWC_net.cpython-37.pyc │ ├── __init__.cpython-36.pyc │ └── __init__.cpython-37.pyc └── networks │ ├── FlowNetC.py │ ├── FlowNetFusion.py │ ├── FlowNetS.py │ ├── FlowNetSD.py │ ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── submodules.cpython-36.pyc │ └── submodules.cpython-37.pyc │ ├── channelnorm_package │ ├── _ext │ │ └── channelnorm │ │ │ ├── __init__.py │ │ │ └── _channelnorm.so │ ├── build.py │ ├── build │ │ ├── lib.linux-x86_64-3.7 │ │ │ └── channelnorm_cuda.cpython-37m-x86_64-linux-gnu.so │ │ └── temp.linux-x86_64-3.7 │ │ │ ├── channelnorm_cuda.o │ │ │ └── channelnorm_kernel.o │ ├── channelnorm.py │ ├── channelnorm_cuda.cc │ ├── channelnorm_cuda.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── channelnorm_kernel.cu │ ├── channelnorm_kernel.cuh │ ├── dist │ │ └── channelnorm_cuda-0.0.0-py3.7-linux-x86_64.egg │ ├── functions │ │ └── channelnorm.py │ ├── make.sh │ ├── modules │ │ └── channelnorm.py │ ├── setup.py │ └── src │ │ ├── ChannelNorm_cuda.c │ │ ├── ChannelNorm_cuda.h │ │ ├── ChannelNorm_kernel.cu │ │ ├── ChannelNorm_kernel.h │ │ └── ChannelNorm_kernel.o │ ├── correlation_package │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-37.pyc │ │ └── correlation.cpython-37.pyc │ ├── _ext │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ └── correlation │ │ │ ├── __init__.py │ │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ │ └── _correlation.so │ ├── build.py │ ├── build │ │ ├── lib.linux-x86_64-3.7 │ │ │ └── correlation_cuda.cpython-37m-x86_64-linux-gnu.so │ │ └── temp.linux-x86_64-3.7 │ │ │ ├── correlation_cuda.o │ │ │ └── correlation_cuda_kernel.o │ ├── correlation.py │ ├── correlation_cuda.cc │ ├── correlation_cuda.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── correlation_cuda_kernel.cu │ ├── correlation_cuda_kernel.cuh │ ├── dist │ │ └── correlation_cuda-0.0.0-py3.7-linux-x86_64.egg │ ├── functions │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── correlation.cpython-36.pyc │ │ └── correlation.py │ ├── make.sh │ ├── modules │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── correlation.cpython-36.pyc │ │ └── correlation.py │ ├── setup.py │ └── src │ │ ├── correlation.c │ │ ├── correlation.h │ │ ├── correlation_cuda.c │ │ ├── correlation_cuda.h │ │ ├── correlation_cuda_kernel.cu │ │ ├── correlation_cuda_kernel.h │ │ └── correlation_cuda_kernel.o │ ├── resample2d_package │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-37.pyc │ │ └── resample2d.cpython-37.pyc │ ├── _ext │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ └── resample2d │ │ │ ├── __init__.py │ │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ │ └── _resample2d.so │ ├── build.py │ ├── build │ │ ├── lib.linux-x86_64-3.7 │ │ │ └── resample2d_cuda.cpython-37m-x86_64-linux-gnu.so │ │ └── temp.linux-x86_64-3.7 │ │ │ ├── resample2d_cuda.o │ │ │ └── resample2d_kernel.o │ ├── dist │ │ └── resample2d_cuda-0.0.0-py3.7-linux-x86_64.egg │ ├── functions │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── resample2d.cpython-36.pyc │ │ └── resample2d.py │ ├── make.sh │ ├── modules │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── resample2d.cpython-36.pyc │ │ └── resample2d.py │ ├── resample2d.py │ ├── resample2d_cuda.cc │ ├── resample2d_cuda.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── resample2d_kernel.cu │ ├── resample2d_kernel.cuh │ ├── setup.py │ └── src │ │ ├── Resample2d_cuda.c │ │ ├── Resample2d_cuda.h │ │ ├── Resample2d_kernel.cu │ │ ├── Resample2d_kernel.h │ │ └── Resample2d_kernel.o │ └── submodules.py ├── test_flow.py ├── test_stereo.py ├── train.py └── utils ├── __pycache__ ├── evaluate_utils.cpython-36.pyc ├── evaluate_utils.cpython-37.pyc ├── scene_dataloader.cpython-36.pyc ├── scene_dataloader.cpython-37.pyc ├── utils.cpython-36.pyc └── utils.cpython-37.pyc ├── evaluate_kitti.py ├── evaluation_utils.py ├── filenames ├── eigen_cycle_8000.txt ├── eigen_test_files.txt ├── kitti_flow_val_files_194_2012.txt ├── kitti_flow_val_files_200.txt ├── kitti_flow_val_files_occ_194_2012.txt ├── kitti_flow_val_files_occ_200.txt ├── kitti_stereo_2015_test_files.txt └── kitti_train_files_png_4frames.txt ├── kitti_archives_to_download.txt ├── scene_dataloader.py └── utils.py /README.md: -------------------------------------------------------------------------------- 1 | # Bridging Stereo Matching and Optical Flow via Spatiotemporal Correspondence 2 | PyTorch implementaton of the following paper. In this paper, we propose a unified model for unsupervised stereo matching and optical flow estimation using a single neural network. 3 | 4 | # Paper 5 | [Bridging Stereo Matching and Optical Flow via Spatiotemporal Correspondence](https://arxiv.org/abs/1905.09265)
6 | [Hsueh-Ying Lai](https://lelimite4444.github.io/), [Yi-Hsuan Tsai](https://sites.google.com/site/yihsuantsai/home), [Wei-Chen Chiu](https://walonchiu.github.io)
7 | IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019. 8 | 9 | Please cite our paper if you find it useful for your research. [[Project Page]](https://lelimite4444.github.io/BridgeDepthFlow-Project-Page/) 10 | 11 | ``` 12 | @inproceedings{lai19cvpr, 13 | title = {Bridging Stereo Matching and Optical Flow via Spatiotemporal Correspondence}, 14 | author = {Hsueh-Ying Lai and Yi-Hsuan Tsai and Wei-Chen Chiu}, 15 | booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 16 | year = {2019} 17 | } 18 | ``` 19 | 20 | ## Example Results 21 | 22 | ![](figure/qualitative.png) 23 | 24 | ## KITTI Dataset 25 | * Our model requires rectified stereo pairs with different timestamps from KITTI for training. \ 26 | We use two different splits of KITTI 2015, **kitti** and **eigen**, for both training and testing. For additional testing, we test on the validation set of KITTI 2012. You can find them in the [filenames](utils/filenames) folder. 27 | * Download the raw data of [KITTI Dataset](http://www.cvlibs.net/datasets/kitti/raw_data.php). This dataset is for training and eigen split evaluation. 28 | ``` 29 | wget -i utils/kitti_archives_to_download.txt -P ~/dataset/ 30 | ``` 31 | * Download [KITTI 2015 scene flow dataset](http://www.cvlibs.net/datasets/kitti/eval_scene_flow.php) and save it in KITTI_PATH. This dataset is for optical flow and kitti split evaluation. 32 | ```shell 33 | KITTI_PATH/ 34 | |--training/ 35 | |--testing/ 36 | ``` 37 | 38 | ## Installation 39 | * This code was developed using Python 3.7 & PyTorch 1.0.0 & CUDA 8.0. 40 | * Other requirements: cv2, matplotlib 41 | * Clone this repo 42 | ```shell 43 | git clone https://github.com/lelimite4444/BridgeDepthFlow 44 | cd BridgeDepthFlow 45 | ``` 46 | 47 | ## Training 48 | We use **kitti** split as example. 49 | ```shell 50 | python train.py --data_path ~/dataset/ 51 | --filenames_file ./utils/filenames/kitti_train_files_png_4frames.txt 52 | --checkpoint_path YOUR_CHECKPOINT_PATH 53 | ``` 54 | The chosen `--type_of_2warp` from 0 ~ 2 correponds to three types of different 2warp function in Figure 4 of our paper. 55 | The `--model_name` flag allows you to choose which model you want to train on. We provide the PyTorch version of both monodepth and PWC-Net. 56 | 57 | ## Testing 58 | We use the validation set of KITTI 2015 as example. The ground truth of optical flow includes occluded area. 59 | * You can download our pretrained models from [here](https://drive.google.com/drive/folders/1l6HLF2Dq1ObsTjdJaAPKK6NOAZDOI0Rs?usp=sharing), the final character of the model name correponds to the type of 2warp in our paper. 60 | * Test on optical flow 61 | ```shell 62 | python test_flow.py --data_path KITTI_PATH 63 | --filenames_file ./utils/filenames/kitti_flow_val_files_occ_200.txt 64 | --checkpoint_path YOUR_CHECKPOINT_PATH/TRAINED_MODEL_NAME 65 | ``` 66 | * Test on stereo matching 67 | ```shell 68 | python test_stereo.py --data_path KITTI_PATH 69 | --filenames_file ./utils/filenames/kitti_stereo_2015_test_files.txt 70 | --checkpoint_path YOUR_CHECKPOINT_PATH/TRAINED_MODEL_NAME 71 | ``` 72 | The network will output `disparities.npy`, containing all the estimated disparities of test data. You need to evaluate it by running: 73 | ```shell 74 | python utils/evaluate_kitti.py --split kitti --predicted_disp_path ./disparities.npy --gt_path ~/dataset/ 75 | ``` 76 | 77 | ## Acknowledgement 78 | * The evaluation code of stereo matching and the structure of monodepth is borrowed from [monodepth](https://github.com/mrharicot/monodepth) 79 | * The PWC-Net is implemented by [NVlabs-PWC-Net](https://github.com/NVlabs/PWC-Net/tree/master/PyTorch) 80 | * The warping function `Resample2d` and custom layers `Correlation` which PWC-Net relys on are implemented by [NVIDIA-flownet2](https://github.com/NVIDIA/flownet2-pytorch) 81 | -------------------------------------------------------------------------------- /_config.yml: -------------------------------------------------------------------------------- 1 | theme: jekyll-theme-cayman -------------------------------------------------------------------------------- /figure/qualitative.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/figure/qualitative.png -------------------------------------------------------------------------------- /figure/quantitative_flow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/figure/quantitative_flow.png -------------------------------------------------------------------------------- /figure/quantitative_stereo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/figure/quantitative_stereo.png -------------------------------------------------------------------------------- /models/MonodepthModel.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | import torchvision.transforms as transforms 4 | import matplotlib.pyplot as plt 5 | import numpy as np 6 | from torch.autograd import Variable 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.optim as optim 10 | import math 11 | import argparse 12 | import matplotlib.pyplot as plt 13 | import matplotlib.image as mpimg 14 | from PIL import Image 15 | 16 | class Bottleneck(nn.Module): 17 | expansion = 4 18 | 19 | def __init__(self, inplanes, planes, stride, downsample=None): 20 | super(Bottleneck, self).__init__() 21 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1) 22 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1) 23 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1) 24 | self.elu = nn.ELU() 25 | self.downsample = downsample 26 | self.stride = stride 27 | 28 | self.shortcut = nn.Sequential(nn.Conv2d(inplanes, self.expansion*planes, kernel_size=1, stride=stride)) 29 | 30 | def forward(self, x): 31 | residual = x 32 | 33 | out = self.conv1(x) 34 | out = self.elu(out) 35 | 36 | out = self.conv2(out) 37 | out = self.elu(out) 38 | 39 | out = self.conv3(out) 40 | 41 | out += self.shortcut(x) 42 | out = self.elu(out) 43 | 44 | return out 45 | 46 | class DecoderBlock(nn.Module): 47 | 48 | def __init__(self, inplanes, planes, mid, stride=1): 49 | super(DecoderBlock, self).__init__() 50 | self.upsample = nn.Upsample(scale_factor=2) 51 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=1, padding=1) 52 | self.conv2 = nn.Conv2d(mid, planes, kernel_size=3, stride=1, padding=1) 53 | self.elu = nn.ELU() 54 | 55 | def forward(self, x, skip, udisp=None): 56 | x = self.upsample(x) 57 | x = self.elu(self.conv1(x)) 58 | #concat 59 | if udisp is not None: 60 | x = torch.cat((x, skip, udisp), 1) 61 | else: 62 | x = torch.cat((x, skip), 1) 63 | x = self.elu(self.conv2(x)) 64 | return x 65 | 66 | class DispBlock(nn.Module): 67 | 68 | def __init__(self, inplanes, planes=2, kernel=3, stride=1): 69 | super(DispBlock, self).__init__() 70 | self.conv1 = nn.Conv2d(inplanes, 2, kernel_size=kernel, stride=stride, padding=1) 71 | self.sigmoid = nn.Sigmoid() 72 | self.tanh = nn.Tanh() 73 | self.upsample = nn.Upsample(scale_factor=2) 74 | 75 | def forward(self, x): 76 | disp = 0.3 * self.tanh(self.conv1(x)) 77 | udisp = self.upsample(disp) 78 | out = torch.cat((disp[:,0,:,:].unsqueeze(1) * x.shape[3], disp[:,1,:,:].unsqueeze(1) * x.shape[2]), 1) 79 | return out, disp, udisp 80 | 81 | class MonodepthNet(nn.Module): 82 | 83 | def __init__(self, block=Bottleneck, layers=[3, 4, 6, 3], num_classes=1000, do_stereo = 1): 84 | self.inplanes = 64 85 | super(MonodepthNet, self).__init__() 86 | if do_stereo: 87 | self.conv1 = nn.Conv2d(6, 64, kernel_size=7, stride=2, padding=3) 88 | else: 89 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) 90 | self.elu = nn.ELU() 91 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 92 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2) 93 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 94 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 95 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 96 | self.up6 = DecoderBlock(2048, 512, 1536) 97 | self.up5 = DecoderBlock(512, 256, 768) 98 | self.up4 = DecoderBlock(256, 128, 384) 99 | self.get_disp4 = DispBlock(128) 100 | self.up3 = DecoderBlock(128, 64, 130) 101 | self.get_disp3 = DispBlock(64) 102 | self.up2 = DecoderBlock(64, 32, 98) 103 | self.get_disp2 = DispBlock(32) 104 | self.up1 = DecoderBlock(32, 16, 18) 105 | self.get_disp1 = DispBlock(16) 106 | 107 | 108 | def _make_layer(self, block, planes, blocks, stride=1): 109 | layers = [] 110 | layers.append(block(self.inplanes, planes, 1)) 111 | self.inplanes = planes * block.expansion 112 | for i in range(1, blocks - 1): 113 | layers.append(block(self.inplanes, planes, 1)) 114 | layers.append(block(self.inplanes, planes, 2)) 115 | 116 | return nn.Sequential(*layers) 117 | 118 | def forward(self, x): 119 | #encoder 120 | x = self.conv1(x) #2 121 | conv1 = self.elu(x) 122 | pool1 = self.maxpool(conv1) #4 123 | conv2 = self.layer1(pool1) #8 124 | conv3 = self.layer2(conv2) #16 125 | conv4 = self.layer3(conv3) #32 126 | conv5 = self.layer4(conv4) #64 127 | 128 | #skip 129 | skip1 = conv1 130 | skip2 = pool1 131 | skip3 = conv2 132 | skip4 = conv3 133 | skip5 = conv4 134 | 135 | #decoder 136 | upconv6 = self.up6(conv5, skip5) 137 | upconv5 = self.up5(upconv6, skip4) 138 | upconv4 = self.up4(upconv5, skip3) 139 | self.disp4_scale, self.disp4, udisp4 = self.get_disp4(upconv4) 140 | 141 | upconv3 = self.up3(upconv4, skip2, udisp4) 142 | self.disp3_scale, self.disp3, udisp3 = self.get_disp3(upconv3) 143 | 144 | upconv2 = self.up2(upconv3, skip1, udisp3) 145 | self.disp2_scale, self.disp2, udisp2 = self.get_disp2(upconv2) 146 | 147 | upconv1 = self.up1(upconv2, udisp2) 148 | self.disp1_scale, self.disp1, udisp1 = self.get_disp1(upconv1) 149 | 150 | return [self.disp1_scale, self.disp2_scale, self.disp3_scale, self.disp4_scale], [self.disp1, self.disp2, self.disp3, self.disp4] 151 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from .MonodepthModel import * 2 | from .PWC_net import * 3 | from .PWC_net import PWCDCNet 4 | -------------------------------------------------------------------------------- /models/__pycache__/MonodepthModel.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/__pycache__/MonodepthModel.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/MonodepthModel.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/__pycache__/MonodepthModel.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/PWC_net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/__pycache__/PWC_net.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/PWC_net.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/__pycache__/PWC_net.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/FlowNetC.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | 5 | import math 6 | import numpy as np 7 | 8 | from .correlation_package.correlation import Correlation 9 | 10 | from .submodules import * 11 | 'Parameter count , 39,175,298 ' 12 | 13 | class FlowNetC(nn.Module): 14 | def __init__(self,args, batchNorm=True, div_flow = 20): 15 | super(FlowNetC,self).__init__() 16 | 17 | self.batchNorm = batchNorm 18 | self.div_flow = div_flow 19 | 20 | self.conv1 = conv(self.batchNorm, 3, 64, kernel_size=7, stride=2) 21 | self.conv2 = conv(self.batchNorm, 64, 128, kernel_size=5, stride=2) 22 | self.conv3 = conv(self.batchNorm, 128, 256, kernel_size=5, stride=2) 23 | self.conv_redir = conv(self.batchNorm, 256, 32, kernel_size=1, stride=1) 24 | 25 | if args.fp16: 26 | self.corr = nn.Sequential( 27 | tofp32(), 28 | Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1), 29 | tofp16()) 30 | else: 31 | self.corr = Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1) 32 | 33 | self.corr_activation = nn.LeakyReLU(0.1,inplace=True) 34 | self.conv3_1 = conv(self.batchNorm, 473, 256) 35 | self.conv4 = conv(self.batchNorm, 256, 512, stride=2) 36 | self.conv4_1 = conv(self.batchNorm, 512, 512) 37 | self.conv5 = conv(self.batchNorm, 512, 512, stride=2) 38 | self.conv5_1 = conv(self.batchNorm, 512, 512) 39 | self.conv6 = conv(self.batchNorm, 512, 1024, stride=2) 40 | self.conv6_1 = conv(self.batchNorm,1024, 1024) 41 | 42 | self.deconv5 = deconv(1024,512) 43 | self.deconv4 = deconv(1026,256) 44 | self.deconv3 = deconv(770,128) 45 | self.deconv2 = deconv(386,64) 46 | 47 | self.predict_flow6 = predict_flow(1024) 48 | self.predict_flow5 = predict_flow(1026) 49 | self.predict_flow4 = predict_flow(770) 50 | self.predict_flow3 = predict_flow(386) 51 | self.predict_flow2 = predict_flow(194) 52 | 53 | self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True) 54 | self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True) 55 | self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True) 56 | self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True) 57 | 58 | for m in self.modules(): 59 | if isinstance(m, nn.Conv2d): 60 | if m.bias is not None: 61 | init.uniform_(m.bias) 62 | init.xavier_uniform_(m.weight) 63 | 64 | if isinstance(m, nn.ConvTranspose2d): 65 | if m.bias is not None: 66 | init.uniform_(m.bias) 67 | init.xavier_uniform_(m.weight) 68 | # init_deconv_bilinear(m.weight) 69 | self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear') 70 | 71 | def forward(self, x): 72 | x1 = x[:,0:3,:,:] 73 | x2 = x[:,3::,:,:] 74 | 75 | out_conv1a = self.conv1(x1) 76 | out_conv2a = self.conv2(out_conv1a) 77 | out_conv3a = self.conv3(out_conv2a) 78 | 79 | # FlownetC bottom input stream 80 | out_conv1b = self.conv1(x2) 81 | 82 | out_conv2b = self.conv2(out_conv1b) 83 | out_conv3b = self.conv3(out_conv2b) 84 | 85 | # Merge streams 86 | out_corr = self.corr(out_conv3a, out_conv3b) # False 87 | out_corr = self.corr_activation(out_corr) 88 | 89 | # Redirect top input stream and concatenate 90 | out_conv_redir = self.conv_redir(out_conv3a) 91 | 92 | in_conv3_1 = torch.cat((out_conv_redir, out_corr), 1) 93 | 94 | # Merged conv layers 95 | out_conv3_1 = self.conv3_1(in_conv3_1) 96 | 97 | out_conv4 = self.conv4_1(self.conv4(out_conv3_1)) 98 | 99 | out_conv5 = self.conv5_1(self.conv5(out_conv4)) 100 | out_conv6 = self.conv6_1(self.conv6(out_conv5)) 101 | 102 | flow6 = self.predict_flow6(out_conv6) 103 | flow6_up = self.upsampled_flow6_to_5(flow6) 104 | out_deconv5 = self.deconv5(out_conv6) 105 | 106 | concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1) 107 | 108 | flow5 = self.predict_flow5(concat5) 109 | flow5_up = self.upsampled_flow5_to_4(flow5) 110 | out_deconv4 = self.deconv4(concat5) 111 | concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1) 112 | 113 | flow4 = self.predict_flow4(concat4) 114 | flow4_up = self.upsampled_flow4_to_3(flow4) 115 | out_deconv3 = self.deconv3(concat4) 116 | concat3 = torch.cat((out_conv3_1,out_deconv3,flow4_up),1) 117 | 118 | flow3 = self.predict_flow3(concat3) 119 | flow3_up = self.upsampled_flow3_to_2(flow3) 120 | out_deconv2 = self.deconv2(concat3) 121 | concat2 = torch.cat((out_conv2a,out_deconv2,flow3_up),1) 122 | 123 | flow2 = self.predict_flow2(concat2) 124 | 125 | if self.training: 126 | return flow2,flow3,flow4,flow5,flow6 127 | else: 128 | return flow2, 129 | -------------------------------------------------------------------------------- /models/networks/FlowNetFusion.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | 5 | import math 6 | import numpy as np 7 | 8 | from .submodules import * 9 | 'Parameter count = 581,226' 10 | 11 | class FlowNetFusion(nn.Module): 12 | def __init__(self,args, batchNorm=True): 13 | super(FlowNetFusion,self).__init__() 14 | 15 | self.batchNorm = batchNorm 16 | self.conv0 = conv(self.batchNorm, 11, 64) 17 | self.conv1 = conv(self.batchNorm, 64, 64, stride=2) 18 | self.conv1_1 = conv(self.batchNorm, 64, 128) 19 | self.conv2 = conv(self.batchNorm, 128, 128, stride=2) 20 | self.conv2_1 = conv(self.batchNorm, 128, 128) 21 | 22 | self.deconv1 = deconv(128,32) 23 | self.deconv0 = deconv(162,16) 24 | 25 | self.inter_conv1 = i_conv(self.batchNorm, 162, 32) 26 | self.inter_conv0 = i_conv(self.batchNorm, 82, 16) 27 | 28 | self.predict_flow2 = predict_flow(128) 29 | self.predict_flow1 = predict_flow(32) 30 | self.predict_flow0 = predict_flow(16) 31 | 32 | self.upsampled_flow2_to_1 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 33 | self.upsampled_flow1_to_0 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 34 | 35 | for m in self.modules(): 36 | if isinstance(m, nn.Conv2d): 37 | if m.bias is not None: 38 | init.uniform_(m.bias) 39 | init.xavier_uniform_(m.weight) 40 | 41 | if isinstance(m, nn.ConvTranspose2d): 42 | if m.bias is not None: 43 | init.uniform_(m.bias) 44 | init.xavier_uniform_(m.weight) 45 | # init_deconv_bilinear(m.weight) 46 | 47 | def forward(self, x): 48 | out_conv0 = self.conv0(x) 49 | out_conv1 = self.conv1_1(self.conv1(out_conv0)) 50 | out_conv2 = self.conv2_1(self.conv2(out_conv1)) 51 | 52 | flow2 = self.predict_flow2(out_conv2) 53 | flow2_up = self.upsampled_flow2_to_1(flow2) 54 | out_deconv1 = self.deconv1(out_conv2) 55 | 56 | concat1 = torch.cat((out_conv1,out_deconv1,flow2_up),1) 57 | out_interconv1 = self.inter_conv1(concat1) 58 | flow1 = self.predict_flow1(out_interconv1) 59 | flow1_up = self.upsampled_flow1_to_0(flow1) 60 | out_deconv0 = self.deconv0(concat1) 61 | 62 | concat0 = torch.cat((out_conv0,out_deconv0,flow1_up),1) 63 | out_interconv0 = self.inter_conv0(concat0) 64 | flow0 = self.predict_flow0(out_interconv0) 65 | 66 | return flow0 67 | 68 | -------------------------------------------------------------------------------- /models/networks/FlowNetS.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Portions of this code copyright 2017, Clement Pinard 3 | ''' 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.nn import init 8 | 9 | import math 10 | import numpy as np 11 | 12 | from .submodules import * 13 | 'Parameter count : 38,676,504 ' 14 | 15 | class FlowNetS(nn.Module): 16 | def __init__(self, args, input_channels = 12, batchNorm=True): 17 | super(FlowNetS,self).__init__() 18 | 19 | self.batchNorm = batchNorm 20 | self.conv1 = conv(self.batchNorm, input_channels, 64, kernel_size=7, stride=2) 21 | self.conv2 = conv(self.batchNorm, 64, 128, kernel_size=5, stride=2) 22 | self.conv3 = conv(self.batchNorm, 128, 256, kernel_size=5, stride=2) 23 | self.conv3_1 = conv(self.batchNorm, 256, 256) 24 | self.conv4 = conv(self.batchNorm, 256, 512, stride=2) 25 | self.conv4_1 = conv(self.batchNorm, 512, 512) 26 | self.conv5 = conv(self.batchNorm, 512, 512, stride=2) 27 | self.conv5_1 = conv(self.batchNorm, 512, 512) 28 | self.conv6 = conv(self.batchNorm, 512, 1024, stride=2) 29 | self.conv6_1 = conv(self.batchNorm,1024, 1024) 30 | 31 | self.deconv5 = deconv(1024,512) 32 | self.deconv4 = deconv(1026,256) 33 | self.deconv3 = deconv(770,128) 34 | self.deconv2 = deconv(386,64) 35 | 36 | self.predict_flow6 = predict_flow(1024) 37 | self.predict_flow5 = predict_flow(1026) 38 | self.predict_flow4 = predict_flow(770) 39 | self.predict_flow3 = predict_flow(386) 40 | self.predict_flow2 = predict_flow(194) 41 | 42 | self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False) 43 | self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False) 44 | self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False) 45 | self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False) 46 | 47 | for m in self.modules(): 48 | if isinstance(m, nn.Conv2d): 49 | if m.bias is not None: 50 | init.uniform_(m.bias) 51 | init.xavier_uniform_(m.weight) 52 | 53 | if isinstance(m, nn.ConvTranspose2d): 54 | if m.bias is not None: 55 | init.uniform_(m.bias) 56 | init.xavier_uniform_(m.weight) 57 | # init_deconv_bilinear(m.weight) 58 | self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear') 59 | 60 | def forward(self, x): 61 | out_conv1 = self.conv1(x) 62 | 63 | out_conv2 = self.conv2(out_conv1) 64 | out_conv3 = self.conv3_1(self.conv3(out_conv2)) 65 | out_conv4 = self.conv4_1(self.conv4(out_conv3)) 66 | out_conv5 = self.conv5_1(self.conv5(out_conv4)) 67 | out_conv6 = self.conv6_1(self.conv6(out_conv5)) 68 | 69 | flow6 = self.predict_flow6(out_conv6) 70 | flow6_up = self.upsampled_flow6_to_5(flow6) 71 | out_deconv5 = self.deconv5(out_conv6) 72 | 73 | concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1) 74 | flow5 = self.predict_flow5(concat5) 75 | flow5_up = self.upsampled_flow5_to_4(flow5) 76 | out_deconv4 = self.deconv4(concat5) 77 | 78 | concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1) 79 | flow4 = self.predict_flow4(concat4) 80 | flow4_up = self.upsampled_flow4_to_3(flow4) 81 | out_deconv3 = self.deconv3(concat4) 82 | 83 | concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1) 84 | flow3 = self.predict_flow3(concat3) 85 | flow3_up = self.upsampled_flow3_to_2(flow3) 86 | out_deconv2 = self.deconv2(concat3) 87 | 88 | concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1) 89 | flow2 = self.predict_flow2(concat2) 90 | 91 | if self.training: 92 | return flow2,flow3,flow4,flow5,flow6 93 | else: 94 | return flow2, 95 | 96 | -------------------------------------------------------------------------------- /models/networks/FlowNetSD.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | 5 | import math 6 | import numpy as np 7 | 8 | from .submodules import * 9 | 'Parameter count = 45,371,666' 10 | 11 | class FlowNetSD(nn.Module): 12 | def __init__(self, args, batchNorm=True): 13 | super(FlowNetSD,self).__init__() 14 | 15 | self.batchNorm = batchNorm 16 | self.conv0 = conv(self.batchNorm, 6, 64) 17 | self.conv1 = conv(self.batchNorm, 64, 64, stride=2) 18 | self.conv1_1 = conv(self.batchNorm, 64, 128) 19 | self.conv2 = conv(self.batchNorm, 128, 128, stride=2) 20 | self.conv2_1 = conv(self.batchNorm, 128, 128) 21 | self.conv3 = conv(self.batchNorm, 128, 256, stride=2) 22 | self.conv3_1 = conv(self.batchNorm, 256, 256) 23 | self.conv4 = conv(self.batchNorm, 256, 512, stride=2) 24 | self.conv4_1 = conv(self.batchNorm, 512, 512) 25 | self.conv5 = conv(self.batchNorm, 512, 512, stride=2) 26 | self.conv5_1 = conv(self.batchNorm, 512, 512) 27 | self.conv6 = conv(self.batchNorm, 512, 1024, stride=2) 28 | self.conv6_1 = conv(self.batchNorm,1024, 1024) 29 | 30 | self.deconv5 = deconv(1024,512) 31 | self.deconv4 = deconv(1026,256) 32 | self.deconv3 = deconv(770,128) 33 | self.deconv2 = deconv(386,64) 34 | 35 | self.inter_conv5 = i_conv(self.batchNorm, 1026, 512) 36 | self.inter_conv4 = i_conv(self.batchNorm, 770, 256) 37 | self.inter_conv3 = i_conv(self.batchNorm, 386, 128) 38 | self.inter_conv2 = i_conv(self.batchNorm, 194, 64) 39 | 40 | self.predict_flow6 = predict_flow(1024) 41 | self.predict_flow5 = predict_flow(512) 42 | self.predict_flow4 = predict_flow(256) 43 | self.predict_flow3 = predict_flow(128) 44 | self.predict_flow2 = predict_flow(64) 45 | 46 | self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 47 | self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 48 | self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 49 | self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1) 50 | 51 | for m in self.modules(): 52 | if isinstance(m, nn.Conv2d): 53 | if m.bias is not None: 54 | init.uniform_(m.bias) 55 | init.xavier_uniform_(m.weight) 56 | 57 | if isinstance(m, nn.ConvTranspose2d): 58 | if m.bias is not None: 59 | init.uniform_(m.bias) 60 | init.xavier_uniform_(m.weight) 61 | # init_deconv_bilinear(m.weight) 62 | self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear') 63 | 64 | 65 | 66 | def forward(self, x): 67 | out_conv0 = self.conv0(x) 68 | out_conv1 = self.conv1_1(self.conv1(out_conv0)) 69 | out_conv2 = self.conv2_1(self.conv2(out_conv1)) 70 | 71 | out_conv3 = self.conv3_1(self.conv3(out_conv2)) 72 | out_conv4 = self.conv4_1(self.conv4(out_conv3)) 73 | out_conv5 = self.conv5_1(self.conv5(out_conv4)) 74 | out_conv6 = self.conv6_1(self.conv6(out_conv5)) 75 | 76 | flow6 = self.predict_flow6(out_conv6) 77 | flow6_up = self.upsampled_flow6_to_5(flow6) 78 | out_deconv5 = self.deconv5(out_conv6) 79 | 80 | concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1) 81 | out_interconv5 = self.inter_conv5(concat5) 82 | flow5 = self.predict_flow5(out_interconv5) 83 | 84 | flow5_up = self.upsampled_flow5_to_4(flow5) 85 | out_deconv4 = self.deconv4(concat5) 86 | 87 | concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1) 88 | out_interconv4 = self.inter_conv4(concat4) 89 | flow4 = self.predict_flow4(out_interconv4) 90 | flow4_up = self.upsampled_flow4_to_3(flow4) 91 | out_deconv3 = self.deconv3(concat4) 92 | 93 | concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1) 94 | out_interconv3 = self.inter_conv3(concat3) 95 | flow3 = self.predict_flow3(out_interconv3) 96 | flow3_up = self.upsampled_flow3_to_2(flow3) 97 | out_deconv2 = self.deconv2(concat3) 98 | 99 | concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1) 100 | out_interconv2 = self.inter_conv2(concat2) 101 | flow2 = self.predict_flow2(out_interconv2) 102 | 103 | if self.training: 104 | return flow2,flow3,flow4,flow5,flow6 105 | else: 106 | return flow2, 107 | -------------------------------------------------------------------------------- /models/networks/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/__pycache__/submodules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/__pycache__/submodules.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/__pycache__/submodules.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/__pycache__/submodules.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/channelnorm_package/_ext/channelnorm/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from torch.utils.ffi import _wrap_function 3 | from ._channelnorm import lib as _lib, ffi as _ffi 4 | 5 | __all__ = [] 6 | def _import_symbols(locals): 7 | for symbol in dir(_lib): 8 | fn = getattr(_lib, symbol) 9 | if callable(fn): 10 | locals[symbol] = _wrap_function(fn, _ffi) 11 | else: 12 | locals[symbol] = fn 13 | __all__.append(symbol) 14 | 15 | _import_symbols(locals()) 16 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/_ext/channelnorm/_channelnorm.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/channelnorm_package/_ext/channelnorm/_channelnorm.so -------------------------------------------------------------------------------- /models/networks/channelnorm_package/build.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.utils.ffi 4 | 5 | this_folder = os.path.dirname(os.path.abspath(__file__)) + '/' 6 | 7 | Headers = [] 8 | Sources = [] 9 | Defines = [] 10 | Objects = [] 11 | 12 | if torch.cuda.is_available() == True: 13 | Headers += ['src/ChannelNorm_cuda.h'] 14 | Sources += ['src/ChannelNorm_cuda.c'] 15 | Defines += [('WITH_CUDA', None)] 16 | Objects += ['src/ChannelNorm_kernel.o'] 17 | 18 | ffi = torch.utils.ffi.create_extension( 19 | name='_ext.channelnorm', 20 | headers=Headers, 21 | sources=Sources, 22 | verbose=False, 23 | with_cuda=True, 24 | package=False, 25 | relative_to=this_folder, 26 | define_macros=Defines, 27 | extra_objects=[os.path.join(this_folder, Object) for Object in Objects] 28 | ) 29 | 30 | if __name__ == '__main__': 31 | ffi.build() -------------------------------------------------------------------------------- /models/networks/channelnorm_package/build/lib.linux-x86_64-3.7/channelnorm_cuda.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/channelnorm_package/build/lib.linux-x86_64-3.7/channelnorm_cuda.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /models/networks/channelnorm_package/build/temp.linux-x86_64-3.7/channelnorm_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/channelnorm_package/build/temp.linux-x86_64-3.7/channelnorm_cuda.o -------------------------------------------------------------------------------- /models/networks/channelnorm_package/build/temp.linux-x86_64-3.7/channelnorm_kernel.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/channelnorm_package/build/temp.linux-x86_64-3.7/channelnorm_kernel.o -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function, Variable 2 | from torch.nn.modules.module import Module 3 | import channelnorm_cuda 4 | 5 | class ChannelNormFunction(Function): 6 | 7 | @staticmethod 8 | def forward(ctx, input1, norm_deg=2): 9 | assert input1.is_contiguous() 10 | b, _, h, w = input1.size() 11 | output = input1.new(b, 1, h, w).zero_() 12 | 13 | channelnorm_cuda.forward(input1, output, norm_deg) 14 | ctx.save_for_backward(input1, output) 15 | ctx.norm_deg = norm_deg 16 | 17 | return output 18 | 19 | @staticmethod 20 | def backward(ctx, grad_output): 21 | input1, output = ctx.saved_tensors 22 | 23 | grad_input1 = Variable(input1.new(input1.size()).zero_()) 24 | 25 | channelnorm_cuda.backward(input1, output, grad_output.data, 26 | grad_input1.data, ctx.norm_deg) 27 | 28 | return grad_input1, None 29 | 30 | 31 | class ChannelNorm(Module): 32 | 33 | def __init__(self, norm_deg=2): 34 | super(ChannelNorm, self).__init__() 35 | self.norm_deg = norm_deg 36 | 37 | def forward(self, input1): 38 | return ChannelNormFunction.apply(input1, self.norm_deg) 39 | 40 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "channelnorm_kernel.cuh" 5 | 6 | int channelnorm_cuda_forward( 7 | at::Tensor& input1, 8 | at::Tensor& output, 9 | int norm_deg) { 10 | 11 | channelnorm_kernel_forward(input1, output, norm_deg); 12 | return 1; 13 | } 14 | 15 | 16 | int channelnorm_cuda_backward( 17 | at::Tensor& input1, 18 | at::Tensor& output, 19 | at::Tensor& gradOutput, 20 | at::Tensor& gradInput1, 21 | int norm_deg) { 22 | 23 | channelnorm_kernel_backward(input1, output, gradOutput, gradInput1, norm_deg); 24 | return 1; 25 | } 26 | 27 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 28 | m.def("forward", &channelnorm_cuda_forward, "Channel norm forward (CUDA)"); 29 | m.def("backward", &channelnorm_cuda_backward, "Channel norm backward (CUDA)"); 30 | } 31 | 32 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_cuda.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: channelnorm-cuda 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_cuda.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | channelnorm_cuda.cc 2 | channelnorm_kernel.cu 3 | setup.py 4 | channelnorm_cuda.egg-info/PKG-INFO 5 | channelnorm_cuda.egg-info/SOURCES.txt 6 | channelnorm_cuda.egg-info/dependency_links.txt 7 | channelnorm_cuda.egg-info/top_level.txt -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_cuda.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_cuda.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | channelnorm_cuda 2 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "channelnorm_kernel.cuh" 6 | 7 | #define CUDA_NUM_THREADS 512 8 | 9 | #define DIM0(TENSOR) ((TENSOR).x) 10 | #define DIM1(TENSOR) ((TENSOR).y) 11 | #define DIM2(TENSOR) ((TENSOR).z) 12 | #define DIM3(TENSOR) ((TENSOR).w) 13 | 14 | #define DIM3_INDEX(TENSOR, xx, yy, zz, ww) ((TENSOR)[((xx) * (TENSOR##_stride.x)) + ((yy) * (TENSOR##_stride.y)) + ((zz) * (TENSOR##_stride.z)) + ((ww) * (TENSOR##_stride.w))]) 15 | 16 | using at::Half; 17 | 18 | template 19 | __global__ void kernel_channelnorm_update_output( 20 | const int n, 21 | const scalar_t* __restrict__ input1, 22 | const long4 input1_size, 23 | const long4 input1_stride, 24 | scalar_t* __restrict__ output, 25 | const long4 output_size, 26 | const long4 output_stride, 27 | int norm_deg) { 28 | 29 | int index = blockIdx.x * blockDim.x + threadIdx.x; 30 | 31 | if (index >= n) { 32 | return; 33 | } 34 | 35 | int dim_b = DIM0(output_size); 36 | int dim_c = DIM1(output_size); 37 | int dim_h = DIM2(output_size); 38 | int dim_w = DIM3(output_size); 39 | int dim_chw = dim_c * dim_h * dim_w; 40 | 41 | int b = ( index / dim_chw ) % dim_b; 42 | int y = ( index / dim_w ) % dim_h; 43 | int x = ( index ) % dim_w; 44 | 45 | int i1dim_c = DIM1(input1_size); 46 | int i1dim_h = DIM2(input1_size); 47 | int i1dim_w = DIM3(input1_size); 48 | int i1dim_chw = i1dim_c * i1dim_h * i1dim_w; 49 | int i1dim_hw = i1dim_h * i1dim_w; 50 | 51 | float result = 0.0; 52 | 53 | for (int c = 0; c < i1dim_c; ++c) { 54 | int i1Index = b * i1dim_chw + c * i1dim_hw + y * i1dim_w + x; 55 | scalar_t val = input1[i1Index]; 56 | result += static_cast(val * val); 57 | } 58 | result = sqrt(result); 59 | output[index] = static_cast(result); 60 | } 61 | 62 | 63 | template 64 | __global__ void kernel_channelnorm_backward_input1( 65 | const int n, 66 | const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride, 67 | const scalar_t* __restrict__ output, const long4 output_size, const long4 output_stride, 68 | const scalar_t* __restrict__ gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride, 69 | scalar_t* __restrict__ gradInput, const long4 gradInput_size, const long4 gradInput_stride, 70 | int norm_deg) { 71 | 72 | int index = blockIdx.x * blockDim.x + threadIdx.x; 73 | 74 | if (index >= n) { 75 | return; 76 | } 77 | 78 | float val = 0.0; 79 | 80 | int dim_b = DIM0(gradInput_size); 81 | int dim_c = DIM1(gradInput_size); 82 | int dim_h = DIM2(gradInput_size); 83 | int dim_w = DIM3(gradInput_size); 84 | int dim_chw = dim_c * dim_h * dim_w; 85 | int dim_hw = dim_h * dim_w; 86 | 87 | int b = ( index / dim_chw ) % dim_b; 88 | int y = ( index / dim_w ) % dim_h; 89 | int x = ( index ) % dim_w; 90 | 91 | 92 | int outIndex = b * dim_hw + y * dim_w + x; 93 | val = static_cast(gradOutput[outIndex]) * static_cast(input1[index]) / (static_cast(output[outIndex])+1e-9); 94 | gradInput[index] = static_cast(val); 95 | 96 | } 97 | 98 | void channelnorm_kernel_forward( 99 | at::Tensor& input1, 100 | at::Tensor& output, 101 | int norm_deg) { 102 | 103 | const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3)); 104 | const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3)); 105 | 106 | const long4 output_size = make_long4(output.size(0), output.size(1), output.size(2), output.size(3)); 107 | const long4 output_stride = make_long4(output.stride(0), output.stride(1), output.stride(2), output.stride(3)); 108 | 109 | int n = output.numel(); 110 | 111 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), "channelnorm_forward", ([&] { 112 | 113 | kernel_channelnorm_update_output<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>( 114 | //at::globalContext().getCurrentCUDAStream() >>>( 115 | n, 116 | input1.data(), 117 | input1_size, 118 | input1_stride, 119 | output.data(), 120 | output_size, 121 | output_stride, 122 | norm_deg); 123 | 124 | })); 125 | 126 | // TODO: ATen-equivalent check 127 | 128 | // THCudaCheck(cudaGetLastError()); 129 | } 130 | 131 | void channelnorm_kernel_backward( 132 | at::Tensor& input1, 133 | at::Tensor& output, 134 | at::Tensor& gradOutput, 135 | at::Tensor& gradInput1, 136 | int norm_deg) { 137 | 138 | const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3)); 139 | const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3)); 140 | 141 | const long4 output_size = make_long4(output.size(0), output.size(1), output.size(2), output.size(3)); 142 | const long4 output_stride = make_long4(output.stride(0), output.stride(1), output.stride(2), output.stride(3)); 143 | 144 | const long4 gradOutput_size = make_long4(gradOutput.size(0), gradOutput.size(1), gradOutput.size(2), gradOutput.size(3)); 145 | const long4 gradOutput_stride = make_long4(gradOutput.stride(0), gradOutput.stride(1), gradOutput.stride(2), gradOutput.stride(3)); 146 | 147 | const long4 gradInput1_size = make_long4(gradInput1.size(0), gradInput1.size(1), gradInput1.size(2), gradInput1.size(3)); 148 | const long4 gradInput1_stride = make_long4(gradInput1.stride(0), gradInput1.stride(1), gradInput1.stride(2), gradInput1.stride(3)); 149 | 150 | int n = gradInput1.numel(); 151 | 152 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(input1.type(), "channelnorm_backward_input1", ([&] { 153 | 154 | kernel_channelnorm_backward_input1<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>( 155 | //at::globalContext().getCurrentCUDAStream() >>>( 156 | n, 157 | input1.data(), 158 | input1_size, 159 | input1_stride, 160 | output.data(), 161 | output_size, 162 | output_stride, 163 | gradOutput.data(), 164 | gradOutput_size, 165 | gradOutput_stride, 166 | gradInput1.data(), 167 | gradInput1_size, 168 | gradInput1_stride, 169 | norm_deg 170 | ); 171 | 172 | })); 173 | 174 | // TODO: Add ATen-equivalent check 175 | 176 | // THCudaCheck(cudaGetLastError()); 177 | } 178 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/channelnorm_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | void channelnorm_kernel_forward( 6 | at::Tensor& input1, 7 | at::Tensor& output, 8 | int norm_deg); 9 | 10 | 11 | void channelnorm_kernel_backward( 12 | at::Tensor& input1, 13 | at::Tensor& output, 14 | at::Tensor& gradOutput, 15 | at::Tensor& gradInput1, 16 | int norm_deg); 17 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/dist/channelnorm_cuda-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/channelnorm_package/dist/channelnorm_cuda-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /models/networks/channelnorm_package/functions/channelnorm.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function, Variable 2 | from .._ext import channelnorm 3 | 4 | 5 | class ChannelNormFunction(Function): 6 | 7 | @staticmethod 8 | def forward(ctx, input1, norm_deg=2): 9 | assert input1.is_contiguous() 10 | b, _, h, w = input1.size() 11 | output = input1.new(b, 1, h, w).zero_() 12 | 13 | channelnorm.ChannelNorm_cuda_forward(input1, output, norm_deg) 14 | ctx.save_for_backward(input1, output) 15 | ctx.norm_deg = norm_deg 16 | 17 | return output 18 | 19 | @staticmethod 20 | def backward(ctx, grad_output): 21 | input1, output = ctx.saved_tensors 22 | 23 | grad_input1 = Variable(input1.new(input1.size()).zero_()) 24 | 25 | channelnorm.ChannelNorm_cuda_backward(input1, output, grad_output.data, 26 | grad_input1.data, ctx.norm_deg) 27 | 28 | return grad_input1, None 29 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/make.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))") 3 | 4 | cd src 5 | echo "Compiling channelnorm kernels by nvcc..." 6 | rm ChannelNorm_kernel.o 7 | rm -r ../_ext 8 | 9 | nvcc -c -o ChannelNorm_kernel.o ChannelNorm_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52 -I ${TORCH}/lib/include/TH -I ${TORCH}/lib/include/THC 10 | 11 | cd ../ 12 | python build.py -------------------------------------------------------------------------------- /models/networks/channelnorm_package/modules/channelnorm.py: -------------------------------------------------------------------------------- 1 | from torch.nn.modules.module import Module 2 | 3 | from ..functions.channelnorm import ChannelNormFunction 4 | 5 | 6 | class ChannelNorm(Module): 7 | 8 | def __init__(self, norm_deg=2): 9 | super(ChannelNorm, self).__init__() 10 | self.norm_deg = norm_deg 11 | 12 | def forward(self, input1): 13 | return ChannelNormFunction.apply(input1, self.norm_deg) 14 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import os 3 | import torch 4 | 5 | from setuptools import setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | cxx_args = ['-std=c++11'] 9 | 10 | nvcc_args = [ 11 | '-gencode', 'arch=compute_52,code=sm_52', 12 | '-gencode', 'arch=compute_60,code=sm_60', 13 | '-gencode', 'arch=compute_61,code=sm_61', 14 | '-gencode', 'arch=compute_70,code=sm_70', 15 | '-gencode', 'arch=compute_70,code=compute_70' 16 | ] 17 | 18 | setup( 19 | name='channelnorm_cuda', 20 | ext_modules=[ 21 | CUDAExtension('channelnorm_cuda', [ 22 | 'channelnorm_cuda.cc', 23 | 'channelnorm_kernel.cu' 24 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 25 | ], 26 | cmdclass={ 27 | 'build_ext': BuildExtension 28 | }) 29 | -------------------------------------------------------------------------------- /models/networks/channelnorm_package/src/ChannelNorm_cuda.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "ChannelNorm_kernel.h" 5 | 6 | extern THCState* state; 7 | 8 | int ChannelNorm_cuda_forward(THCudaTensor* input1, THCudaTensor* output, int norm_deg) { 9 | ChannelNorm_kernel_forward(state, input1, output, norm_deg); 10 | return 1; 11 | } 12 | 13 | 14 | int ChannelNorm_cuda_backward(THCudaTensor* input1, THCudaTensor* output, THCudaTensor* gradOutput, THCudaTensor* gradInput1, int norm_deg) { 15 | ChannelNorm_kernel_backward(state, input1, output, gradOutput, gradInput1, norm_deg); 16 | return 1; 17 | } -------------------------------------------------------------------------------- /models/networks/channelnorm_package/src/ChannelNorm_cuda.h: -------------------------------------------------------------------------------- 1 | int ChannelNorm_cuda_forward(THCudaTensor* input1, THCudaTensor* output, int norm_deg); 2 | 3 | int ChannelNorm_cuda_backward(THCudaTensor* input1, THCudaTensor* output, THCudaTensor* gradOutput, THCudaTensor* gradInput1, int norm_deg); -------------------------------------------------------------------------------- /models/networks/channelnorm_package/src/ChannelNorm_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #define CUDA_NUM_THREADS 512 5 | #define THREADS_PER_BLOCK 64 6 | 7 | #define DIM0(TENSOR) ((TENSOR).x) 8 | #define DIM1(TENSOR) ((TENSOR).y) 9 | #define DIM2(TENSOR) ((TENSOR).z) 10 | #define DIM3(TENSOR) ((TENSOR).w) 11 | 12 | #define DIM3_INDEX(TENSOR, xx, yy, zz, ww) ((TENSOR)[((xx) * (TENSOR##_stride.x)) + ((yy) * (TENSOR##_stride.y)) + ((zz) * (TENSOR##_stride.z)) + ((ww) * (TENSOR##_stride.w))]) 13 | 14 | 15 | #ifdef __cplusplus 16 | extern "C" { 17 | #endif 18 | 19 | __global__ void kernel_ChannelNorm_updateOutput(const int n, const float* input1, const long4 input1_size, const long4 input1_stride, float* output, const long4 output_size, const long4 output_stride, int norm_deg) { 20 | int index = blockIdx.x * blockDim.x + threadIdx.x; 21 | 22 | if (index >= n) { 23 | return; 24 | } 25 | 26 | int dim_b = DIM0(output_size); 27 | int dim_c = DIM1(output_size); 28 | int dim_h = DIM2(output_size); 29 | int dim_w = DIM3(output_size); 30 | int dim_chw = dim_c * dim_h * dim_w; 31 | 32 | int b = ( index / dim_chw ) % dim_b; 33 | int y = ( index / dim_w ) % dim_h; 34 | int x = ( index ) % dim_w; 35 | 36 | int i1dim_c = DIM1(input1_size); 37 | int i1dim_h = DIM2(input1_size); 38 | int i1dim_w = DIM3(input1_size); 39 | int i1dim_chw = i1dim_c * i1dim_h * i1dim_w; 40 | int i1dim_hw = i1dim_h * i1dim_w; 41 | 42 | float result = 0.0; 43 | 44 | for (int c = 0; c < i1dim_c; ++c) { 45 | int i1Index = b * i1dim_chw + c * i1dim_hw + y * i1dim_w + x; 46 | float val = input1[i1Index]; 47 | result += val * val; 48 | } 49 | result = sqrt(result); 50 | output[index] = result; 51 | } 52 | 53 | 54 | __global__ void kernel_ChannelNorm_backward_input1(const int n, const float* input1, const long4 input1_size, const long4 input1_stride, 55 | const float* output, const long4 output_size, const long4 output_stride, const float* gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride, 56 | float* gradInput, const long4 gradInput_size, const long4 gradInput_stride, int norm_deg) { 57 | int index = blockIdx.x * blockDim.x + threadIdx.x; 58 | 59 | if (index >= n) { 60 | return; 61 | } 62 | 63 | float val = 0.0; 64 | 65 | int dim_b = DIM0(gradInput_size); 66 | int dim_c = DIM1(gradInput_size); 67 | int dim_h = DIM2(gradInput_size); 68 | int dim_w = DIM3(gradInput_size); 69 | int dim_chw = dim_c * dim_h * dim_w; 70 | int dim_hw = dim_h * dim_w; 71 | 72 | int b = ( index / dim_chw ) % dim_b; 73 | int y = ( index / dim_w ) % dim_h; 74 | int x = ( index ) % dim_w; 75 | 76 | 77 | int outIndex = b * dim_hw + y * dim_w + x; 78 | val = gradOutput[outIndex] * input1[index] / (output[outIndex]+1e-9); 79 | gradInput[index] = val; 80 | 81 | } 82 | 83 | void ChannelNorm_kernel_forward(THCState* state, THCudaTensor* input1, THCudaTensor* output, int norm_deg) { 84 | int n = 0; 85 | 86 | const long4 input1_size = make_long4(input1->size[0], input1->size[1], input1->size[2], input1->size[3]); 87 | const long4 input1_stride = make_long4(input1->stride[0], input1->stride[1], input1->stride[2], input1->stride[3]); 88 | 89 | const long4 output_size = make_long4(output->size[0], output->size[1], output->size[2], output->size[3]); 90 | const long4 output_stride = make_long4(output->stride[0], output->stride[1], output->stride[2], output->stride[3]); 91 | 92 | n = THCudaTensor_nElement(state, output); 93 | kernel_ChannelNorm_updateOutput<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>( 94 | n, THCudaTensor_data(state, input1), input1_size, input1_stride, THCudaTensor_data(state, output), output_size, output_stride, 95 | norm_deg); 96 | 97 | THCudaCheck(cudaGetLastError()); 98 | } 99 | 100 | void ChannelNorm_kernel_backward(THCState* state, THCudaTensor* input1, THCudaTensor* output, THCudaTensor* gradOutput, THCudaTensor* gradInput1, int norm_deg) { 101 | int n = 0; 102 | 103 | const long4 input1_size = make_long4(input1->size[0], input1->size[1], input1->size[2], input1->size[3]); 104 | const long4 input1_stride = make_long4(input1->stride[0], input1->stride[1], input1->stride[2], input1->stride[3]); 105 | 106 | const long4 output_size = make_long4(output->size[0], output->size[1], output->size[2], output->size[3]); 107 | const long4 output_stride = make_long4(output->stride[0], output->stride[1], output->stride[2], output->stride[3]); 108 | 109 | const long4 gradOutput_size = make_long4(gradOutput->size[0], gradOutput->size[1], gradOutput->size[2], gradOutput->size[3]); 110 | const long4 gradOutput_stride = make_long4(gradOutput->stride[0], gradOutput->stride[1], gradOutput->stride[2], gradOutput->stride[3]); 111 | 112 | const long4 gradInput1_size = make_long4(gradInput1->size[0], gradInput1->size[1], gradInput1->size[2], gradInput1->size[3]); 113 | const long4 gradInput1_stride = make_long4(gradInput1->stride[0], gradInput1->stride[1], gradInput1->stride[2], gradInput1->stride[3]); 114 | 115 | n = THCudaTensor_nElement(state, gradInput1); 116 | kernel_ChannelNorm_backward_input1<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>( 117 | n, THCudaTensor_data(state, input1), input1_size, input1_stride, THCudaTensor_data(state, output), output_size, output_stride, 118 | THCudaTensor_data(state, gradOutput), gradOutput_size, gradOutput_stride, THCudaTensor_data(state, gradInput1), gradInput1_size, gradInput1_stride, 119 | norm_deg 120 | ); 121 | 122 | THCudaCheck(cudaGetLastError()); 123 | } 124 | 125 | #ifdef __cplusplus 126 | } 127 | #endif -------------------------------------------------------------------------------- /models/networks/channelnorm_package/src/ChannelNorm_kernel.h: -------------------------------------------------------------------------------- 1 | #ifdef __cplusplus 2 | extern "C" { 3 | #endif 4 | 5 | void ChannelNorm_kernel_forward(THCState* state, THCudaTensor* input1, THCudaTensor* output, int norm_deg); 6 | 7 | 8 | void ChannelNorm_kernel_backward(THCState* state, THCudaTensor* input1, THCudaTensor* output, THCudaTensor* gradOutput, THCudaTensor* gradInput1, int norm_deg); 9 | 10 | #ifdef __cplusplus 11 | } 12 | #endif -------------------------------------------------------------------------------- /models/networks/channelnorm_package/src/ChannelNorm_kernel.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/channelnorm_package/src/ChannelNorm_kernel.o -------------------------------------------------------------------------------- /models/networks/correlation_package/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/__pycache__/correlation.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/__pycache__/correlation.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/_ext/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/_ext/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/_ext/correlation/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from torch.utils.ffi import _wrap_function 3 | from ._correlation import lib as _lib, ffi as _ffi 4 | 5 | __all__ = [] 6 | def _import_symbols(locals): 7 | for symbol in dir(_lib): 8 | fn = getattr(_lib, symbol) 9 | if callable(fn): 10 | locals[symbol] = _wrap_function(fn, _ffi) 11 | else: 12 | locals[symbol] = fn 13 | __all__.append(symbol) 14 | 15 | _import_symbols(locals()) 16 | -------------------------------------------------------------------------------- /models/networks/correlation_package/_ext/correlation/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/_ext/correlation/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/_ext/correlation/_correlation.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/_ext/correlation/_correlation.so -------------------------------------------------------------------------------- /models/networks/correlation_package/build.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.utils.ffi 4 | 5 | this_folder = os.path.dirname(os.path.abspath(__file__)) + '/' 6 | 7 | Headers = [] 8 | Sources = [] 9 | Defines = [] 10 | Objects = [] 11 | 12 | if torch.cuda.is_available() == True: 13 | Headers += ['src/correlation_cuda.h'] 14 | Sources += ['src/correlation_cuda.c'] 15 | Defines += [('WITH_CUDA', None)] 16 | Objects += ['src/correlation_cuda_kernel.o'] 17 | 18 | ffi = torch.utils.ffi.create_extension( 19 | name='_ext.correlation', 20 | headers=Headers, 21 | sources=Sources, 22 | verbose=False, 23 | with_cuda=True, 24 | package=False, 25 | relative_to=this_folder, 26 | define_macros=Defines, 27 | extra_objects=[os.path.join(this_folder, Object) for Object in Objects] 28 | ) 29 | 30 | if __name__ == '__main__': 31 | ffi.build() -------------------------------------------------------------------------------- /models/networks/correlation_package/build/lib.linux-x86_64-3.7/correlation_cuda.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/build/lib.linux-x86_64-3.7/correlation_cuda.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /models/networks/correlation_package/build/temp.linux-x86_64-3.7/correlation_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/build/temp.linux-x86_64-3.7/correlation_cuda.o -------------------------------------------------------------------------------- /models/networks/correlation_package/build/temp.linux-x86_64-3.7/correlation_cuda_kernel.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/build/temp.linux-x86_64-3.7/correlation_cuda_kernel.o -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn.modules.module import Module 3 | from torch.autograd import Function 4 | import correlation_cuda 5 | 6 | class CorrelationFunction(Function): 7 | 8 | def __init__(self, pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1): 9 | super(CorrelationFunction, self).__init__() 10 | self.pad_size = pad_size 11 | self.kernel_size = kernel_size 12 | self.max_displacement = max_displacement 13 | self.stride1 = stride1 14 | self.stride2 = stride2 15 | self.corr_multiply = corr_multiply 16 | # self.out_channel = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1) 17 | 18 | def forward(self, input1, input2): 19 | self.save_for_backward(input1, input2) 20 | 21 | with torch.cuda.device_of(input1): 22 | rbot1 = input1.new() 23 | rbot2 = input2.new() 24 | output = input1.new() 25 | 26 | correlation_cuda.forward(input1, input2, rbot1, rbot2, output, 27 | self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply) 28 | 29 | return output 30 | 31 | def backward(self, grad_output): 32 | input1, input2 = self.saved_tensors 33 | 34 | with torch.cuda.device_of(input1): 35 | rbot1 = input1.new() 36 | rbot2 = input2.new() 37 | 38 | grad_input1 = input1.new() 39 | grad_input2 = input2.new() 40 | 41 | correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2, 42 | self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply) 43 | 44 | return grad_input1, grad_input2 45 | 46 | 47 | class Correlation(Module): 48 | def __init__(self, pad_size=0, kernel_size=0, max_displacement=0, stride1=1, stride2=2, corr_multiply=1): 49 | super(Correlation, self).__init__() 50 | self.pad_size = pad_size 51 | self.kernel_size = kernel_size 52 | self.max_displacement = max_displacement 53 | self.stride1 = stride1 54 | self.stride2 = stride2 55 | self.corr_multiply = corr_multiply 56 | 57 | def forward(self, input1, input2): 58 | 59 | result = CorrelationFunction(self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)(input1, input2) 60 | 61 | return result 62 | 63 | -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include "correlation_cuda_kernel.cuh" 9 | 10 | int correlation_forward_cuda(at::Tensor& input1, at::Tensor& input2, at::Tensor& rInput1, at::Tensor& rInput2, at::Tensor& output, 11 | int pad_size, 12 | int kernel_size, 13 | int max_displacement, 14 | int stride1, 15 | int stride2, 16 | int corr_type_multiply) 17 | { 18 | 19 | int batchSize = input1.size(0); 20 | 21 | int nInputChannels = input1.size(1); 22 | int inputHeight = input1.size(2); 23 | int inputWidth = input1.size(3); 24 | 25 | int kernel_radius = (kernel_size - 1) / 2; 26 | int border_radius = kernel_radius + max_displacement; 27 | 28 | int paddedInputHeight = inputHeight + 2 * pad_size; 29 | int paddedInputWidth = inputWidth + 2 * pad_size; 30 | 31 | int nOutputChannels = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1); 32 | 33 | int outputHeight = ceil(static_cast(paddedInputHeight - 2 * border_radius) / static_cast(stride1)); 34 | int outputwidth = ceil(static_cast(paddedInputWidth - 2 * border_radius) / static_cast(stride1)); 35 | 36 | rInput1.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 37 | rInput2.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 38 | output.resize_({batchSize, nOutputChannels, outputHeight, outputwidth}); 39 | 40 | rInput1.fill_(0); 41 | rInput2.fill_(0); 42 | output.fill_(0); 43 | 44 | int success = correlation_forward_cuda_kernel( 45 | output, 46 | output.size(0), 47 | output.size(1), 48 | output.size(2), 49 | output.size(3), 50 | output.stride(0), 51 | output.stride(1), 52 | output.stride(2), 53 | output.stride(3), 54 | input1, 55 | input1.size(1), 56 | input1.size(2), 57 | input1.size(3), 58 | input1.stride(0), 59 | input1.stride(1), 60 | input1.stride(2), 61 | input1.stride(3), 62 | input2, 63 | input2.size(1), 64 | input2.stride(0), 65 | input2.stride(1), 66 | input2.stride(2), 67 | input2.stride(3), 68 | rInput1, 69 | rInput2, 70 | pad_size, 71 | kernel_size, 72 | max_displacement, 73 | stride1, 74 | stride2, 75 | corr_type_multiply, 76 | at::cuda::getCurrentCUDAStream() 77 | //at::globalContext().getCurrentCUDAStream() 78 | ); 79 | 80 | //check for errors 81 | if (!success) { 82 | AT_ERROR("CUDA call failed"); 83 | } 84 | 85 | return 1; 86 | 87 | } 88 | 89 | int correlation_backward_cuda(at::Tensor& input1, at::Tensor& input2, at::Tensor& rInput1, at::Tensor& rInput2, at::Tensor& gradOutput, 90 | at::Tensor& gradInput1, at::Tensor& gradInput2, 91 | int pad_size, 92 | int kernel_size, 93 | int max_displacement, 94 | int stride1, 95 | int stride2, 96 | int corr_type_multiply) 97 | { 98 | 99 | int batchSize = input1.size(0); 100 | int nInputChannels = input1.size(1); 101 | int paddedInputHeight = input1.size(2)+ 2 * pad_size; 102 | int paddedInputWidth = input1.size(3)+ 2 * pad_size; 103 | 104 | int height = input1.size(2); 105 | int width = input1.size(3); 106 | 107 | rInput1.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 108 | rInput2.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 109 | gradInput1.resize_({batchSize, nInputChannels, height, width}); 110 | gradInput2.resize_({batchSize, nInputChannels, height, width}); 111 | 112 | rInput1.fill_(0); 113 | rInput2.fill_(0); 114 | gradInput1.fill_(0); 115 | gradInput2.fill_(0); 116 | 117 | int success = correlation_backward_cuda_kernel(gradOutput, 118 | gradOutput.size(0), 119 | gradOutput.size(1), 120 | gradOutput.size(2), 121 | gradOutput.size(3), 122 | gradOutput.stride(0), 123 | gradOutput.stride(1), 124 | gradOutput.stride(2), 125 | gradOutput.stride(3), 126 | input1, 127 | input1.size(1), 128 | input1.size(2), 129 | input1.size(3), 130 | input1.stride(0), 131 | input1.stride(1), 132 | input1.stride(2), 133 | input1.stride(3), 134 | input2, 135 | input2.stride(0), 136 | input2.stride(1), 137 | input2.stride(2), 138 | input2.stride(3), 139 | gradInput1, 140 | gradInput1.stride(0), 141 | gradInput1.stride(1), 142 | gradInput1.stride(2), 143 | gradInput1.stride(3), 144 | gradInput2, 145 | gradInput2.size(1), 146 | gradInput2.stride(0), 147 | gradInput2.stride(1), 148 | gradInput2.stride(2), 149 | gradInput2.stride(3), 150 | rInput1, 151 | rInput2, 152 | pad_size, 153 | kernel_size, 154 | max_displacement, 155 | stride1, 156 | stride2, 157 | corr_type_multiply, 158 | at::cuda::getCurrentCUDAStream() 159 | //at::globalContext().getCurrentCUDAStream() 160 | ); 161 | 162 | if (!success) { 163 | AT_ERROR("CUDA call failed"); 164 | } 165 | 166 | return 1; 167 | } 168 | 169 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 170 | m.def("forward", &correlation_forward_cuda, "Correlation forward (CUDA)"); 171 | m.def("backward", &correlation_backward_cuda, "Correlation backward (CUDA)"); 172 | } 173 | 174 | -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation_cuda.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: correlation-cuda 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation_cuda.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | correlation_cuda.cc 2 | correlation_cuda_kernel.cu 3 | setup.py 4 | correlation_cuda.egg-info/PKG-INFO 5 | correlation_cuda.egg-info/SOURCES.txt 6 | correlation_cuda.egg-info/dependency_links.txt 7 | correlation_cuda.egg-info/top_level.txt -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation_cuda.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation_cuda.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | correlation_cuda 2 | -------------------------------------------------------------------------------- /models/networks/correlation_package/correlation_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int correlation_forward_cuda_kernel(at::Tensor& output, 8 | int ob, 9 | int oc, 10 | int oh, 11 | int ow, 12 | int osb, 13 | int osc, 14 | int osh, 15 | int osw, 16 | 17 | at::Tensor& input1, 18 | int ic, 19 | int ih, 20 | int iw, 21 | int isb, 22 | int isc, 23 | int ish, 24 | int isw, 25 | 26 | at::Tensor& input2, 27 | int gc, 28 | int gsb, 29 | int gsc, 30 | int gsh, 31 | int gsw, 32 | 33 | at::Tensor& rInput1, 34 | at::Tensor& rInput2, 35 | int pad_size, 36 | int kernel_size, 37 | int max_displacement, 38 | int stride1, 39 | int stride2, 40 | int corr_type_multiply, 41 | cudaStream_t stream); 42 | 43 | 44 | int correlation_backward_cuda_kernel( 45 | at::Tensor& gradOutput, 46 | int gob, 47 | int goc, 48 | int goh, 49 | int gow, 50 | int gosb, 51 | int gosc, 52 | int gosh, 53 | int gosw, 54 | 55 | at::Tensor& input1, 56 | int ic, 57 | int ih, 58 | int iw, 59 | int isb, 60 | int isc, 61 | int ish, 62 | int isw, 63 | 64 | at::Tensor& input2, 65 | int gsb, 66 | int gsc, 67 | int gsh, 68 | int gsw, 69 | 70 | at::Tensor& gradInput1, 71 | int gisb, 72 | int gisc, 73 | int gish, 74 | int gisw, 75 | 76 | at::Tensor& gradInput2, 77 | int ggc, 78 | int ggsb, 79 | int ggsc, 80 | int ggsh, 81 | int ggsw, 82 | 83 | at::Tensor& rInput1, 84 | at::Tensor& rInput2, 85 | int pad_size, 86 | int kernel_size, 87 | int max_displacement, 88 | int stride1, 89 | int stride2, 90 | int corr_type_multiply, 91 | cudaStream_t stream); 92 | -------------------------------------------------------------------------------- /models/networks/correlation_package/dist/correlation_cuda-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/dist/correlation_cuda-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /models/networks/correlation_package/functions/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/functions/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/functions/__pycache__/correlation.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/functions/__pycache__/correlation.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/functions/correlation.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function, Variable 2 | from .._ext import correlation 3 | 4 | 5 | class CorrelationFunction(Function): 6 | 7 | @staticmethod 8 | def forward(ctx, 9 | input1, 10 | input2, 11 | pad_size=3, 12 | kernel_size=3, 13 | max_displacement=20, 14 | stride1=1, 15 | stride2=2, 16 | corr_multiply=1): 17 | assert input1.is_contiguous() 18 | assert input2.is_contiguous() 19 | 20 | ctx.save_for_backward(input1, input2) 21 | ctx.pad_size = pad_size 22 | ctx.kernel_size = kernel_size 23 | ctx.max_displacement = max_displacement 24 | ctx.stride1 = stride1 25 | ctx.stride2 = stride2 26 | ctx.corr_multiply = corr_multiply 27 | 28 | rbot1 = input1.new() 29 | rbot2 = input2.new() 30 | output = input1.new() 31 | 32 | correlation.Correlation_forward_cuda( 33 | input1, input2, rbot1, rbot2, output, pad_size, kernel_size, 34 | max_displacement, stride1, stride2, corr_multiply) 35 | 36 | return output 37 | 38 | @staticmethod 39 | def backward(ctx, grad_output): 40 | assert grad_output.is_contiguous() 41 | 42 | input1, input2 = ctx.saved_tensors 43 | 44 | rbot1 = input1.new() 45 | rbot2 = input2.new() 46 | 47 | grad_input1 = Variable(input1.new()) 48 | grad_input2 = Variable(input2.new()) 49 | 50 | correlation.Correlation_backward_cuda( 51 | input1, input2, rbot1, rbot2, grad_output.data, grad_input1.data, 52 | grad_input2.data, ctx.pad_size, ctx.kernel_size, 53 | ctx.max_displacement, ctx.stride1, ctx.stride2, ctx.corr_multiply) 54 | 55 | return (grad_input1, grad_input2) + (None, ) * 6 56 | -------------------------------------------------------------------------------- /models/networks/correlation_package/make.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))") 3 | 4 | cd src 5 | 6 | echo "Compiling correlation kernels by nvcc..." 7 | 8 | rm correlation_cuda_kernel.o 9 | rm -r ../_ext 10 | 11 | nvcc -c -o correlation_cuda_kernel.o correlation_cuda_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52 12 | 13 | cd ../ 14 | python build.py 15 | -------------------------------------------------------------------------------- /models/networks/correlation_package/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/modules/__pycache__/correlation.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/modules/__pycache__/correlation.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/correlation_package/modules/correlation.py: -------------------------------------------------------------------------------- 1 | from torch.nn.modules.module import Module 2 | 3 | from ..functions.correlation import CorrelationFunction 4 | 5 | 6 | class Correlation(Module): 7 | 8 | def __init__(self, 9 | pad_size=0, 10 | kernel_size=0, 11 | max_displacement=0, 12 | stride1=1, 13 | stride2=2, 14 | corr_multiply=1): 15 | super(Correlation, self).__init__() 16 | self.pad_size = pad_size 17 | self.kernel_size = kernel_size 18 | self.max_displacement = max_displacement 19 | self.stride1 = stride1 20 | self.stride2 = stride2 21 | self.corr_multiply = corr_multiply 22 | 23 | def forward(self, input1, input2): 24 | return CorrelationFunction.apply(input1, input2, self.pad_size, 25 | self.kernel_size, 26 | self.max_displacement, self.stride1, 27 | self.stride2, self.corr_multiply) 28 | -------------------------------------------------------------------------------- /models/networks/correlation_package/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import os 3 | import torch 4 | 5 | from setuptools import setup, find_packages 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | cxx_args = ['-std=c++11'] 9 | 10 | nvcc_args = [ 11 | '-gencode', 'arch=compute_50,code=sm_50', 12 | '-gencode', 'arch=compute_52,code=sm_52', 13 | '-gencode', 'arch=compute_60,code=sm_60', 14 | '-gencode', 'arch=compute_61,code=sm_61', 15 | '-gencode', 'arch=compute_70,code=sm_70', 16 | '-gencode', 'arch=compute_70,code=compute_70' 17 | ] 18 | 19 | setup( 20 | name='correlation_cuda', 21 | ext_modules=[ 22 | CUDAExtension('correlation_cuda', [ 23 | 'correlation_cuda.cc', 24 | 'correlation_cuda_kernel.cu' 25 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 26 | ], 27 | cmdclass={ 28 | 'build_ext': BuildExtension 29 | }) 30 | -------------------------------------------------------------------------------- /models/networks/correlation_package/src/correlation.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | int Correlation_forward_cpu(THFloatTensor *input1, 4 | THFloatTensor *input2, 5 | THFloatTensor *rInput1, 6 | THFloatTensor *rInput2, 7 | THFloatTensor *output, 8 | int pad_size, 9 | int kernel_size, 10 | int max_displacement, 11 | int stride1, 12 | int stride2, 13 | int corr_type_multiply) 14 | { 15 | return 1; 16 | } 17 | 18 | int Correlation_backward_cpu(THFloatTensor *input1, 19 | THFloatTensor *input2, 20 | THFloatTensor *rInput1, 21 | THFloatTensor *rInput2, 22 | THFloatTensor *gradOutput, 23 | THFloatTensor *gradInput1, 24 | THFloatTensor *gradInput2, 25 | int pad_size, 26 | int kernel_size, 27 | int max_displacement, 28 | int stride1, 29 | int stride2, 30 | int corr_type_multiply) 31 | { 32 | return 1; 33 | } 34 | -------------------------------------------------------------------------------- /models/networks/correlation_package/src/correlation.h: -------------------------------------------------------------------------------- 1 | int Correlation_forward_cpu(THFloatTensor *input1, 2 | THFloatTensor *input2, 3 | THFloatTensor *rInput1, 4 | THFloatTensor *rInput2, 5 | THFloatTensor *output, 6 | int pad_size, 7 | int kernel_size, 8 | int max_displacement, 9 | int stride1, 10 | int stride2, 11 | int corr_type_multiply); 12 | 13 | int Correlation_backward_cpu(THFloatTensor *input1, 14 | THFloatTensor *input2, 15 | THFloatTensor *rInput1, 16 | THFloatTensor *rInput2, 17 | THFloatTensor *gradOutput, 18 | THFloatTensor *gradInput1, 19 | THFloatTensor *gradInput2, 20 | int pad_size, 21 | int kernel_size, 22 | int max_displacement, 23 | int stride1, 24 | int stride2, 25 | int corr_type_multiply); 26 | -------------------------------------------------------------------------------- /models/networks/correlation_package/src/correlation_cuda.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "correlation_cuda_kernel.h" 5 | 6 | #define real float 7 | 8 | // symbol to be automatically resolved by PyTorch libs 9 | extern THCState *state; 10 | 11 | int Correlation_forward_cuda(THCudaTensor *input1, THCudaTensor *input2, THCudaTensor *rInput1, THCudaTensor *rInput2, THCudaTensor *output, 12 | int pad_size, 13 | int kernel_size, 14 | int max_displacement, 15 | int stride1, 16 | int stride2, 17 | int corr_type_multiply) 18 | { 19 | 20 | int batchSize = input1->size[0]; 21 | int nInputChannels = input1->size[1]; 22 | int inputHeight = input1->size[2]; 23 | int inputWidth = input1->size[3]; 24 | 25 | int kernel_radius = (kernel_size - 1) / 2; 26 | int border_radius = kernel_radius + max_displacement; 27 | 28 | int paddedInputHeight = inputHeight + 2 * pad_size; 29 | int paddedInputWidth = inputWidth + 2 * pad_size; 30 | 31 | int nOutputChannels = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1); 32 | 33 | int outputHeight = ceil((float)(paddedInputHeight - 2 * border_radius) / (float)stride1); 34 | int outputwidth = ceil((float)(paddedInputWidth - 2 * border_radius) / (float)stride1); 35 | 36 | THCudaTensor_resize4d(state, rInput1, batchSize, paddedInputHeight, paddedInputWidth, nInputChannels); 37 | THCudaTensor_resize4d(state, rInput2, batchSize, paddedInputHeight, paddedInputWidth, nInputChannels); 38 | THCudaTensor_resize4d(state, output, batchSize, nOutputChannels, outputHeight, outputwidth); 39 | 40 | THCudaTensor_fill(state, rInput1, 0); 41 | THCudaTensor_fill(state, rInput2, 0); 42 | THCudaTensor_fill(state, output, 0); 43 | 44 | int success = 0; 45 | success = Correlation_forward_cuda_kernel( THCudaTensor_data(state, output), 46 | THCudaTensor_size(state, output, 0), 47 | THCudaTensor_size(state, output, 1), 48 | THCudaTensor_size(state, output, 2), 49 | THCudaTensor_size(state, output, 3), 50 | THCudaTensor_stride(state, output, 0), 51 | THCudaTensor_stride(state, output, 1), 52 | THCudaTensor_stride(state, output, 2), 53 | THCudaTensor_stride(state, output, 3), 54 | 55 | THCudaTensor_data(state, input1), 56 | THCudaTensor_size(state, input1, 1), 57 | THCudaTensor_size(state, input1, 2), 58 | THCudaTensor_size(state, input1, 3), 59 | THCudaTensor_stride(state, input1, 0), 60 | THCudaTensor_stride(state, input1, 1), 61 | THCudaTensor_stride(state, input1, 2), 62 | THCudaTensor_stride(state, input1, 3), 63 | 64 | THCudaTensor_data(state, input2), 65 | THCudaTensor_size(state, input2, 1), 66 | THCudaTensor_stride(state, input2, 0), 67 | THCudaTensor_stride(state, input2, 1), 68 | THCudaTensor_stride(state, input2, 2), 69 | THCudaTensor_stride(state, input2, 3), 70 | 71 | THCudaTensor_data(state, rInput1), 72 | THCudaTensor_data(state, rInput2), 73 | 74 | pad_size, 75 | kernel_size, 76 | max_displacement, 77 | stride1, 78 | stride2, 79 | corr_type_multiply, 80 | 81 | THCState_getCurrentStream(state)); 82 | 83 | THCudaTensor_free(state, rInput1); 84 | THCudaTensor_free(state, rInput2); 85 | 86 | //check for errors 87 | if (!success) { 88 | THError("aborting"); 89 | } 90 | 91 | return 1; 92 | 93 | } 94 | 95 | int Correlation_backward_cuda(THCudaTensor *input1, THCudaTensor *input2, THCudaTensor *rInput1, THCudaTensor *rInput2, THCudaTensor *gradOutput, 96 | THCudaTensor *gradInput1, THCudaTensor *gradInput2, 97 | int pad_size, 98 | int kernel_size, 99 | int max_displacement, 100 | int stride1, 101 | int stride2, 102 | int corr_type_multiply) 103 | { 104 | 105 | int batchSize = input1->size[0]; 106 | int nInputChannels = input1->size[1]; 107 | int paddedInputHeight = input1->size[2]+ 2 * pad_size; 108 | int paddedInputWidth = input1->size[3]+ 2 * pad_size; 109 | 110 | int height = input1->size[2]; 111 | int width = input1->size[3]; 112 | 113 | THCudaTensor_resize4d(state, rInput1, batchSize, paddedInputHeight, paddedInputWidth, nInputChannels); 114 | THCudaTensor_resize4d(state, rInput2, batchSize, paddedInputHeight, paddedInputWidth, nInputChannels); 115 | THCudaTensor_resize4d(state, gradInput1, batchSize, nInputChannels, height, width); 116 | THCudaTensor_resize4d(state, gradInput2, batchSize, nInputChannels, height, width); 117 | 118 | THCudaTensor_fill(state, rInput1, 0); 119 | THCudaTensor_fill(state, rInput2, 0); 120 | THCudaTensor_fill(state, gradInput1, 0); 121 | THCudaTensor_fill(state, gradInput2, 0); 122 | 123 | int success = 0; 124 | success = Correlation_backward_cuda_kernel( 125 | THCudaTensor_data(state, gradOutput), 126 | THCudaTensor_size(state, gradOutput, 0), 127 | THCudaTensor_size(state, gradOutput, 1), 128 | THCudaTensor_size(state, gradOutput, 2), 129 | THCudaTensor_size(state, gradOutput, 3), 130 | THCudaTensor_stride(state, gradOutput, 0), 131 | THCudaTensor_stride(state, gradOutput, 1), 132 | THCudaTensor_stride(state, gradOutput, 2), 133 | THCudaTensor_stride(state, gradOutput, 3), 134 | 135 | THCudaTensor_data(state, input1), 136 | THCudaTensor_size(state, input1, 1), 137 | THCudaTensor_size(state, input1, 2), 138 | THCudaTensor_size(state, input1, 3), 139 | THCudaTensor_stride(state, input1, 0), 140 | THCudaTensor_stride(state, input1, 1), 141 | THCudaTensor_stride(state, input1, 2), 142 | THCudaTensor_stride(state, input1, 3), 143 | 144 | THCudaTensor_data(state, input2), 145 | THCudaTensor_stride(state, input2, 0), 146 | THCudaTensor_stride(state, input2, 1), 147 | THCudaTensor_stride(state, input2, 2), 148 | THCudaTensor_stride(state, input2, 3), 149 | 150 | THCudaTensor_data(state, gradInput1), 151 | THCudaTensor_stride(state, gradInput1, 0), 152 | THCudaTensor_stride(state, gradInput1, 1), 153 | THCudaTensor_stride(state, gradInput1, 2), 154 | THCudaTensor_stride(state, gradInput1, 3), 155 | 156 | THCudaTensor_data(state, gradInput2), 157 | THCudaTensor_size(state, gradInput2, 1), 158 | THCudaTensor_stride(state, gradInput2, 0), 159 | THCudaTensor_stride(state, gradInput2, 1), 160 | THCudaTensor_stride(state, gradInput2, 2), 161 | THCudaTensor_stride(state, gradInput2, 3), 162 | 163 | THCudaTensor_data(state, rInput1), 164 | THCudaTensor_data(state, rInput2), 165 | pad_size, 166 | kernel_size, 167 | max_displacement, 168 | stride1, 169 | stride2, 170 | corr_type_multiply, 171 | THCState_getCurrentStream(state)); 172 | 173 | THCudaTensor_free(state, rInput1); 174 | THCudaTensor_free(state, rInput2); 175 | 176 | if (!success) { 177 | THError("aborting"); 178 | } 179 | return 1; 180 | } 181 | -------------------------------------------------------------------------------- /models/networks/correlation_package/src/correlation_cuda.h: -------------------------------------------------------------------------------- 1 | int Correlation_forward_cuda(THCudaTensor *input1, THCudaTensor *input2, THCudaTensor *rInput1, THCudaTensor *rInput2, 2 | THCudaTensor *output, 3 | int pad_size, 4 | int kernel_size, 5 | int max_displacement, 6 | int stride1, 7 | int stride2, 8 | int corr_type_multiply); 9 | 10 | int Correlation_backward_cuda(THCudaTensor *input1, THCudaTensor *input2, THCudaTensor *rInput1, THCudaTensor *rInput2, 11 | THCudaTensor *gradOutput, THCudaTensor *gradInput1, THCudaTensor *gradInput2, 12 | int pad_size, 13 | int kernel_size, 14 | int max_displacement, 15 | int stride1, 16 | int stride2, 17 | int corr_type_multiply); 18 | 19 | -------------------------------------------------------------------------------- /models/networks/correlation_package/src/correlation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifdef __cplusplus 2 | extern "C" { 3 | #endif 4 | 5 | int Correlation_forward_cuda_kernel(/*THCudaTensor_data(state, output)*/ float *output, 6 | /*THCudaTensor_size(state, output, 0)*/ int ob, 7 | /*THCudaTensor_size(state, output, 1)*/ int oc, 8 | /*THCudaTensor_size(state, output, 2)*/ int oh, 9 | /*THCudaTensor_size(state, output, 3)*/ int ow, 10 | /*THCudaTensor_stride(state, output, 0)*/ int osb, 11 | /*THCudaTensor_stride(state, output, 1)*/ int osc, 12 | /*THCudaTensor_stride(state, output, 2)*/ int osh, 13 | /*THCudaTensor_stride(state, output, 3)*/ int osw, 14 | 15 | /*THCudaTensor_data(state, input1)*/ float *input1, 16 | /*THCudaTensor_size(state, input1, 1)*/ int ic, 17 | /*THCudaTensor_size(state, input1, 2)*/ int ih, 18 | /*THCudaTensor_size(state, input1, 3)*/ int iw, 19 | /*THCudaTensor_stride(state, input1, 0)*/ int isb, 20 | /*THCudaTensor_stride(state, input1, 1)*/ int isc, 21 | /*THCudaTensor_stride(state, input1, 2)*/ int ish, 22 | /*THCudaTensor_stride(state, input1, 3)*/ int isw, 23 | 24 | /*THCudaTensor_data(state, input2)*/ float *input2, 25 | /*THCudaTensor_size(state, input2, 1)*/ int gc, 26 | /*THCudaTensor_stride(state, input2, 0)*/ int gsb, 27 | /*THCudaTensor_stride(state, input2, 1)*/ int gsc, 28 | /*THCudaTensor_stride(state, input2, 2)*/ int gsh, 29 | /*THCudaTensor_stride(state, input2, 3)*/ int gsw, 30 | 31 | /*THCudaTensor_data(state, rInput1)*/ float *rInput1, 32 | /*THCudaTensor_data(state, rInput2)*/ float *rInput2, 33 | int pad_size, 34 | int kernel_size, 35 | int max_displacement, 36 | int stride1, 37 | int stride2, 38 | int corr_type_multiply, 39 | /*THCState_getCurrentStream(state)*/ cudaStream_t stream); 40 | 41 | int Correlation_backward_cuda_kernel( 42 | /*THCudaTensor_data(state, gradOutput)*/ float *gradOutput, 43 | /*THCudaTensor_size(state, gradOutput, 0)*/ int gob, 44 | /*THCudaTensor_size(state, gradOutput, 1)*/ int goc, 45 | /*THCudaTensor_size(state, gradOutput, 2)*/ int goh, 46 | /*THCudaTensor_size(state, gradOutput, 3)*/ int gow, 47 | /*THCudaTensor_stride(state, gradOutput, 0)*/ int gosb, 48 | /*THCudaTensor_stride(state, gradOutput, 1)*/ int gosc, 49 | /*THCudaTensor_stride(state, gradOutput, 2)*/ int gosh, 50 | /*THCudaTensor_stride(state, gradOutput, 3)*/ int gosw, 51 | 52 | /*THCudaTensor_data(state, input1)*/ float* input1, 53 | /*THCudaTensor_size(state, input1, 1)*/ int ic, 54 | /*THCudaTensor_size(state, input1, 2)*/ int ih, 55 | /*THCudaTensor_size(state, input1, 3)*/ int iw, 56 | /*THCudaTensor_stride(state, input1, 0)*/ int isb, 57 | /*THCudaTensor_stride(state, input1, 1)*/ int isc, 58 | /*THCudaTensor_stride(state, input1, 2)*/ int ish, 59 | /*THCudaTensor_stride(state, input1, 3)*/ int isw, 60 | 61 | /*THCudaTensor_data(state, input2)*/ float *input2, 62 | /*THCudaTensor_stride(state, input2, 0)*/ int gsb, 63 | /*THCudaTensor_stride(state, input2, 1)*/ int gsc, 64 | /*THCudaTensor_stride(state, input2, 2)*/ int gsh, 65 | /*THCudaTensor_stride(state, input2, 3)*/ int gsw, 66 | 67 | /*THCudaTensor_data(state, gradInput1)*/ float *gradInput1, 68 | /*THCudaTensor_stride(state, gradInput1, 0)*/ int gisb, 69 | /*THCudaTensor_stride(state, gradInput1, 1)*/ int gisc, 70 | /*THCudaTensor_stride(state, gradInput1, 2)*/ int gish, 71 | /*THCudaTensor_stride(state, gradInput1, 3)*/ int gisw, 72 | 73 | /*THCudaTensor_data(state, gradInput2)*/ float *gradInput2, 74 | /*THCudaTensor_size(state, gradInput2, 1)*/ int ggc, 75 | /*THCudaTensor_stride(state, gradInput2, 0)*/ int ggsb, 76 | /*THCudaTensor_stride(state, gradInput2, 1)*/ int ggsc, 77 | /*THCudaTensor_stride(state, gradInput2, 2)*/ int ggsh, 78 | /*THCudaTensor_stride(state, gradInput2, 3)*/ int ggsw, 79 | 80 | /*THCudaTensor_data(state, rInput1)*/ float *rInput1, 81 | /*THCudaTensor_data(state, rInput2)*/ float *rInput2, 82 | int pad_size, 83 | int kernel_size, 84 | int max_displacement, 85 | int stride1, 86 | int stride2, 87 | int corr_type_multiply, 88 | /*THCState_getCurrentStream(state)*/cudaStream_t stream); 89 | 90 | #ifdef __cplusplus 91 | } 92 | #endif 93 | -------------------------------------------------------------------------------- /models/networks/correlation_package/src/correlation_cuda_kernel.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/correlation_package/src/correlation_cuda_kernel.o -------------------------------------------------------------------------------- /models/networks/resample2d_package/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/__pycache__/resample2d.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/__pycache__/resample2d.cpython-37.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/_ext/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/_ext/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/_ext/resample2d/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from torch.utils.ffi import _wrap_function 3 | from ._resample2d import lib as _lib, ffi as _ffi 4 | 5 | __all__ = [] 6 | def _import_symbols(locals): 7 | for symbol in dir(_lib): 8 | fn = getattr(_lib, symbol) 9 | if callable(fn): 10 | locals[symbol] = _wrap_function(fn, _ffi) 11 | else: 12 | locals[symbol] = fn 13 | __all__.append(symbol) 14 | 15 | _import_symbols(locals()) 16 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/_ext/resample2d/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/_ext/resample2d/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/_ext/resample2d/_resample2d.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/_ext/resample2d/_resample2d.so -------------------------------------------------------------------------------- /models/networks/resample2d_package/build.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.utils.ffi 4 | 5 | this_folder = os.path.dirname(os.path.abspath(__file__)) + '/' 6 | 7 | Headers = [] 8 | Sources = [] 9 | Defines = [] 10 | Objects = [] 11 | 12 | if torch.cuda.is_available() == True: 13 | Headers += ['src/Resample2d_cuda.h'] 14 | Sources += ['src/Resample2d_cuda.c'] 15 | Defines += [('WITH_CUDA', None)] 16 | Objects += ['src/Resample2d_kernel.o'] 17 | 18 | ffi = torch.utils.ffi.create_extension( 19 | name='_ext.resample2d', 20 | headers=Headers, 21 | sources=Sources, 22 | verbose=False, 23 | with_cuda=True, 24 | package=False, 25 | relative_to=this_folder, 26 | define_macros=Defines, 27 | extra_objects=[os.path.join(this_folder, Object) for Object in Objects] 28 | ) 29 | 30 | if __name__ == '__main__': 31 | ffi.build() -------------------------------------------------------------------------------- /models/networks/resample2d_package/build/lib.linux-x86_64-3.7/resample2d_cuda.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/build/lib.linux-x86_64-3.7/resample2d_cuda.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /models/networks/resample2d_package/build/temp.linux-x86_64-3.7/resample2d_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/build/temp.linux-x86_64-3.7/resample2d_cuda.o -------------------------------------------------------------------------------- /models/networks/resample2d_package/build/temp.linux-x86_64-3.7/resample2d_kernel.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/build/temp.linux-x86_64-3.7/resample2d_kernel.o -------------------------------------------------------------------------------- /models/networks/resample2d_package/dist/resample2d_cuda-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/dist/resample2d_cuda-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /models/networks/resample2d_package/functions/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/functions/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/functions/__pycache__/resample2d.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/functions/__pycache__/resample2d.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/functions/resample2d.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function, Variable 2 | from .._ext import resample2d 3 | 4 | 5 | class Resample2dFunction(Function): 6 | 7 | @staticmethod 8 | def forward(ctx, input1, input2, kernel_size=1): 9 | assert input1.is_contiguous() 10 | assert input2.is_contiguous() 11 | 12 | ctx.save_for_backward(input1, input2) 13 | ctx.kernel_size = kernel_size 14 | 15 | _, d, _, _ = input1.size() 16 | b, _, h, w = input2.size() 17 | output = input1.new(b, d, h, w).zero_() 18 | #print("what the hell") 19 | resample2d.Resample2d_cuda_forward(input1, input2, output, kernel_size) 20 | 21 | return output 22 | 23 | @staticmethod 24 | def backward(ctx, grad_output): 25 | assert grad_output.is_contiguous() 26 | 27 | input1, input2 = ctx.saved_tensors 28 | 29 | grad_input1 = Variable(input1.new(input1.size()).zero_()) 30 | grad_input2 = Variable(input1.new(input2.size()).zero_()) 31 | 32 | resample2d.Resample2d_cuda_backward(input1, input2, grad_output.data, 33 | grad_input1.data, grad_input2.data, 34 | ctx.kernel_size) 35 | #print("gradient", grad_input1.min(), grad_input1.max()) 36 | return grad_input1, grad_input2, None 37 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/make.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))") 3 | 4 | cd src 5 | echo "Compiling resample2d kernels by nvcc..." 6 | rm Resample2d_kernel.o 7 | rm -r ../_ext 8 | 9 | nvcc -c -o Resample2d_kernel.o Resample2d_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52 -I ${TORCH}/lib/include/TH -I ${TORCH}/lib/include/THC 10 | 11 | cd ../ 12 | python build.py -------------------------------------------------------------------------------- /models/networks/resample2d_package/modules/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/modules/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/modules/__pycache__/resample2d.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/modules/__pycache__/resample2d.cpython-36.pyc -------------------------------------------------------------------------------- /models/networks/resample2d_package/modules/resample2d.py: -------------------------------------------------------------------------------- 1 | from torch.nn.modules.module import Module 2 | 3 | from ..functions.resample2d import Resample2dFunction 4 | 5 | 6 | class Resample2d(Module): 7 | 8 | def __init__(self, kernel_size=1): 9 | super(Resample2d, self).__init__() 10 | self.kernel_size = kernel_size 11 | 12 | def forward(self, input1, input2): 13 | input1_c = input1.contiguous() 14 | return Resample2dFunction.apply(input1_c, input2, self.kernel_size) 15 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d.py: -------------------------------------------------------------------------------- 1 | from torch.nn.modules.module import Module 2 | from torch.autograd import Function, Variable 3 | import resample2d_cuda 4 | 5 | class Resample2dFunction(Function): 6 | 7 | @staticmethod 8 | def forward(ctx, input1, input2, kernel_size=1): 9 | assert input1.is_contiguous() 10 | assert input2.is_contiguous() 11 | 12 | ctx.save_for_backward(input1, input2) 13 | ctx.kernel_size = kernel_size 14 | 15 | _, d, _, _ = input1.size() 16 | b, _, h, w = input2.size() 17 | output = input1.new(b, d, h, w).zero_() 18 | 19 | resample2d_cuda.forward(input1, input2, output, kernel_size) 20 | 21 | return output 22 | 23 | @staticmethod 24 | def backward(ctx, grad_output): 25 | assert grad_output.is_contiguous() 26 | 27 | input1, input2 = ctx.saved_tensors 28 | 29 | grad_input1 = Variable(input1.new(input1.size()).zero_()) 30 | grad_input2 = Variable(input1.new(input2.size()).zero_()) 31 | 32 | resample2d_cuda.backward(input1, input2, grad_output.data, 33 | grad_input1.data, grad_input2.data, 34 | ctx.kernel_size) 35 | 36 | return grad_input1, grad_input2, None 37 | 38 | class Resample2d(Module): 39 | 40 | def __init__(self, kernel_size=1): 41 | super(Resample2d, self).__init__() 42 | self.kernel_size = kernel_size 43 | 44 | def forward(self, input1, input2): 45 | input1_c = input1.contiguous() 46 | return Resample2dFunction.apply(input1_c, input2, self.kernel_size) 47 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "resample2d_kernel.cuh" 5 | 6 | int resample2d_cuda_forward( 7 | at::Tensor& input1, 8 | at::Tensor& input2, 9 | at::Tensor& output, 10 | int kernel_size) { 11 | resample2d_kernel_forward(input1, input2, output, kernel_size); 12 | return 1; 13 | } 14 | 15 | int resample2d_cuda_backward( 16 | at::Tensor& input1, 17 | at::Tensor& input2, 18 | at::Tensor& gradOutput, 19 | at::Tensor& gradInput1, 20 | at::Tensor& gradInput2, 21 | int kernel_size) { 22 | resample2d_kernel_backward(input1, input2, gradOutput, gradInput1, gradInput2, kernel_size); 23 | return 1; 24 | } 25 | 26 | 27 | 28 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 29 | m.def("forward", &resample2d_cuda_forward, "Resample2D forward (CUDA)"); 30 | m.def("backward", &resample2d_cuda_backward, "Resample2D backward (CUDA)"); 31 | } 32 | 33 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_cuda.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: resample2d-cuda 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_cuda.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | resample2d_cuda.cc 2 | resample2d_kernel.cu 3 | setup.py 4 | resample2d_cuda.egg-info/PKG-INFO 5 | resample2d_cuda.egg-info/SOURCES.txt 6 | resample2d_cuda.egg-info/dependency_links.txt 7 | resample2d_cuda.egg-info/top_level.txt -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_cuda.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_cuda.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | resample2d_cuda 2 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #define CUDA_NUM_THREADS 512 6 | #define THREADS_PER_BLOCK 64 7 | 8 | #define DIM0(TENSOR) ((TENSOR).x) 9 | #define DIM1(TENSOR) ((TENSOR).y) 10 | #define DIM2(TENSOR) ((TENSOR).z) 11 | #define DIM3(TENSOR) ((TENSOR).w) 12 | 13 | #define DIM3_INDEX(TENSOR, xx, yy, zz, ww) ((TENSOR)[((xx) * (TENSOR##_stride.x)) + ((yy) * (TENSOR##_stride.y)) + ((zz) * (TENSOR##_stride.z)) + ((ww) * (TENSOR##_stride.w))]) 14 | 15 | template 16 | __global__ void kernel_resample2d_update_output(const int n, 17 | const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride, 18 | const scalar_t* __restrict__ input2, const long4 input2_size, const long4 input2_stride, 19 | scalar_t* __restrict__ output, const long4 output_size, const long4 output_stride, int kernel_size) { 20 | int index = blockIdx.x * blockDim.x + threadIdx.x; 21 | 22 | if (index >= n) { 23 | return; 24 | } 25 | 26 | scalar_t val = 0.0f; 27 | 28 | int dim_b = DIM0(output_size); 29 | int dim_c = DIM1(output_size); 30 | int dim_h = DIM2(output_size); 31 | int dim_w = DIM3(output_size); 32 | int dim_chw = dim_c * dim_h * dim_w; 33 | int dim_hw = dim_h * dim_w; 34 | 35 | int b = ( index / dim_chw ) % dim_b; 36 | int c = ( index / dim_hw ) % dim_c; 37 | int y = ( index / dim_w ) % dim_h; 38 | int x = ( index ) % dim_w; 39 | 40 | scalar_t dx = DIM3_INDEX(input2, b, 0, y, x); 41 | scalar_t dy = DIM3_INDEX(input2, b, 1, y, x); 42 | 43 | scalar_t xf = static_cast(x) + dx; 44 | scalar_t yf = static_cast(y) + dy; 45 | scalar_t alpha = xf - floor(xf); // alpha 46 | scalar_t beta = yf - floor(yf); // beta 47 | 48 | int xL = max(min( int (floor(xf)), dim_w-1), 0); 49 | int xR = max(min( int (floor(xf)+1), dim_w -1), 0); 50 | int yT = max(min( int (floor(yf)), dim_h-1), 0); 51 | int yB = max(min( int (floor(yf)+1), dim_h-1), 0); 52 | 53 | for (int fy = 0; fy < kernel_size; fy += 1) { 54 | for (int fx = 0; fx < kernel_size; fx += 1) { 55 | val += static_cast((1. - alpha)*(1. - beta) * DIM3_INDEX(input1, b, c, yT + fy, xL + fx)); 56 | val += static_cast((alpha)*(1. - beta) * DIM3_INDEX(input1, b, c, yT + fy, xR + fx)); 57 | val += static_cast((1. - alpha)*(beta) * DIM3_INDEX(input1, b, c, yB + fy, xL + fx)); 58 | val += static_cast((alpha)*(beta) * DIM3_INDEX(input1, b, c, yB + fy, xR + fx)); 59 | } 60 | } 61 | 62 | output[index] = val; 63 | 64 | } 65 | 66 | 67 | template 68 | __global__ void kernel_resample2d_backward_input1( 69 | const int n, const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride, 70 | const scalar_t* __restrict__ input2, const long4 input2_size, const long4 input2_stride, 71 | const scalar_t* __restrict__ gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride, 72 | scalar_t* __restrict__ gradInput, const long4 gradInput_size, const long4 gradInput_stride, int kernel_size) { 73 | 74 | int index = blockIdx.x * blockDim.x + threadIdx.x; 75 | 76 | if (index >= n) { 77 | return; 78 | } 79 | 80 | int dim_b = DIM0(gradOutput_size); 81 | int dim_c = DIM1(gradOutput_size); 82 | int dim_h = DIM2(gradOutput_size); 83 | int dim_w = DIM3(gradOutput_size); 84 | int dim_chw = dim_c * dim_h * dim_w; 85 | int dim_hw = dim_h * dim_w; 86 | 87 | int b = ( index / dim_chw ) % dim_b; 88 | int c = ( index / dim_hw ) % dim_c; 89 | int y = ( index / dim_w ) % dim_h; 90 | int x = ( index ) % dim_w; 91 | 92 | scalar_t dx = DIM3_INDEX(input2, b, 0, y, x); 93 | scalar_t dy = DIM3_INDEX(input2, b, 1, y, x); 94 | 95 | scalar_t xf = static_cast(x) + dx; 96 | scalar_t yf = static_cast(y) + dy; 97 | scalar_t alpha = xf - int(xf); // alpha 98 | scalar_t beta = yf - int(yf); // beta 99 | 100 | int idim_h = DIM2(input1_size); 101 | int idim_w = DIM3(input1_size); 102 | 103 | int xL = max(min( int (floor(xf)), idim_w-1), 0); 104 | int xR = max(min( int (floor(xf)+1), idim_w -1), 0); 105 | int yT = max(min( int (floor(yf)), idim_h-1), 0); 106 | int yB = max(min( int (floor(yf)+1), idim_h-1), 0); 107 | 108 | for (int fy = 0; fy < kernel_size; fy += 1) { 109 | for (int fx = 0; fx < kernel_size; fx += 1) { 110 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yT + fy), (xL + fx)), (1-alpha)*(1-beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 111 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yT + fy), (xR + fx)), (alpha)*(1-beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 112 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yB + fy), (xL + fx)), (1-alpha)*(beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 113 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yB + fy), (xR + fx)), (alpha)*(beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 114 | } 115 | } 116 | 117 | } 118 | 119 | template 120 | __global__ void kernel_resample2d_backward_input2( 121 | const int n, const scalar_t* __restrict__ input1, const long4 input1_size, const long4 input1_stride, 122 | const scalar_t* __restrict__ input2, const long4 input2_size, const long4 input2_stride, 123 | const scalar_t* __restrict__ gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride, 124 | scalar_t* __restrict__ gradInput, const long4 gradInput_size, const long4 gradInput_stride, int kernel_size) { 125 | 126 | int index = blockIdx.x * blockDim.x + threadIdx.x; 127 | 128 | if (index >= n) { 129 | return; 130 | } 131 | 132 | scalar_t output = 0.0; 133 | int kernel_rad = (kernel_size - 1)/2; 134 | 135 | int dim_b = DIM0(gradInput_size); 136 | int dim_c = DIM1(gradInput_size); 137 | int dim_h = DIM2(gradInput_size); 138 | int dim_w = DIM3(gradInput_size); 139 | int dim_chw = dim_c * dim_h * dim_w; 140 | int dim_hw = dim_h * dim_w; 141 | 142 | int b = ( index / dim_chw ) % dim_b; 143 | int c = ( index / dim_hw ) % dim_c; 144 | int y = ( index / dim_w ) % dim_h; 145 | int x = ( index ) % dim_w; 146 | 147 | int odim_c = DIM1(gradOutput_size); 148 | 149 | scalar_t dx = DIM3_INDEX(input2, b, 0, y, x); 150 | scalar_t dy = DIM3_INDEX(input2, b, 1, y, x); 151 | 152 | scalar_t xf = static_cast(x) + dx; 153 | scalar_t yf = static_cast(y) + dy; 154 | 155 | int xL = max(min( int (floor(xf)), dim_w-1), 0); 156 | int xR = max(min( int (floor(xf)+1), dim_w -1), 0); 157 | int yT = max(min( int (floor(yf)), dim_h-1), 0); 158 | int yB = max(min( int (floor(yf)+1), dim_h-1), 0); 159 | 160 | if (c % 2) { 161 | float gamma = 1 - (xf - floor(xf)); // alpha 162 | for (int i = 0; i <= 2*kernel_rad; ++i) { 163 | for (int j = 0; j <= 2*kernel_rad; ++j) { 164 | for (int ch = 0; ch < odim_c; ++ch) { 165 | output += (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xL + i)); 166 | output -= (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xL + i)); 167 | output += (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xR + i)); 168 | output -= (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xR + i)); 169 | } 170 | } 171 | } 172 | } 173 | else { 174 | float gamma = 1 - (yf - floor(yf)); // alpha 175 | for (int i = 0; i <= 2*kernel_rad; ++i) { 176 | for (int j = 0; j <= 2*kernel_rad; ++j) { 177 | for (int ch = 0; ch < odim_c; ++ch) { 178 | output += (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xR + i)); 179 | output -= (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xL + i)); 180 | output += (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xR + i)); 181 | output -= (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xL + i)); 182 | } 183 | } 184 | } 185 | 186 | } 187 | 188 | gradInput[index] = output; 189 | 190 | } 191 | 192 | void resample2d_kernel_forward( 193 | at::Tensor& input1, 194 | at::Tensor& input2, 195 | at::Tensor& output, 196 | int kernel_size) { 197 | 198 | int n = output.numel(); 199 | 200 | const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3)); 201 | const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3)); 202 | 203 | const long4 input2_size = make_long4(input2.size(0), input2.size(1), input2.size(2), input2.size(3)); 204 | const long4 input2_stride = make_long4(input2.stride(0), input2.stride(1), input2.stride(2), input2.stride(3)); 205 | 206 | const long4 output_size = make_long4(output.size(0), output.size(1), output.size(2), output.size(3)); 207 | const long4 output_stride = make_long4(output.stride(0), output.stride(1), output.stride(2), output.stride(3)); 208 | 209 | // TODO: when atomicAdd gets resolved, change to AT_DISPATCH_FLOATING_TYPES_AND_HALF 210 | // AT_DISPATCH_FLOATING_TYPES(input1.type(), "resample_forward_kernel", ([&] { 211 | 212 | kernel_resample2d_update_output<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>( 213 | //at::globalContext().getCurrentCUDAStream() >>>( 214 | n, 215 | input1.data(), 216 | input1_size, 217 | input1_stride, 218 | input2.data(), 219 | input2_size, 220 | input2_stride, 221 | output.data(), 222 | output_size, 223 | output_stride, 224 | kernel_size); 225 | 226 | // })); 227 | 228 | // TODO: ATen-equivalent check 229 | 230 | // THCudaCheck(cudaGetLastError()); 231 | 232 | } 233 | 234 | void resample2d_kernel_backward( 235 | at::Tensor& input1, 236 | at::Tensor& input2, 237 | at::Tensor& gradOutput, 238 | at::Tensor& gradInput1, 239 | at::Tensor& gradInput2, 240 | int kernel_size) { 241 | 242 | int n = gradOutput.numel(); 243 | 244 | const long4 input1_size = make_long4(input1.size(0), input1.size(1), input1.size(2), input1.size(3)); 245 | const long4 input1_stride = make_long4(input1.stride(0), input1.stride(1), input1.stride(2), input1.stride(3)); 246 | 247 | const long4 input2_size = make_long4(input2.size(0), input2.size(1), input2.size(2), input2.size(3)); 248 | const long4 input2_stride = make_long4(input2.stride(0), input2.stride(1), input2.stride(2), input2.stride(3)); 249 | 250 | const long4 gradOutput_size = make_long4(gradOutput.size(0), gradOutput.size(1), gradOutput.size(2), gradOutput.size(3)); 251 | const long4 gradOutput_stride = make_long4(gradOutput.stride(0), gradOutput.stride(1), gradOutput.stride(2), gradOutput.stride(3)); 252 | 253 | const long4 gradInput1_size = make_long4(gradInput1.size(0), gradInput1.size(1), gradInput1.size(2), gradInput1.size(3)); 254 | const long4 gradInput1_stride = make_long4(gradInput1.stride(0), gradInput1.stride(1), gradInput1.stride(2), gradInput1.stride(3)); 255 | 256 | // AT_DISPATCH_FLOATING_TYPES(input1.type(), "resample_backward_input1", ([&] { 257 | 258 | kernel_resample2d_backward_input1<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>( 259 | //at::globalContext().getCurrentCUDAStream() >>>( 260 | n, 261 | input1.data(), 262 | input1_size, 263 | input1_stride, 264 | input2.data(), 265 | input2_size, 266 | input2_stride, 267 | gradOutput.data(), 268 | gradOutput_size, 269 | gradOutput_stride, 270 | gradInput1.data(), 271 | gradInput1_size, 272 | gradInput1_stride, 273 | kernel_size 274 | ); 275 | 276 | // })); 277 | 278 | const long4 gradInput2_size = make_long4(gradInput2.size(0), gradInput2.size(1), gradInput2.size(2), gradInput2.size(3)); 279 | const long4 gradInput2_stride = make_long4(gradInput2.stride(0), gradInput2.stride(1), gradInput2.stride(2), gradInput2.stride(3)); 280 | 281 | n = gradInput2.numel(); 282 | 283 | // AT_DISPATCH_FLOATING_TYPES(gradInput2.type(), "resample_backward_input2", ([&] { 284 | 285 | 286 | kernel_resample2d_backward_input2<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, at::cuda::getCurrentCUDAStream() >>>( 287 | //at::globalContext().getCurrentCUDAStream() >>>( 288 | n, 289 | input1.data(), 290 | input1_size, 291 | input1_stride, 292 | input2.data(), 293 | input2_size, 294 | input2_stride, 295 | gradOutput.data(), 296 | gradOutput_size, 297 | gradOutput_stride, 298 | gradInput2.data(), 299 | gradInput2_size, 300 | gradInput2_stride, 301 | kernel_size 302 | ); 303 | 304 | // })); 305 | 306 | // TODO: Use the ATen equivalent to get last error 307 | 308 | // THCudaCheck(cudaGetLastError()); 309 | 310 | } 311 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/resample2d_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | void resample2d_kernel_forward( 6 | at::Tensor& input1, 7 | at::Tensor& input2, 8 | at::Tensor& output, 9 | int kernel_size); 10 | 11 | void resample2d_kernel_backward( 12 | at::Tensor& input1, 13 | at::Tensor& input2, 14 | at::Tensor& gradOutput, 15 | at::Tensor& gradInput1, 16 | at::Tensor& gradInput2, 17 | int kernel_size); 18 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import os 3 | import torch 4 | 5 | from setuptools import setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | cxx_args = ['-std=c++11'] 9 | 10 | nvcc_args = [ 11 | '-gencode', 'arch=compute_50,code=sm_50', 12 | '-gencode', 'arch=compute_52,code=sm_52', 13 | '-gencode', 'arch=compute_60,code=sm_60', 14 | '-gencode', 'arch=compute_61,code=sm_61', 15 | '-gencode', 'arch=compute_70,code=sm_70', 16 | '-gencode', 'arch=compute_70,code=compute_70' 17 | ] 18 | 19 | setup( 20 | name='resample2d_cuda', 21 | ext_modules=[ 22 | CUDAExtension('resample2d_cuda', [ 23 | 'resample2d_cuda.cc', 24 | 'resample2d_kernel.cu' 25 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 26 | ], 27 | cmdclass={ 28 | 'build_ext': BuildExtension 29 | }) 30 | -------------------------------------------------------------------------------- /models/networks/resample2d_package/src/Resample2d_cuda.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "Resample2d_kernel.h" 5 | 6 | extern THCState* state; 7 | 8 | int Resample2d_cuda_forward(THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* output, int kernel_size) { 9 | Resample2d_kernel_forward(state, input1, input2, output, kernel_size); 10 | return 1; 11 | } 12 | 13 | int Resample2d_cuda_backward(THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* gradOutput, THCudaTensor* gradInput1, THCudaTensor* gradInput2, int kernel_size) { 14 | Resample2d_kernel_backward(state, input1, input2, gradOutput, gradInput1, gradInput2, kernel_size); 15 | 16 | return 1; 17 | } -------------------------------------------------------------------------------- /models/networks/resample2d_package/src/Resample2d_cuda.h: -------------------------------------------------------------------------------- 1 | int Resample2d_cuda_forward(THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* output, int kernel_size); 2 | 3 | int Resample2d_cuda_backward(THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* gradOutput, THCudaTensor* gradInput1, THCudaTensor* gradInput2, int kernel_size); -------------------------------------------------------------------------------- /models/networks/resample2d_package/src/Resample2d_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #define CUDA_NUM_THREADS 512 7 | #define THREADS_PER_BLOCK 64 8 | 9 | #define DIM0(TENSOR) ((TENSOR).x) 10 | #define DIM1(TENSOR) ((TENSOR).y) 11 | #define DIM2(TENSOR) ((TENSOR).z) 12 | #define DIM3(TENSOR) ((TENSOR).w) 13 | 14 | #define DIM3_INDEX(TENSOR, xx, yy, zz, ww) ((TENSOR)[((xx) * (TENSOR##_stride.x)) + ((yy) * (TENSOR##_stride.y)) + ((zz) * (TENSOR##_stride.z)) + ((ww) * (TENSOR##_stride.w))]) 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | __global__ void kernel_Resample2d_updateOutput(const int n, const float* input1, const long4 input1_size, const long4 input1_stride, 21 | const float* input2, const long4 input2_size, const long4 input2_stride, float* output, const long4 output_size, const long4 output_stride, int kernel_size) { 22 | int index = blockIdx.x * blockDim.x + threadIdx.x; 23 | 24 | if (index >= n) { 25 | return; 26 | } 27 | 28 | float val = 0.0; 29 | 30 | int dim_b = DIM0(output_size); 31 | int dim_c = DIM1(output_size); 32 | int dim_h = DIM2(output_size); 33 | int dim_w = DIM3(output_size); 34 | int dim_chw = dim_c * dim_h * dim_w; 35 | int dim_hw = dim_h * dim_w; 36 | 37 | int b = ( index / dim_chw ) % dim_b; 38 | int c = ( index / dim_hw ) % dim_c; 39 | int y = ( index / dim_w ) % dim_h; 40 | int x = ( index ) % dim_w; 41 | 42 | float dx = DIM3_INDEX(input2, b, 0, y, x); 43 | float dy = DIM3_INDEX(input2, b, 1, y, x); 44 | 45 | float xf = float(x) + dx; 46 | float yf = float(y) + dy; 47 | float alpha = xf - floor(xf); // alpha 48 | float beta = yf - floor(yf); // beta 49 | 50 | int xL = max(min( int (floor(xf)), dim_w-1), 0); 51 | int xR = max(min( int (floor(xf)+1), dim_w -1), 0); 52 | int yT = max(min( int (floor(yf)), dim_h-1), 0); 53 | int yB = max(min( int (floor(yf)+1), dim_h-1), 0); 54 | 55 | for (int fy = 0; fy < kernel_size; fy += 1) { 56 | for (int fx = 0; fx < kernel_size; fx += 1) { 57 | val += (1. - alpha)*(1. - beta) * DIM3_INDEX(input1, b, c, yT + fy, xL + fx); 58 | val += (alpha)*(1. - beta) * DIM3_INDEX(input1, b, c, yT + fy, xR + fx); 59 | val += (1. - alpha)*(beta) * DIM3_INDEX(input1, b, c, yB + fy, xL + fx); 60 | val += (alpha)*(beta) * DIM3_INDEX(input1, b, c, yB + fy, xR + fx); 61 | } 62 | } 63 | 64 | output[index] = val; 65 | 66 | } 67 | 68 | 69 | __global__ void kernel_Resample2d_backward_input1( 70 | const int n, const float* input1, const long4 input1_size, const long4 input1_stride, const float* input2, const long4 input2_size, const long4 input2_stride, 71 | const float* gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride, float* gradInput, const long4 gradInput_size, const long4 gradInput_stride, int kernel_size) { 72 | 73 | int index = blockIdx.x * blockDim.x + threadIdx.x; 74 | 75 | if (index >= n) { 76 | return; 77 | } 78 | 79 | int dim_b = DIM0(gradOutput_size); 80 | int dim_c = DIM1(gradOutput_size); 81 | int dim_h = DIM2(gradOutput_size); 82 | int dim_w = DIM3(gradOutput_size); 83 | int dim_chw = dim_c * dim_h * dim_w; 84 | int dim_hw = dim_h * dim_w; 85 | 86 | int b = ( index / dim_chw ) % dim_b; 87 | int c = ( index / dim_hw ) % dim_c; 88 | int y = ( index / dim_w ) % dim_h; 89 | int x = ( index ) % dim_w; 90 | 91 | float dx = DIM3_INDEX(input2, b, 0, y, x); 92 | float dy = DIM3_INDEX(input2, b, 1, y, x); 93 | 94 | float xf = float(x) + dx; 95 | float yf = float(y) + dy; 96 | float alpha = xf - int(xf); // alpha 97 | float beta = yf - int(yf); // beta 98 | 99 | int idim_h = DIM2(input1_size); 100 | int idim_w = DIM3(input1_size); 101 | 102 | int xL = max(min( int (floor(xf)), idim_w-1), 0); 103 | int xR = max(min( int (floor(xf)+1), idim_w -1), 0); 104 | int yT = max(min( int (floor(yf)), idim_h-1), 0); 105 | int yB = max(min( int (floor(yf)+1), idim_h-1), 0); 106 | 107 | for (int fy = 0; fy < kernel_size; fy += 1) { 108 | for (int fx = 0; fx < kernel_size; fx += 1) { 109 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yT + fy), (xL + fx)), (1-alpha)*(1-beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 110 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yT + fy), (xR + fx)), (alpha)*(1-beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 111 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yB + fy), (xL + fx)), (1-alpha)*(beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 112 | atomicAdd(&DIM3_INDEX(gradInput, b, c, (yB + fy), (xR + fx)), (alpha)*(beta) * DIM3_INDEX(gradOutput, b, c, y, x)); 113 | } 114 | } 115 | 116 | } 117 | 118 | __global__ void kernel_Resample2d_backward_input2( 119 | const int n, const float* input1, const long4 input1_size, const long4 input1_stride, const float* input2, const long4 input2_size, const long4 input2_stride, 120 | const float* gradOutput, const long4 gradOutput_size, const long4 gradOutput_stride, float* gradInput, const long4 gradInput_size, const long4 gradInput_stride, int kernel_size) { 121 | 122 | int index = blockIdx.x * blockDim.x + threadIdx.x; 123 | 124 | if (index >= n) { 125 | return; 126 | } 127 | 128 | float output = 0.0; 129 | int kernel_rad = (kernel_size - 1)/2; 130 | 131 | int dim_b = DIM0(gradInput_size); 132 | int dim_c = DIM1(gradInput_size); 133 | int dim_h = DIM2(gradInput_size); 134 | int dim_w = DIM3(gradInput_size); 135 | int dim_chw = dim_c * dim_h * dim_w; 136 | int dim_hw = dim_h * dim_w; 137 | 138 | int b = ( index / dim_chw ) % dim_b; 139 | int c = ( index / dim_hw ) % dim_c; 140 | int y = ( index / dim_w ) % dim_h; 141 | int x = ( index ) % dim_w; 142 | 143 | int odim_c = DIM1(gradOutput_size); 144 | 145 | float dx = DIM3_INDEX(input2, b, 0, y, x); 146 | float dy = DIM3_INDEX(input2, b, 1, y, x); 147 | 148 | float xf = float(x) + dx; 149 | float yf = float(y) + dy; 150 | 151 | int xL = max(min( int (floor(xf)), dim_w-1), 0); 152 | int xR = max(min( int (floor(xf)+1), dim_w -1), 0); 153 | int yT = max(min( int (floor(yf)), dim_h-1), 0); 154 | int yB = max(min( int (floor(yf)+1), dim_h-1), 0); 155 | 156 | if (c % 2) { 157 | float gamma = 1 - (xf - floor(xf)); // alpha 158 | for (int i = 0; i <= 2*kernel_rad; ++i) { 159 | for (int j = 0; j <= 2*kernel_rad; ++j) { 160 | for (int ch = 0; ch < odim_c; ++ch) { 161 | output += (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xL + i)); 162 | output -= (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xL + i)); 163 | output += (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xR + i)); 164 | output -= (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xR + i)); 165 | } 166 | } 167 | } 168 | } 169 | else { 170 | float gamma = 1 - (yf - floor(yf)); // alpha 171 | for (int i = 0; i <= 2*kernel_rad; ++i) { 172 | for (int j = 0; j <= 2*kernel_rad; ++j) { 173 | for (int ch = 0; ch < odim_c; ++ch) { 174 | output += (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xR + i)); 175 | output -= (gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yT + j), (xL + i)); 176 | output += (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xR + i)); 177 | output -= (1-gamma) * DIM3_INDEX(gradOutput, b, ch, y, x) * DIM3_INDEX(input1, b, ch, (yB + j), (xL + i)); 178 | } 179 | } 180 | } 181 | 182 | } 183 | 184 | gradInput[index] = output; 185 | 186 | } 187 | 188 | void Resample2d_kernel_forward(THCState* state, THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* output, int kernel_size) { 189 | int n = 0; 190 | 191 | const long4 input1_size = make_long4(input1->size[0], input1->size[1], input1->size[2], input1->size[3]); 192 | const long4 input1_stride = make_long4(input1->stride[0], input1->stride[1], input1->stride[2], input1->stride[3]); 193 | 194 | const long4 input2_size = make_long4(input2->size[0], input2->size[1], input2->size[2], input2->size[3]); 195 | const long4 input2_stride = make_long4(input2->stride[0], input2->stride[1], input2->stride[2], input2->stride[3]); 196 | 197 | const long4 output_size = make_long4(output->size[0], output->size[1], output->size[2], output->size[3]); 198 | const long4 output_stride = make_long4(output->stride[0], output->stride[1], output->stride[2], output->stride[3]); 199 | 200 | n = THCudaTensor_nElement(state, output); 201 | kernel_Resample2d_updateOutput<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>( 202 | n, THCudaTensor_data(state, input1), input1_size, input1_stride, THCudaTensor_data(state, input2), input2_size, input2_stride, 203 | THCudaTensor_data(state, output), output_size, output_stride, kernel_size); 204 | 205 | THCudaCheck(cudaGetLastError()); 206 | } 207 | 208 | void Resample2d_kernel_backward(THCState* state, THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* gradOutput, THCudaTensor* gradInput1, THCudaTensor* gradInput2, int kernel_size) { 209 | int n = 0; 210 | 211 | const long4 input1_size = make_long4(input1->size[0], input1->size[1], input1->size[2], input1->size[3]); 212 | const long4 input1_stride = make_long4(input1->stride[0], input1->stride[1], input1->stride[2], input1->stride[3]); 213 | 214 | const long4 input2_size = make_long4(input2->size[0], input2->size[1], input2->size[2], input2->size[3]); 215 | const long4 input2_stride = make_long4(input2->stride[0], input2->stride[1], input2->stride[2], input2->stride[3]); 216 | 217 | const long4 gradOutput_size = make_long4(gradOutput->size[0], gradOutput->size[1], gradOutput->size[2], gradOutput->size[3]); 218 | const long4 gradOutput_stride = make_long4(gradOutput->stride[0], gradOutput->stride[1], gradOutput->stride[2], gradOutput->stride[3]); 219 | 220 | const long4 gradInput1_size = make_long4(gradInput1->size[0], gradInput1->size[1], gradInput1->size[2], gradInput1->size[3]); 221 | const long4 gradInput1_stride = make_long4(gradInput1->stride[0], gradInput1->stride[1], gradInput1->stride[2], gradInput1->stride[3]); 222 | 223 | n = THCudaTensor_nElement(state, gradOutput); 224 | kernel_Resample2d_backward_input1<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>( 225 | n, THCudaTensor_data(state, input1), input1_size, input1_stride, THCudaTensor_data(state, input2), input2_size, input2_stride, 226 | THCudaTensor_data(state, gradOutput), gradOutput_size, gradOutput_stride, THCudaTensor_data(state, gradInput1), gradInput1_size, gradInput1_stride, kernel_size 227 | ); 228 | 229 | const long4 gradInput2_size = make_long4(gradInput2->size[0], gradInput2->size[1], gradInput2->size[2], gradInput2->size[3]); 230 | const long4 gradInput2_stride = make_long4(gradInput2->stride[0], gradInput2->stride[1], gradInput2->stride[2], gradInput2->stride[3]); 231 | 232 | n = THCudaTensor_nElement(state, gradInput2); 233 | kernel_Resample2d_backward_input2<<< (n + CUDA_NUM_THREADS - 1)/CUDA_NUM_THREADS, CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>( 234 | n, THCudaTensor_data(state, input1), input1_size, input1_stride, THCudaTensor_data(state, input2), input2_size, input2_stride, 235 | THCudaTensor_data(state, gradOutput), gradOutput_size, gradOutput_stride, THCudaTensor_data(state, gradInput2), gradInput2_size, gradInput2_stride, kernel_size 236 | ); 237 | THCudaCheck(cudaGetLastError()); 238 | } 239 | 240 | #ifdef __cplusplus 241 | } 242 | #endif -------------------------------------------------------------------------------- /models/networks/resample2d_package/src/Resample2d_kernel.h: -------------------------------------------------------------------------------- 1 | #ifdef __cplusplus 2 | extern "C" { 3 | #endif 4 | 5 | void Resample2d_kernel_forward(THCState* state, THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* output, int kernel_size); 6 | 7 | void Resample2d_kernel_backward(THCState* state, THCudaTensor* input1, THCudaTensor* input2, THCudaTensor* gradOutput, THCudaTensor* gradInput1, THCudaTensor* gradInput2, int kernel_size); 8 | 9 | #ifdef __cplusplus 10 | } 11 | #endif -------------------------------------------------------------------------------- /models/networks/resample2d_package/src/Resample2d_kernel.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/models/networks/resample2d_package/src/Resample2d_kernel.o -------------------------------------------------------------------------------- /models/networks/submodules.py: -------------------------------------------------------------------------------- 1 | # freda (todo) : 2 | 3 | import torch.nn as nn 4 | import torch 5 | import numpy as np 6 | 7 | def conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1): 8 | if batchNorm: 9 | return nn.Sequential( 10 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False), 11 | nn.BatchNorm2d(out_planes), 12 | nn.LeakyReLU(0.1,inplace=True) 13 | ) 14 | else: 15 | return nn.Sequential( 16 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True), 17 | nn.LeakyReLU(0.1,inplace=True) 18 | ) 19 | 20 | def i_conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1, bias = True): 21 | if batchNorm: 22 | return nn.Sequential( 23 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=bias), 24 | nn.BatchNorm2d(out_planes), 25 | ) 26 | else: 27 | return nn.Sequential( 28 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=bias), 29 | ) 30 | 31 | def predict_flow(in_planes): 32 | return nn.Conv2d(in_planes,2,kernel_size=3,stride=1,padding=1,bias=True) 33 | 34 | def deconv(in_planes, out_planes): 35 | return nn.Sequential( 36 | nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True), 37 | nn.LeakyReLU(0.1,inplace=True) 38 | ) 39 | 40 | class tofp16(nn.Module): 41 | def __init__(self): 42 | super(tofp16, self).__init__() 43 | 44 | def forward(self, input): 45 | return input.half() 46 | 47 | 48 | class tofp32(nn.Module): 49 | def __init__(self): 50 | super(tofp32, self).__init__() 51 | 52 | def forward(self, input): 53 | return input.float() 54 | 55 | 56 | def init_deconv_bilinear(weight): 57 | f_shape = weight.size() 58 | heigh, width = f_shape[-2], f_shape[-1] 59 | f = np.ceil(width/2.0) 60 | c = (2 * f - 1 - f % 2) / (2.0 * f) 61 | bilinear = np.zeros([heigh, width]) 62 | for x in range(width): 63 | for y in range(heigh): 64 | value = (1 - abs(x / f - c)) * (1 - abs(y / f - c)) 65 | bilinear[x, y] = value 66 | weight.data.fill_(0.) 67 | for i in range(f_shape[0]): 68 | for j in range(f_shape[1]): 69 | weight.data[i,j,:,:] = torch.from_numpy(bilinear) 70 | 71 | 72 | def save_grad(grads, name): 73 | def hook(grad): 74 | grads[name] = grad 75 | return hook 76 | 77 | ''' 78 | def save_grad(grads, name): 79 | def hook(grad): 80 | grads[name] = grad 81 | return hook 82 | import torch 83 | from channelnorm_package.modules.channelnorm import ChannelNorm 84 | model = ChannelNorm().cuda() 85 | grads = {} 86 | a = 100*torch.autograd.Variable(torch.randn((1,3,5,5)).cuda(), requires_grad=True) 87 | a.register_hook(save_grad(grads, 'a')) 88 | b = model(a) 89 | y = torch.mean(b) 90 | y.backward() 91 | 92 | ''' 93 | -------------------------------------------------------------------------------- /test_flow.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | import torchvision.transforms as transforms 4 | import torch.utils.data as data 5 | import numpy as np 6 | from torch.autograd import Variable 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.optim as optim 10 | import argparse 11 | import random 12 | from PIL import Image 13 | import matplotlib.pyplot as plt 14 | import cv2 15 | from models.MonodepthModel import * 16 | from models.PWC_net import * 17 | from models.PWC_net import PWCDCNet 18 | from utils.scene_dataloader import * 19 | from utils.utils import * 20 | 21 | def get_args(): 22 | parser = argparse.ArgumentParser() 23 | 24 | parser.add_argument('--model_name', type=str, help='model name', default='monodepth') 25 | parser.add_argument('--data_path', type=str, help='path to the data', required=True) 26 | parser.add_argument('--filenames_file', type=str, help='path to the filenames text file', required=True) 27 | parser.add_argument('--input_height', type=int, help='input height', default=256) 28 | parser.add_argument('--input_width', type=int, help='input width', default=512) 29 | parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', required=True) 30 | args = parser.parse_args() 31 | return args 32 | 33 | args = get_args() 34 | 35 | checkpoint = torch.load(args.checkpoint_path) 36 | if args.model_name == 'monodepth': 37 | net = MonodepthNet().cuda() 38 | elif args.model_name == 'pwc': 39 | net = pwc_dc_net().cuda() 40 | args.input_width = 832 41 | net.load_state_dict(checkpoint['state_dict']) 42 | 43 | former_test, latter_test, flow = get_flow_data(args.filenames_file, args.data_path) 44 | TestFlowLoader = torch.utils.data.DataLoader( 45 | myImageFolder(former_test, latter_test, flow, args), 46 | batch_size = 1, shuffle = False, num_workers = 1, drop_last = False) 47 | 48 | total_error = 0 49 | fl_error = 0 50 | num_test = 0 51 | for batch_idx, (left, right, gt, mask, h, w) in enumerate(TestFlowLoader, 0): 52 | 53 | left_batch = torch.cat((left, torch.from_numpy(np.flip(left.numpy(), 3).copy())), 0) 54 | right_batch = torch.cat((right, torch.from_numpy(np.flip(right.numpy(), 3).copy())), 0) 55 | 56 | left = Variable(left_batch.cuda()) 57 | right = Variable(right_batch.cuda()) 58 | 59 | model_input = torch.cat((left, right), 1) 60 | if args.model_name == 'monodepth': 61 | disp_est_scale, disp_est = net(model_input) 62 | elif args.model_name == 'pwc': 63 | disp_est_scale = net(model_input) 64 | 65 | mask = np.ceil(np.clip(np.abs(gt[0,0]), 0, 1)) 66 | 67 | disp_ori_scale = nn.UpsamplingBilinear2d(size=(int(h), int(w)))(disp_est_scale[0][:1]) 68 | disp_ori_scale[0,0] = disp_ori_scale[0,0] * int(w) / args.input_width 69 | disp_ori_scale[0,1] = disp_ori_scale[0,1] * int(h) / args.input_height 70 | 71 | error, fl = evaluate_flow(disp_ori_scale[0].data.cpu().numpy(), gt[0].numpy(), mask.numpy()) 72 | total_error += error 73 | fl_error += fl 74 | num_test += 1 75 | 76 | total_error /= num_test 77 | fl_error /= num_test 78 | print("The average EPE is : ", total_error) 79 | print("The average Fl is : ", fl_error) -------------------------------------------------------------------------------- /test_stereo.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | import torchvision.transforms as transforms 4 | import torch.utils.data as data 5 | import numpy as np 6 | from torch.autograd import Variable 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.optim as optim 10 | import argparse 11 | import random 12 | from PIL import Image 13 | import matplotlib.pyplot as plt 14 | import cv2 15 | from models.MonodepthModel import * 16 | from models.PWC_net import * 17 | from models.PWC_net import PWCDCNet 18 | from utils.scene_dataloader import * 19 | from utils.utils import * 20 | 21 | def get_args(): 22 | parser = argparse.ArgumentParser() 23 | 24 | parser.add_argument('--model_name', type=str, help='model name', default='monodepth') 25 | parser.add_argument('--split', type=str, help='validation set', default='kitti') 26 | parser.add_argument('--data_path', type=str, help='path to the data', required=True) 27 | parser.add_argument('--filenames_file', type=str, help='path to the filenames text file', required=True) 28 | parser.add_argument('--input_height', type=int, help='input height', default=256) 29 | parser.add_argument('--input_width', type=int, help='input width', default=512) 30 | parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', required=True) 31 | args = parser.parse_args() 32 | return args 33 | 34 | args = get_args() 35 | 36 | checkpoint = torch.load(args.checkpoint_path) 37 | if args.model_name == 'monodepth': 38 | net = MonodepthNet().cuda() 39 | elif args.model_name == 'pwc': 40 | net = pwc_dc_net().cuda() 41 | args.input_width = 832 42 | net.load_state_dict(checkpoint['state_dict']) 43 | 44 | left_image_test, right_image_test = get_data(args.filenames_file, args.data_path) 45 | TestImageLoader = torch.utils.data.DataLoader( 46 | myImageFolder(left_image_test, right_image_test, None, args), 47 | batch_size = 1, shuffle = False, num_workers = 1, drop_last =False) 48 | 49 | if args.split == 'kitti': 50 | disparities = np.zeros((200, args.input_height, args.input_width), dtype=np.float32) 51 | elif args.split == 'eigen': 52 | disparities = np.zeros((697, args.input_height, args.input_width), dtype=np.float32) 53 | 54 | for batch_idx, (left, right) in enumerate(TestImageLoader, 0): 55 | print(batch_idx) 56 | left_batch = torch.cat((left, torch.from_numpy(np.flip(left.numpy(), 3).copy())), 0) 57 | right_batch = torch.cat((right, torch.from_numpy(np.flip(right.numpy(), 3).copy())), 0) 58 | 59 | left_pyramid = make_pyramid(left_batch, 4) 60 | 61 | model_input = Variable(torch.cat((left_batch, right_batch), 1).cuda()) 62 | if args.model_name == 'monodepth': 63 | disp_est_scale, disp_est= net(model_input) 64 | elif args.model_name == 'pwc': 65 | disp_est_scale = net(model_input) 66 | disp_est = [torch.cat((disp_est_scale[i][:,0,:,:].unsqueeze(1) / disp_est_scale[i].shape[3], disp_est_scale[i][:,1,:,:].unsqueeze(1) / disp_est_scale[i].shape[2]), 1) for i in range(4)] 67 | 68 | disparities[batch_idx] = -disp_est[0][0,0,:,:].data.cpu().numpy() 69 | print('done') 70 | np.save('./disparities.npy', disparities) 71 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | import torchvision.transforms as transforms 4 | import torch.utils.data as data 5 | import numpy as np 6 | from torch.autograd import Variable 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.optim as optim 10 | import argparse 11 | import random 12 | from PIL import Image 13 | import matplotlib.pyplot as plt 14 | import cv2 15 | from models.MonodepthModel import * 16 | from models.PWC_net import * 17 | from models.PWC_net import PWCDCNet 18 | from utils.scene_dataloader import * 19 | from utils.utils import * 20 | from models.networks.submodules import * 21 | from models.networks.resample2d_package.resample2d import Resample2d 22 | 23 | def get_args(): 24 | parser = argparse.ArgumentParser() 25 | 26 | parser.add_argument('--model_name', type=str, help='model name', default='monodepth') 27 | parser.add_argument('--data_path', type=str, help='path to the data', required=True) 28 | parser.add_argument('--filenames_file', type=str, help='path to the filenames text file', required=True) 29 | parser.add_argument('--input_height', type=int, help='input height', default=256) 30 | parser.add_argument('--input_width', type=int, help='input width', default=512) 31 | parser.add_argument('--batch_size', type=int, help='batch size', default=2) 32 | parser.add_argument('--num_epochs', type=int, help='number of epochs', default=80) 33 | parser.add_argument('--learning_rate', type=float, help='initial learning rate', default=1e-4) 34 | parser.add_argument('--lr_loss_weight', type=float, help='left-right consistency weight', default=0.5) 35 | parser.add_argument('--alpha_image_loss', type=float, help='weight between SSIM and L1 in the image loss', default=0.85) 36 | parser.add_argument('--disp_gradient_loss_weight', type=float, help='disparity smoothness weigth', default=0.1) 37 | parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=8) 38 | parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', default='') 39 | parser.add_argument('--type_of_2warp', type=int, help='2warp type', default=0) 40 | args = parser.parse_args() 41 | return args 42 | 43 | args = get_args() 44 | 45 | if args.model_name == 'monodepth': 46 | net = MonodepthNet().cuda() 47 | elif args.model_name == 'pwc': 48 | net = pwc_dc_net().cuda() 49 | args.input_width = 832 50 | 51 | left_image_1, left_image_2, right_image_1, right_image_2 = get_kitti_cycle_data(args.filenames_file, args.data_path) 52 | CycleLoader = torch.utils.data.DataLoader( 53 | myCycleImageFolder(left_image_1, left_image_2, right_image_1, right_image_2, True, args), 54 | batch_size = args.batch_size, shuffle = True, num_workers = args.num_threads, drop_last = False) 55 | optimizer = optim.Adam(net.parameters(), lr = args.learning_rate) 56 | scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[4, 7, 10, 13], gamma=0.5) 57 | 58 | for epoch in range(args.num_epochs): 59 | 60 | scheduler.step() 61 | 62 | for batch_idx, (left_image_1, left_image_2, right_image_1, right_image_2) in enumerate(CycleLoader, 0): 63 | 64 | optimizer.zero_grad() 65 | 66 | former = torch.cat((left_image_2, left_image_1, right_image_1, left_image_1), 0) 67 | latter = torch.cat((right_image_2, left_image_2, right_image_2, right_image_1), 0) 68 | 69 | left_pyramid = make_pyramid(former, 4) 70 | right_pyramid = make_pyramid(latter, 4) 71 | 72 | model_input = Variable(torch.cat((former, latter), 1).cuda()) 73 | model_input_2 = Variable(torch.cat((latter, former), 1).cuda()) 74 | 75 | if args.model_name == 'monodepth': 76 | disp_est_scale, disp_est = net(model_input) 77 | disp_est_scale_2, disp_est_2 = net(model_input_2) 78 | 79 | elif args.model_name == 'pwc': 80 | disp_est_scale = net(model_input) 81 | disp_est = [torch.cat((disp_est_scale[i][:,0,:,:].unsqueeze(1) / disp_est_scale[i].shape[3], 82 | disp_est_scale[i][:,1,:,:].unsqueeze(1) / disp_est_scale[i].shape[2]), 1) for i in range(4)] 83 | disp_est_scale_2 = net(model_input_2) 84 | disp_est_2 = [torch.cat((disp_est_scale_2[i][:,0,:,:].unsqueeze(1) / disp_est_scale_2[i].shape[3], 85 | disp_est_scale_2[i][:,1,:,:].unsqueeze(1) / disp_est_scale_2[i].shape[2]), 1) for i in range(4)] 86 | 87 | border_mask = [create_border_mask(left_pyramid[i], 0.1) for i in range(4)] 88 | 89 | fw_mask = [] 90 | bw_mask = [] 91 | for i in range(4): 92 | fw, bw, diff_fw, diff_bw = get_mask(disp_est_scale[i], disp_est_scale_2[i], border_mask[i]) 93 | fw += 1e-3 94 | bw += 1e-3 95 | fw[[0,1,6,7]] = fw[[0,1,6,7]] * 0 + 1 96 | bw[[0,1,6,7]] = bw[[0,1,6,7]] * 0 + 1 97 | fw_detached = fw.clone().detach() 98 | bw_detached = bw.clone().detach() 99 | fw_mask.append(fw_detached) 100 | bw_mask.append(bw_detached) 101 | 102 | #reconstruction from right to left 103 | left_est = [Resample2d()(right_pyramid[i], disp_est_scale[i]) for i in range(4)] 104 | l1_left = [torch.abs(left_est[i] - left_pyramid[i]) * fw_mask[i] for i in range(4)] 105 | l1_reconstruction_loss_left = [torch.mean(l1_left[i]) / torch.mean(fw_mask[i]) for i in range(4)] 106 | ssim_left = [SSIM(left_est[i] * fw_mask[i], left_pyramid[i] * fw_mask[i]) for i in range(4)] 107 | ssim_loss_left = [torch.mean(ssim_left[i]) / torch.mean(fw_mask[i]) for i in range(4)] 108 | image_loss_left = [args.alpha_image_loss * ssim_loss_left[i] + 109 | (1 - args.alpha_image_loss) * l1_reconstruction_loss_left[i] for i in range(4)] 110 | image_loss = image_loss_left[0] + image_loss_left[1] + image_loss_left[2] + image_loss_left[3] 111 | 112 | disp_loss = [cal_grad2_error(disp_est_scale[i] / 20, left_pyramid[i], 1.0) for i in range(4)] 113 | disp_gradient_loss = disp_loss[0] + disp_loss[1] + disp_loss[2] + disp_loss[3] 114 | 115 | #reconstruction from left to right 116 | right_est = [Resample2d()(left_pyramid[i], disp_est_scale_2[i]) for i in range(4)] 117 | l1_right = [torch.abs(right_est[i] - right_pyramid[i]) * bw_mask[i] for i in range(4)] 118 | l1_reconstruction_loss_right = [torch.mean(l1_right[i]) / torch.mean(bw_mask[i]) for i in range(4)] 119 | ssim_right = [SSIM(right_est[i] * bw_mask[i], right_pyramid[i] * bw_mask[i]) for i in range(4)] 120 | ssim_loss_right = [torch.mean(ssim_right[i]) / torch.mean(bw_mask[i]) for i in range(4)] 121 | image_loss_right = [args.alpha_image_loss * ssim_loss_right[i] + 122 | (1 - args.alpha_image_loss) * l1_reconstruction_loss_right[i] for i in range(4)] 123 | image_loss_2 = image_loss_right[0] + image_loss_right[1] + image_loss_right[2] + image_loss_right[3] 124 | 125 | disp_loss_2 = [cal_grad2_error(disp_est_scale_2[i] / 20, right_pyramid[i], 1.0) for i in range(4)] 126 | disp_gradient_loss_2 = disp_loss_2[0] + disp_loss_2[1] + disp_loss_2[2] + disp_loss_2[3] 127 | 128 | #LR consistency 129 | right_to_left_disp = [- Resample2d()(disp_est_2[i], disp_est_scale[i]) for i in range(4)] 130 | left_to_right_disp = [- Resample2d()(disp_est[i], disp_est_scale_2[i]) for i in range(4)] 131 | 132 | lr_left_loss = [torch.mean(torch.abs(right_to_left_disp[i][[0,1,6,7]] - disp_est[i][[0,1,6,7]])) for i in range(4)] 133 | lr_right_loss = [torch.mean(torch.abs(left_to_right_disp[i][[0,1,6,7]] - disp_est_2[i][[0,1,6,7]])) for i in range(4)] 134 | lr_loss = sum(lr_left_loss + lr_right_loss) 135 | 136 | loss = image_loss + image_loss_2 + 10 * (disp_gradient_loss + disp_gradient_loss_2) + args.lr_loss_weight * lr_loss 137 | 138 | """ 139 | ########################################################################################## 140 | # # 141 | # batch 7,8 mask for the direction of the reconstruction # 142 | # forward L_t ------------> R_t # 143 | # | | mask : L_t+1 ---> L_t ---> R_t # 144 | # 3,4 | | 5,6 mask_2 : L_t+1 ---> R_t+1 ---> R_t # 145 | # | | mask_3 : R_t+1 ---> R_t ---> L_t # 146 | # v v mask_4 : R_t+1 ---> L_t+1 ---> L_t # 147 | # L_t+1 ----------> R_t+1 mask_5 : R_t ---> L_t ---> L_t+1 # 148 | # 1,2 # 149 | # # 150 | ########################################################################################## 151 | """ 152 | 153 | if args.type_of_2warp == 1: 154 | mask_4 = [fw_mask[i][[2,3]] for i in range(4)] 155 | warp2_est_4 = [Resample2d()(left_est[i][[0,1]], disp_est_scale[i][[2,3]]) for i in range(4)] 156 | loss += 0.1 * sum([warp_2(warp2_est_4[i], left_pyramid[i][[6,7]], mask_4[i], args) for i in range(4)]) 157 | mask_5 = [bw_mask[i][[2,3]] for i in range(4)] 158 | warp2_est_5 = [Resample2d()(left_est[i][[6,7]], disp_est_scale_2[i][[2,3]]) for i in range(4)] 159 | loss += 0.1 * sum([warp_2(warp2_est_5[i], left_pyramid[i][[0,1]], mask_5[i], args) for i in range(4)]) 160 | 161 | elif args.type_of_2warp == 2: 162 | mask = [Resample2d()(fw_mask[i][[2,3]], disp_est_scale_2[i][[0,1]]) for i in range(4)] 163 | warp2_est = [Resample2d()(left_est[i][[2,3]], disp_est_scale_2[i][[6,7]]) for i in range(4)] 164 | loss += 0.1 * sum([warp_2(warp2_est[i], right_pyramid[i][[6,7]], mask[i], args) for i in range(4)]) 165 | mask_3 = [Resample2d()(fw_mask[i][[4,5]], disp_est_scale[i][[0,1]]) for i in range(4)] 166 | warp2_est_3 = [Resample2d()(left_est[i][[4,5]], disp_est_scale[i][[6,7]]) for i in range(4)] 167 | loss += 0.1 * sum([warp_2(warp2_est_3[i], left_pyramid[i][[6,7]], mask_3[i], args) for i in range(4)]) 168 | 169 | elif args.type_of_2warp == 3: 170 | mask = [Resample2d()(fw_mask[i][[2,3]], disp_est_scale_2[i][[0,1]]) for i in range(4)] 171 | warp2_est = [Resample2d()(left_est[i][[2,3]], disp_est_scale_2[i][[6,7]]) for i in range(4)] 172 | loss += 0.1 * sum([warp_2(warp2_est[i], right_pyramid[i][[6,7]], mask[i], args) for i in range(4)]) 173 | mask_2 = [fw_mask[i][[4,5]] for i in range(4)] 174 | warp2_est_2 = [Resample2d()(right_est[i][[0,1]], disp_est_scale[i][[4,5]]) for i in range(4)] 175 | loss += 0.1 * sum([warp_2(warp2_est_2[i], right_pyramid[i][[6,7]], mask_2[i], args) for i in range(4)]) 176 | 177 | loss.backward() 178 | 179 | if args.model_name == 'monodepth': 180 | print("Epoch :", epoch) 181 | print("Batch Index :", batch_idx) 182 | print(net.conv1.weight.grad[0,0,0,0]) 183 | elif args.model_name == 'pwc': 184 | print("Epoch :", epoch) 185 | print("Batch Index :", batch_idx) 186 | print(net.conv1a[0].weight.grad[0,0,0,0]) 187 | 188 | optimizer.step() 189 | 190 | if epoch % 1 == 0: 191 | state = {'epoch': epoch, 'state_dict': net.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler} 192 | torch.save(state, args.checkpoint_path + "model_epoch" + str(epoch)) 193 | print("The model of epoch ", epoch, "has been saved.") 194 | 195 | -------------------------------------------------------------------------------- /utils/__pycache__/evaluate_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/utils/__pycache__/evaluate_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/evaluate_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/utils/__pycache__/evaluate_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/scene_dataloader.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/utils/__pycache__/scene_dataloader.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/scene_dataloader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/utils/__pycache__/scene_dataloader.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/utils/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lelimite4444/BridgeDepthFlow/14183b99830e1f41e774c0d43fdb058d07f2a397/utils/__pycache__/utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/evaluate_kitti.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import argparse 4 | from evaluation_utils import * 5 | 6 | parser = argparse.ArgumentParser(description='Evaluation on the KITTI dataset') 7 | parser.add_argument('--split', type=str, help='data split, kitti or eigen', required=True) 8 | parser.add_argument('--predicted_disp_path', type=str, help='path to estimated disparities', required=True) 9 | parser.add_argument('--gt_path', type=str, help='path to ground truth disparities', required=True) 10 | parser.add_argument('--min_depth', type=float, help='minimum depth for evaluation', default=1e-3) 11 | parser.add_argument('--max_depth', type=float, help='maximum depth for evaluation', default=80) 12 | parser.add_argument('--eigen_crop', help='if set, crops according to Eigen NIPS14', action='store_true') 13 | parser.add_argument('--garg_crop', help='if set, crops according to Garg ECCV16', action='store_true') 14 | 15 | args = parser.parse_args() 16 | 17 | if __name__ == '__main__': 18 | 19 | pred_disparities = np.load(args.predicted_disp_path) 20 | 21 | if args.split == 'kitti': 22 | num_samples = 200 23 | 24 | gt_disparities = load_gt_disp_kitti(args.gt_path) 25 | gt_depths, pred_depths, pred_disparities_resized = convert_disps_to_depths_kitti(gt_disparities, pred_disparities) 26 | 27 | elif args.split == 'eigen': 28 | num_samples = 697 29 | test_files = read_text_lines('/home/hylai/monodepth/utils/filenames/eigen_test_files.txt') 30 | gt_files, gt_calib, im_sizes, im_files, cams = read_file_data(test_files, args.gt_path) 31 | 32 | num_test = len(im_files) 33 | gt_depths = [] 34 | pred_depths = [] 35 | for t_id in range(num_samples): 36 | camera_id = cams[t_id] # 2 is left, 3 is right 37 | depth = generate_depth_map(gt_calib[t_id], gt_files[t_id], im_sizes[t_id], camera_id, False, True) 38 | gt_depths.append(depth.astype(np.float32)) 39 | 40 | disp_pred = cv2.resize(pred_disparities[t_id], (im_sizes[t_id][1], im_sizes[t_id][0]), interpolation=cv2.INTER_LINEAR) 41 | disp_pred = disp_pred * disp_pred.shape[1] 42 | 43 | # need to convert from disparity to depth 44 | focal_length, baseline = get_focal_length_baseline(gt_calib[t_id], camera_id) 45 | depth_pred = (baseline * focal_length) / disp_pred 46 | depth_pred[np.isinf(depth_pred)] = 0 47 | 48 | pred_depths.append(depth_pred) 49 | 50 | rms = np.zeros(num_samples, np.float32) 51 | log_rms = np.zeros(num_samples, np.float32) 52 | abs_rel = np.zeros(num_samples, np.float32) 53 | sq_rel = np.zeros(num_samples, np.float32) 54 | d1_all = np.zeros(num_samples, np.float32) 55 | a1 = np.zeros(num_samples, np.float32) 56 | a2 = np.zeros(num_samples, np.float32) 57 | a3 = np.zeros(num_samples, np.float32) 58 | 59 | for i in range(num_samples): 60 | 61 | gt_depth = gt_depths[i] 62 | pred_depth = pred_depths[i] 63 | 64 | pred_depth[pred_depth < args.min_depth] = args.min_depth 65 | pred_depth[pred_depth > args.max_depth] = args.max_depth 66 | 67 | if args.split == 'eigen': 68 | mask = np.logical_and(gt_depth > args.min_depth, gt_depth < args.max_depth) 69 | 70 | 71 | if args.garg_crop or args.eigen_crop: 72 | gt_height, gt_width = gt_depth.shape 73 | 74 | # crop used by Garg ECCV16 75 | # if used on gt_size 370x1224 produces a crop of [-218, -3, 44, 1180] 76 | if args.garg_crop: 77 | crop = np.array([0.40810811 * gt_height, 0.99189189 * gt_height, 78 | 0.03594771 * gt_width, 0.96405229 * gt_width]).astype(np.int32) 79 | # crop we found by trial and error to reproduce Eigen NIPS14 results 80 | elif args.eigen_crop: 81 | crop = np.array([0.3324324 * gt_height, 0.91351351 * gt_height, 82 | 0.0359477 * gt_width, 0.96405229 * gt_width]).astype(np.int32) 83 | 84 | crop_mask = np.zeros(mask.shape) 85 | crop_mask[crop[0]:crop[1],crop[2]:crop[3]] = 1 86 | mask = np.logical_and(mask, crop_mask) 87 | 88 | if args.split == 'kitti': 89 | gt_disp = gt_disparities[i] 90 | mask = gt_disp > 0 91 | pred_disp = pred_disparities_resized[i] 92 | 93 | disp_diff = np.abs(gt_disp[mask] - pred_disp[mask]) 94 | bad_pixels = np.logical_and(disp_diff >= 3, (disp_diff / gt_disp[mask]) >= 0.05) 95 | d1_all[i] = 100.0 * bad_pixels.sum() / mask.sum() 96 | 97 | abs_rel[i], sq_rel[i], rms[i], log_rms[i], a1[i], a2[i], a3[i] = compute_errors(gt_depth[mask], pred_depth[mask]) 98 | 99 | print("{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}".format('abs_rel', 'sq_rel', 'rms', 'log_rms', 'd1_all', 'a1', 'a2', 'a3')) 100 | print("{:10.4f}, {:10.4f}, {:10.3f}, {:10.3f}, {:10.3f}, {:10.3f}, {:10.3f}, {:10.3f}".format(abs_rel.mean(), sq_rel.mean(), rms.mean(), log_rms.mean(), d1_all.mean(), a1.mean(), a2.mean(), a3.mean())) 101 | -------------------------------------------------------------------------------- /utils/evaluation_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import os 4 | import cv2 5 | from collections import Counter 6 | import pickle 7 | 8 | def compute_errors(gt, pred): 9 | thresh = np.maximum((gt / pred), (pred / gt)) 10 | a1 = (thresh < 1.25 ).mean() 11 | a2 = (thresh < 1.25 ** 2).mean() 12 | a3 = (thresh < 1.25 ** 3).mean() 13 | 14 | rmse = (gt - pred) ** 2 15 | rmse = np.sqrt(rmse.mean()) 16 | 17 | rmse_log = (np.log(gt) - np.log(pred)) ** 2 18 | rmse_log = np.sqrt(rmse_log.mean()) 19 | 20 | abs_rel = np.mean(np.abs(gt - pred) / gt) 21 | 22 | sq_rel = np.mean(((gt - pred)**2) / gt) 23 | 24 | return abs_rel, sq_rel, rmse, rmse_log, a1, a2, a3 25 | 26 | ############################################################################### 27 | ####################### KITTI 28 | 29 | width_to_focal = dict() 30 | width_to_focal[1242] = 721.5377 31 | width_to_focal[1241] = 718.856 32 | width_to_focal[1224] = 707.0493 33 | width_to_focal[1238] = 718.3351 34 | 35 | def load_gt_disp_kitti(path): 36 | gt_disparities = [] 37 | for i in range(200): 38 | disp = cv2.imread(path + "/training/disp_occ_0/" + str(i).zfill(6) + "_10.png", -1) 39 | disp = disp.astype(np.float32) / 256 40 | gt_disparities.append(disp) 41 | return gt_disparities 42 | 43 | def convert_disps_to_depths_kitti(gt_disparities, pred_disparities): 44 | gt_depths = [] 45 | pred_depths = [] 46 | pred_disparities_resized = [] 47 | 48 | for i in range(len(gt_disparities)): 49 | gt_disp = gt_disparities[i] 50 | height, width = gt_disp.shape 51 | 52 | pred_disp = pred_disparities[i] 53 | pred_disp = width * cv2.resize(pred_disp, (width, height), interpolation=cv2.INTER_LINEAR) 54 | 55 | pred_disparities_resized.append(pred_disp) 56 | 57 | mask = gt_disp > 0 58 | 59 | gt_depth = width_to_focal[width] * 0.54 / (gt_disp + (1.0 - mask)) 60 | pred_depth = width_to_focal[width] * 0.54 / pred_disp 61 | 62 | gt_depths.append(gt_depth) 63 | pred_depths.append(pred_depth) 64 | return gt_depths, pred_depths, pred_disparities_resized 65 | 66 | 67 | ############################################################################### 68 | ####################### EIGEN 69 | 70 | def read_text_lines(file_path): 71 | f = open(file_path, 'r') 72 | lines = f.readlines() 73 | f.close() 74 | lines = [l.rstrip() for l in lines] 75 | return lines 76 | 77 | def read_file_data(files, data_root): 78 | gt_files = [] 79 | gt_calib = [] 80 | im_sizes = [] 81 | im_files = [] 82 | cams = [] 83 | num_probs = 0 84 | for filename in files: 85 | filename = filename.split()[0] 86 | splits = filename.split('/') 87 | camera_id = np.int32(splits[2][-1:]) # 2 is left, 3 is right 88 | date = splits[0] 89 | im_id = splits[4][:10] 90 | file_root = '{}/{}' 91 | 92 | im = filename 93 | vel = '{}/{}/velodyne_points/data/{}.bin'.format(splits[0], splits[1], im_id) 94 | 95 | if os.path.isfile(data_root + im): 96 | gt_files.append(data_root + vel) 97 | gt_calib.append(data_root + date + '/') 98 | im_sizes.append(cv2.imread(data_root + im).shape[:2]) 99 | im_files.append(data_root + im) 100 | cams.append(2) 101 | else: 102 | num_probs += 1 103 | print('{} missing'.format(data_root + im)) 104 | print (num_probs, 'files missing') 105 | 106 | return gt_files, gt_calib, im_sizes, im_files, cams 107 | 108 | def load_velodyne_points(file_name): 109 | # adapted from https://github.com/hunse/kitti 110 | points = np.fromfile(file_name, dtype=np.float32).reshape(-1, 4) 111 | points[:, 3] = 1.0 # homogeneous 112 | return points 113 | 114 | 115 | def lin_interp(shape, xyd): 116 | # taken from https://github.com/hunse/kitti 117 | m, n = shape 118 | ij, d = xyd[:, 1::-1], xyd[:, 2] 119 | f = LinearNDInterpolator(ij, d, fill_value=0) 120 | J, I = np.meshgrid(np.arange(n), np.arange(m)) 121 | IJ = np.vstack([I.flatten(), J.flatten()]).T 122 | disparity = f(IJ).reshape(shape) 123 | return disparity 124 | 125 | 126 | def read_calib_file(path): 127 | # taken from https://github.com/hunse/kitti 128 | float_chars = set("0123456789.e+- ") 129 | data = {} 130 | with open(path, 'r') as f: 131 | for line in f.readlines(): 132 | key, value = line.split(':', 1) 133 | value = value.strip() 134 | data[key] = value 135 | if float_chars.issuperset(value): 136 | # try to cast to float array 137 | try: 138 | data[key] = np.array(list(map(float, value.split(' ')))) 139 | except ValueError: 140 | # casting error: data[key] already eq. value, so pass 141 | pass 142 | 143 | return data 144 | 145 | 146 | def get_focal_length_baseline(calib_dir, cam): 147 | cam2cam = read_calib_file(calib_dir + 'calib_cam_to_cam.txt') 148 | P2_rect = cam2cam['P_rect_02'].reshape(3,4) 149 | P3_rect = cam2cam['P_rect_03'].reshape(3,4) 150 | 151 | # cam 2 is left of camera 0 -6cm 152 | # cam 3 is to the right +54cm 153 | b2 = P2_rect[0,3] / -P2_rect[0,0] 154 | b3 = P3_rect[0,3] / -P3_rect[0,0] 155 | baseline = b3-b2 156 | 157 | if cam==2: 158 | focal_length = P2_rect[0,0] 159 | elif cam==3: 160 | focal_length = P3_rect[0,0] 161 | 162 | return focal_length, baseline 163 | 164 | 165 | def sub2ind(matrixSize, rowSub, colSub): 166 | m, n = matrixSize 167 | return rowSub * (n-1) + colSub - 1 168 | 169 | def generate_depth_map(calib_dir, velo_file_name, im_shape, cam=2, interp=False, vel_depth=False): 170 | # load calibration files 171 | cam2cam = read_calib_file(calib_dir + 'calib_cam_to_cam.txt') 172 | velo2cam = read_calib_file(calib_dir + 'calib_velo_to_cam.txt') 173 | velo2cam = np.hstack((velo2cam['R'].reshape(3,3), velo2cam['T'][..., np.newaxis])) 174 | velo2cam = np.vstack((velo2cam, np.array([0, 0, 0, 1.0]))) 175 | 176 | # compute projection matrix velodyne->image plane 177 | R_cam2rect = np.eye(4) 178 | R_cam2rect[:3,:3] = cam2cam['R_rect_00'].reshape(3,3) 179 | P_rect = cam2cam['P_rect_0'+str(cam)].reshape(3,4) 180 | P_velo2im = np.dot(np.dot(P_rect, R_cam2rect), velo2cam) 181 | 182 | # load velodyne points and remove all behind image plane (approximation) 183 | # each row of the velodyne data is forward, left, up, reflectance 184 | velo = load_velodyne_points(velo_file_name) 185 | velo = velo[velo[:, 0] >= 0, :] 186 | 187 | # project the points to the camera 188 | velo_pts_im = np.dot(P_velo2im, velo.T).T 189 | velo_pts_im[:, :2] = velo_pts_im[:,:2] / velo_pts_im[:,2][..., np.newaxis] 190 | 191 | if vel_depth: 192 | velo_pts_im[:, 2] = velo[:, 0] 193 | 194 | # check if in bounds 195 | # use minus 1 to get the exact same value as KITTI matlab code 196 | velo_pts_im[:, 0] = np.round(velo_pts_im[:,0]) - 1 197 | velo_pts_im[:, 1] = np.round(velo_pts_im[:,1]) - 1 198 | val_inds = (velo_pts_im[:, 0] >= 0) & (velo_pts_im[:, 1] >= 0) 199 | val_inds = val_inds & (velo_pts_im[:,0] < im_shape[1]) & (velo_pts_im[:,1] < im_shape[0]) 200 | velo_pts_im = velo_pts_im[val_inds, :] 201 | 202 | # project to image 203 | depth = np.zeros((im_shape)) 204 | depth[velo_pts_im[:, 1].astype(np.int), velo_pts_im[:, 0].astype(np.int)] = velo_pts_im[:, 2] 205 | 206 | # find the duplicate points and choose the closest depth 207 | inds = sub2ind(depth.shape, velo_pts_im[:, 1], velo_pts_im[:, 0]) 208 | dupe_inds = [item for item, count in Counter(inds).items() if count > 1] 209 | for dd in dupe_inds: 210 | pts = np.where(inds==dd)[0] 211 | x_loc = int(velo_pts_im[pts[0], 0]) 212 | y_loc = int(velo_pts_im[pts[0], 1]) 213 | depth[y_loc, x_loc] = velo_pts_im[pts, 2].min() 214 | depth[depth<0] = 0 215 | 216 | if interp: 217 | # interpolate the depth map to fill in holes 218 | depth_interp = lin_interp(im_shape, velo_pts_im) 219 | return depth, depth_interp 220 | else: 221 | return depth 222 | 223 | 224 | 225 | -------------------------------------------------------------------------------- /utils/filenames/kitti_stereo_2015_test_files.txt: -------------------------------------------------------------------------------- 1 | training/image_2/000000_10.png training/image_3/000000_10.png 2 | training/image_2/000001_10.png training/image_3/000001_10.png 3 | training/image_2/000002_10.png training/image_3/000002_10.png 4 | training/image_2/000003_10.png training/image_3/000003_10.png 5 | training/image_2/000004_10.png training/image_3/000004_10.png 6 | training/image_2/000005_10.png training/image_3/000005_10.png 7 | training/image_2/000006_10.png training/image_3/000006_10.png 8 | training/image_2/000007_10.png training/image_3/000007_10.png 9 | training/image_2/000008_10.png training/image_3/000008_10.png 10 | training/image_2/000009_10.png training/image_3/000009_10.png 11 | training/image_2/000010_10.png training/image_3/000010_10.png 12 | training/image_2/000011_10.png training/image_3/000011_10.png 13 | training/image_2/000012_10.png training/image_3/000012_10.png 14 | training/image_2/000013_10.png training/image_3/000013_10.png 15 | training/image_2/000014_10.png training/image_3/000014_10.png 16 | training/image_2/000015_10.png training/image_3/000015_10.png 17 | training/image_2/000016_10.png training/image_3/000016_10.png 18 | training/image_2/000017_10.png training/image_3/000017_10.png 19 | training/image_2/000018_10.png training/image_3/000018_10.png 20 | training/image_2/000019_10.png training/image_3/000019_10.png 21 | training/image_2/000020_10.png training/image_3/000020_10.png 22 | training/image_2/000021_10.png training/image_3/000021_10.png 23 | training/image_2/000022_10.png training/image_3/000022_10.png 24 | training/image_2/000023_10.png training/image_3/000023_10.png 25 | training/image_2/000024_10.png training/image_3/000024_10.png 26 | training/image_2/000025_10.png training/image_3/000025_10.png 27 | training/image_2/000026_10.png training/image_3/000026_10.png 28 | training/image_2/000027_10.png training/image_3/000027_10.png 29 | training/image_2/000028_10.png training/image_3/000028_10.png 30 | training/image_2/000029_10.png training/image_3/000029_10.png 31 | training/image_2/000030_10.png training/image_3/000030_10.png 32 | training/image_2/000031_10.png training/image_3/000031_10.png 33 | training/image_2/000032_10.png training/image_3/000032_10.png 34 | training/image_2/000033_10.png training/image_3/000033_10.png 35 | training/image_2/000034_10.png training/image_3/000034_10.png 36 | training/image_2/000035_10.png training/image_3/000035_10.png 37 | training/image_2/000036_10.png training/image_3/000036_10.png 38 | training/image_2/000037_10.png training/image_3/000037_10.png 39 | training/image_2/000038_10.png training/image_3/000038_10.png 40 | training/image_2/000039_10.png training/image_3/000039_10.png 41 | training/image_2/000040_10.png training/image_3/000040_10.png 42 | training/image_2/000041_10.png training/image_3/000041_10.png 43 | training/image_2/000042_10.png training/image_3/000042_10.png 44 | training/image_2/000043_10.png training/image_3/000043_10.png 45 | training/image_2/000044_10.png training/image_3/000044_10.png 46 | training/image_2/000045_10.png training/image_3/000045_10.png 47 | training/image_2/000046_10.png training/image_3/000046_10.png 48 | training/image_2/000047_10.png training/image_3/000047_10.png 49 | training/image_2/000048_10.png training/image_3/000048_10.png 50 | training/image_2/000049_10.png training/image_3/000049_10.png 51 | training/image_2/000050_10.png training/image_3/000050_10.png 52 | training/image_2/000051_10.png training/image_3/000051_10.png 53 | training/image_2/000052_10.png training/image_3/000052_10.png 54 | training/image_2/000053_10.png training/image_3/000053_10.png 55 | training/image_2/000054_10.png training/image_3/000054_10.png 56 | training/image_2/000055_10.png training/image_3/000055_10.png 57 | training/image_2/000056_10.png training/image_3/000056_10.png 58 | training/image_2/000057_10.png training/image_3/000057_10.png 59 | training/image_2/000058_10.png training/image_3/000058_10.png 60 | training/image_2/000059_10.png training/image_3/000059_10.png 61 | training/image_2/000060_10.png training/image_3/000060_10.png 62 | training/image_2/000061_10.png training/image_3/000061_10.png 63 | training/image_2/000062_10.png training/image_3/000062_10.png 64 | training/image_2/000063_10.png training/image_3/000063_10.png 65 | training/image_2/000064_10.png training/image_3/000064_10.png 66 | training/image_2/000065_10.png training/image_3/000065_10.png 67 | training/image_2/000066_10.png training/image_3/000066_10.png 68 | training/image_2/000067_10.png training/image_3/000067_10.png 69 | training/image_2/000068_10.png training/image_3/000068_10.png 70 | training/image_2/000069_10.png training/image_3/000069_10.png 71 | training/image_2/000070_10.png training/image_3/000070_10.png 72 | training/image_2/000071_10.png training/image_3/000071_10.png 73 | training/image_2/000072_10.png training/image_3/000072_10.png 74 | training/image_2/000073_10.png training/image_3/000073_10.png 75 | training/image_2/000074_10.png training/image_3/000074_10.png 76 | training/image_2/000075_10.png training/image_3/000075_10.png 77 | training/image_2/000076_10.png training/image_3/000076_10.png 78 | training/image_2/000077_10.png training/image_3/000077_10.png 79 | training/image_2/000078_10.png training/image_3/000078_10.png 80 | training/image_2/000079_10.png training/image_3/000079_10.png 81 | training/image_2/000080_10.png training/image_3/000080_10.png 82 | training/image_2/000081_10.png training/image_3/000081_10.png 83 | training/image_2/000082_10.png training/image_3/000082_10.png 84 | training/image_2/000083_10.png training/image_3/000083_10.png 85 | training/image_2/000084_10.png training/image_3/000084_10.png 86 | training/image_2/000085_10.png training/image_3/000085_10.png 87 | training/image_2/000086_10.png training/image_3/000086_10.png 88 | training/image_2/000087_10.png training/image_3/000087_10.png 89 | training/image_2/000088_10.png training/image_3/000088_10.png 90 | training/image_2/000089_10.png training/image_3/000089_10.png 91 | training/image_2/000090_10.png training/image_3/000090_10.png 92 | training/image_2/000091_10.png training/image_3/000091_10.png 93 | training/image_2/000092_10.png training/image_3/000092_10.png 94 | training/image_2/000093_10.png training/image_3/000093_10.png 95 | training/image_2/000094_10.png training/image_3/000094_10.png 96 | training/image_2/000095_10.png training/image_3/000095_10.png 97 | training/image_2/000096_10.png training/image_3/000096_10.png 98 | training/image_2/000097_10.png training/image_3/000097_10.png 99 | training/image_2/000098_10.png training/image_3/000098_10.png 100 | training/image_2/000099_10.png training/image_3/000099_10.png 101 | training/image_2/000100_10.png training/image_3/000100_10.png 102 | training/image_2/000101_10.png training/image_3/000101_10.png 103 | training/image_2/000102_10.png training/image_3/000102_10.png 104 | training/image_2/000103_10.png training/image_3/000103_10.png 105 | training/image_2/000104_10.png training/image_3/000104_10.png 106 | training/image_2/000105_10.png training/image_3/000105_10.png 107 | training/image_2/000106_10.png training/image_3/000106_10.png 108 | training/image_2/000107_10.png training/image_3/000107_10.png 109 | training/image_2/000108_10.png training/image_3/000108_10.png 110 | training/image_2/000109_10.png training/image_3/000109_10.png 111 | training/image_2/000110_10.png training/image_3/000110_10.png 112 | training/image_2/000111_10.png training/image_3/000111_10.png 113 | training/image_2/000112_10.png training/image_3/000112_10.png 114 | training/image_2/000113_10.png training/image_3/000113_10.png 115 | training/image_2/000114_10.png training/image_3/000114_10.png 116 | training/image_2/000115_10.png training/image_3/000115_10.png 117 | training/image_2/000116_10.png training/image_3/000116_10.png 118 | training/image_2/000117_10.png training/image_3/000117_10.png 119 | training/image_2/000118_10.png training/image_3/000118_10.png 120 | training/image_2/000119_10.png training/image_3/000119_10.png 121 | training/image_2/000120_10.png training/image_3/000120_10.png 122 | training/image_2/000121_10.png training/image_3/000121_10.png 123 | training/image_2/000122_10.png training/image_3/000122_10.png 124 | training/image_2/000123_10.png training/image_3/000123_10.png 125 | training/image_2/000124_10.png training/image_3/000124_10.png 126 | training/image_2/000125_10.png training/image_3/000125_10.png 127 | training/image_2/000126_10.png training/image_3/000126_10.png 128 | training/image_2/000127_10.png training/image_3/000127_10.png 129 | training/image_2/000128_10.png training/image_3/000128_10.png 130 | training/image_2/000129_10.png training/image_3/000129_10.png 131 | training/image_2/000130_10.png training/image_3/000130_10.png 132 | training/image_2/000131_10.png training/image_3/000131_10.png 133 | training/image_2/000132_10.png training/image_3/000132_10.png 134 | training/image_2/000133_10.png training/image_3/000133_10.png 135 | training/image_2/000134_10.png training/image_3/000134_10.png 136 | training/image_2/000135_10.png training/image_3/000135_10.png 137 | training/image_2/000136_10.png training/image_3/000136_10.png 138 | training/image_2/000137_10.png training/image_3/000137_10.png 139 | training/image_2/000138_10.png training/image_3/000138_10.png 140 | training/image_2/000139_10.png training/image_3/000139_10.png 141 | training/image_2/000140_10.png training/image_3/000140_10.png 142 | training/image_2/000141_10.png training/image_3/000141_10.png 143 | training/image_2/000142_10.png training/image_3/000142_10.png 144 | training/image_2/000143_10.png training/image_3/000143_10.png 145 | training/image_2/000144_10.png training/image_3/000144_10.png 146 | training/image_2/000145_10.png training/image_3/000145_10.png 147 | training/image_2/000146_10.png training/image_3/000146_10.png 148 | training/image_2/000147_10.png training/image_3/000147_10.png 149 | training/image_2/000148_10.png training/image_3/000148_10.png 150 | training/image_2/000149_10.png training/image_3/000149_10.png 151 | training/image_2/000150_10.png training/image_3/000150_10.png 152 | training/image_2/000151_10.png training/image_3/000151_10.png 153 | training/image_2/000152_10.png training/image_3/000152_10.png 154 | training/image_2/000153_10.png training/image_3/000153_10.png 155 | training/image_2/000154_10.png training/image_3/000154_10.png 156 | training/image_2/000155_10.png training/image_3/000155_10.png 157 | training/image_2/000156_10.png training/image_3/000156_10.png 158 | training/image_2/000157_10.png training/image_3/000157_10.png 159 | training/image_2/000158_10.png training/image_3/000158_10.png 160 | training/image_2/000159_10.png training/image_3/000159_10.png 161 | training/image_2/000160_10.png training/image_3/000160_10.png 162 | training/image_2/000161_10.png training/image_3/000161_10.png 163 | training/image_2/000162_10.png training/image_3/000162_10.png 164 | training/image_2/000163_10.png training/image_3/000163_10.png 165 | training/image_2/000164_10.png training/image_3/000164_10.png 166 | training/image_2/000165_10.png training/image_3/000165_10.png 167 | training/image_2/000166_10.png training/image_3/000166_10.png 168 | training/image_2/000167_10.png training/image_3/000167_10.png 169 | training/image_2/000168_10.png training/image_3/000168_10.png 170 | training/image_2/000169_10.png training/image_3/000169_10.png 171 | training/image_2/000170_10.png training/image_3/000170_10.png 172 | training/image_2/000171_10.png training/image_3/000171_10.png 173 | training/image_2/000172_10.png training/image_3/000172_10.png 174 | training/image_2/000173_10.png training/image_3/000173_10.png 175 | training/image_2/000174_10.png training/image_3/000174_10.png 176 | training/image_2/000175_10.png training/image_3/000175_10.png 177 | training/image_2/000176_10.png training/image_3/000176_10.png 178 | training/image_2/000177_10.png training/image_3/000177_10.png 179 | training/image_2/000178_10.png training/image_3/000178_10.png 180 | training/image_2/000179_10.png training/image_3/000179_10.png 181 | training/image_2/000180_10.png training/image_3/000180_10.png 182 | training/image_2/000181_10.png training/image_3/000181_10.png 183 | training/image_2/000182_10.png training/image_3/000182_10.png 184 | training/image_2/000183_10.png training/image_3/000183_10.png 185 | training/image_2/000184_10.png training/image_3/000184_10.png 186 | training/image_2/000185_10.png training/image_3/000185_10.png 187 | training/image_2/000186_10.png training/image_3/000186_10.png 188 | training/image_2/000187_10.png training/image_3/000187_10.png 189 | training/image_2/000188_10.png training/image_3/000188_10.png 190 | training/image_2/000189_10.png training/image_3/000189_10.png 191 | training/image_2/000190_10.png training/image_3/000190_10.png 192 | training/image_2/000191_10.png training/image_3/000191_10.png 193 | training/image_2/000192_10.png training/image_3/000192_10.png 194 | training/image_2/000193_10.png training/image_3/000193_10.png 195 | training/image_2/000194_10.png training/image_3/000194_10.png 196 | training/image_2/000195_10.png training/image_3/000195_10.png 197 | training/image_2/000196_10.png training/image_3/000196_10.png 198 | training/image_2/000197_10.png training/image_3/000197_10.png 199 | training/image_2/000198_10.png training/image_3/000198_10.png 200 | training/image_2/000199_10.png training/image_3/000199_10.png 201 | -------------------------------------------------------------------------------- /utils/kitti_archives_to_download.txt: -------------------------------------------------------------------------------- 1 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_calib.zip 2 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0001/2011_09_26_drive_0001_sync.zip 3 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0002/2011_09_26_drive_0002_sync.zip 4 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0005/2011_09_26_drive_0005_sync.zip 5 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0009/2011_09_26_drive_0009_sync.zip 6 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0011/2011_09_26_drive_0011_sync.zip 7 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0013/2011_09_26_drive_0013_sync.zip 8 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0014/2011_09_26_drive_0014_sync.zip 9 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0015/2011_09_26_drive_0015_sync.zip 10 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0017/2011_09_26_drive_0017_sync.zip 11 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0018/2011_09_26_drive_0018_sync.zip 12 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0019/2011_09_26_drive_0019_sync.zip 13 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0020/2011_09_26_drive_0020_sync.zip 14 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0022/2011_09_26_drive_0022_sync.zip 15 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0023/2011_09_26_drive_0023_sync.zip 16 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0027/2011_09_26_drive_0027_sync.zip 17 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0028/2011_09_26_drive_0028_sync.zip 18 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0029/2011_09_26_drive_0029_sync.zip 19 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0032/2011_09_26_drive_0032_sync.zip 20 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0035/2011_09_26_drive_0035_sync.zip 21 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0036/2011_09_26_drive_0036_sync.zip 22 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0039/2011_09_26_drive_0039_sync.zip 23 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0046/2011_09_26_drive_0046_sync.zip 24 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0048/2011_09_26_drive_0048_sync.zip 25 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0051/2011_09_26_drive_0051_sync.zip 26 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0052/2011_09_26_drive_0052_sync.zip 27 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0056/2011_09_26_drive_0056_sync.zip 28 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0057/2011_09_26_drive_0057_sync.zip 29 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0059/2011_09_26_drive_0059_sync.zip 30 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0060/2011_09_26_drive_0060_sync.zip 31 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0061/2011_09_26_drive_0061_sync.zip 32 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0064/2011_09_26_drive_0064_sync.zip 33 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0070/2011_09_26_drive_0070_sync.zip 34 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0079/2011_09_26_drive_0079_sync.zip 35 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0084/2011_09_26_drive_0084_sync.zip 36 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0086/2011_09_26_drive_0086_sync.zip 37 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0087/2011_09_26_drive_0087_sync.zip 38 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0091/2011_09_26_drive_0091_sync.zip 39 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0093/2011_09_26_drive_0093_sync.zip 40 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0095/2011_09_26_drive_0095_sync.zip 41 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0096/2011_09_26_drive_0096_sync.zip 42 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0101/2011_09_26_drive_0101_sync.zip 43 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0104/2011_09_26_drive_0104_sync.zip 44 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0106/2011_09_26_drive_0106_sync.zip 45 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0113/2011_09_26_drive_0113_sync.zip 46 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_26_drive_0117/2011_09_26_drive_0117_sync.zip 47 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_28_calib.zip 48 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_28_drive_0001/2011_09_28_drive_0001_sync.zip 49 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_28_drive_0002/2011_09_28_drive_0002_sync.zip 50 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_29_calib.zip 51 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_29_drive_0004/2011_09_29_drive_0004_sync.zip 52 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_29_drive_0026/2011_09_29_drive_0026_sync.zip 53 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_29_drive_0071/2011_09_29_drive_0071_sync.zip 54 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_calib.zip 55 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0016/2011_09_30_drive_0016_sync.zip 56 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0018/2011_09_30_drive_0018_sync.zip 57 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0020/2011_09_30_drive_0020_sync.zip 58 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0027/2011_09_30_drive_0027_sync.zip 59 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0028/2011_09_30_drive_0028_sync.zip 60 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0033/2011_09_30_drive_0033_sync.zip 61 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_09_30_drive_0034/2011_09_30_drive_0034_sync.zip 62 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_10_03_calib.zip 63 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_10_03_drive_0027/2011_10_03_drive_0027_sync.zip 64 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_10_03_drive_0034/2011_10_03_drive_0034_sync.zip 65 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_10_03_drive_0042/2011_10_03_drive_0042_sync.zip 66 | http://kitti.is.tue.mpg.de/kitti/raw_data/2011_10_03_drive_0047/2011_10_03_drive_0047_sync.zip 67 | -------------------------------------------------------------------------------- /utils/scene_dataloader.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | import torchvision.transforms as transforms 4 | import torch.utils.data as data 5 | import numpy as np 6 | from torch.autograd import Variable 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.optim as optim 10 | import argparse 11 | import random 12 | from PIL import Image 13 | import matplotlib.pyplot as plt 14 | import cv2 15 | 16 | def get_kitti_cycle_data(file_path_train, path): 17 | f_train = open(file_path_train) 18 | former_left_image_train = list() 19 | latter_left_image_train = list() 20 | former_right_image_train = list() 21 | latter_right_image_train = list() 22 | 23 | for line in f_train: 24 | former_left_image_train.append(path+line.split()[0]) 25 | latter_left_image_train.append(path+line.split()[2]) 26 | former_right_image_train.append(path+line.split()[1]) 27 | latter_right_image_train.append(path+line.split()[3]) 28 | 29 | return former_left_image_train, latter_left_image_train, former_right_image_train, latter_right_image_train 30 | 31 | def get_data(file_path_test, path): 32 | f_test = open(file_path_test) 33 | left_image_test = list() 34 | right_image_test = list() 35 | 36 | for line in f_test: 37 | left_image_test.append(path+line.split()[0]) 38 | right_image_test.append(path+line.split()[1]) 39 | 40 | return left_image_test, right_image_test 41 | 42 | def get_flow_data(file_path_test, path): 43 | f_test = open(file_path_test) 44 | flow_test = list() 45 | former_image_test = list() 46 | latter_image_test = list() 47 | 48 | for line in f_test: 49 | former_image_test.append(path+line.split()[0]) 50 | latter_image_test.append(path+line.split()[1]) 51 | flow_test.append(path+line.split()[2]) 52 | 53 | return former_image_test, latter_image_test, flow_test 54 | 55 | def get_transform(param): 56 | return transforms.Compose([ 57 | transforms.Resize([param.input_height, param.input_width]), 58 | transforms.ToTensor() 59 | ]) 60 | 61 | class myCycleImageFolder(data.Dataset): 62 | def __init__(self, left1, left2, right1, right2, training, param): 63 | self.right1 = right1 64 | self.left1 = left1 65 | self.right2 = right2 66 | self.left2 = left2 67 | self.training = training 68 | self.param = param 69 | 70 | def __getitem__(self, index): 71 | left1 = self.left1[index] 72 | right1 = self.right1[index] 73 | left2 = self.left2[index] 74 | right2 = self.right2[index] 75 | param = self.param 76 | left_image_1 = Image.open(left1).convert('RGB') 77 | right_image_1 = Image.open(right1).convert('RGB') 78 | left_image_2 = Image.open(left2).convert('RGB') 79 | right_image_2 = Image.open(right2).convert('RGB') 80 | 81 | 82 | #augmentation 83 | if self.training: 84 | 85 | #randomly flip 86 | if random.uniform(0, 1) > 0.5: 87 | left_image_1 = left_image_1.transpose(Image.FLIP_LEFT_RIGHT) 88 | right_image_1 = right_image_1.transpose(Image.FLIP_LEFT_RIGHT) 89 | left_image_2 = left_image_2.transpose(Image.FLIP_LEFT_RIGHT) 90 | right_image_2 = right_image_2.transpose(Image.FLIP_LEFT_RIGHT) 91 | 92 | #randomly shift gamma 93 | if random.uniform(0, 1) > 0.5: 94 | gamma = random.uniform(0.8, 1.2) 95 | left_image_1 = Image.fromarray(np.clip((np.array(left_image_1) ** gamma), 0, 255).astype('uint8'), 'RGB') 96 | right_image_1 = Image.fromarray(np.clip((np.array(right_image_1) ** gamma), 0, 255).astype('uint8'), 'RGB') 97 | left_image_2 = Image.fromarray(np.clip((np.array(left_image_2) ** gamma), 0, 255).astype('uint8'), 'RGB') 98 | right_image_2 = Image.fromarray(np.clip((np.array(right_image_2) ** gamma), 0, 255).astype('uint8'), 'RGB') 99 | 100 | #randomly shift brightness 101 | if random.uniform(0, 1) > 0.5: 102 | brightness = random.uniform(0.5, 2.0) 103 | left_image_1 = Image.fromarray(np.clip((np.array(left_image_1) * brightness), 0, 255).astype('uint8'), 'RGB') 104 | right_image_1 = Image.fromarray(np.clip((np.array(right_image_1) * brightness), 0, 255).astype('uint8'), 'RGB') 105 | left_image_2 = Image.fromarray(np.clip((np.array(left_image_2) * brightness), 0, 255).astype('uint8'), 'RGB') 106 | right_image_2 = Image.fromarray(np.clip((np.array(right_image_2) * brightness), 0, 255).astype('uint8'), 'RGB') 107 | 108 | #randomly shift color 109 | if random.uniform(0, 1) > 0.5: 110 | colors = [random.uniform(0.8, 1.2) for i in range(3)] 111 | shape = np.array(left_image_1).shape 112 | white = np.ones((shape[0], shape[1])) 113 | color_image = np.stack([white * colors[i] for i in range(3)], axis=2) 114 | left_image_1 = Image.fromarray(np.clip((np.array(left_image_1) * color_image), 0, 255).astype('uint8'), 'RGB') 115 | right_image_1 = Image.fromarray(np.clip((np.array(right_image_1) * color_image), 0, 255).astype('uint8'), 'RGB') 116 | left_image_2 = Image.fromarray(np.clip((np.array(left_image_2) * color_image), 0, 255).astype('uint8'), 'RGB') 117 | right_image_2 = Image.fromarray(np.clip((np.array(right_image_2) * color_image), 0, 255).astype('uint8'), 'RGB') 118 | 119 | 120 | #transforms 121 | process = get_transform(param) 122 | left_image_1 = process(left_image_1) 123 | right_image_1 = process(right_image_1) 124 | left_image_2 = process(left_image_2) 125 | right_image_2 = process(right_image_2) 126 | 127 | return left_image_1, left_image_2, right_image_1, right_image_2 128 | def __len__(self): 129 | return len(self.left1) 130 | 131 | class myImageFolder(data.Dataset): 132 | def __init__(self, left, right, flow, param): 133 | self.right = right 134 | self.left = left 135 | self.flow = flow 136 | self.param = param 137 | 138 | def __getitem__(self, index): 139 | left = self.left[index] 140 | right = self.right[index] 141 | param = self.param 142 | left_image = Image.open(left).convert('RGB') 143 | right_image = Image.open(right).convert('RGB') 144 | 145 | process = get_transform(param) 146 | left_image = process(left_image) 147 | right_image = process(right_image) 148 | 149 | if self.flow is not None: 150 | flow = self.flow[index] 151 | flow_image = cv2.imread(flow, -1) 152 | h, w, _ = flow_image.shape 153 | flo_img = flow_image[:,:,2:0:-1].astype(np.float32) 154 | invalid = (flow_image[:,:,0] == 0) 155 | 156 | flo_img = (flo_img - 32768) / 64 157 | flo_img[np.abs(flo_img) < 1e-10] = 1e-10 158 | flo_img[invalid, :] = 0 159 | 160 | f = torch.from_numpy(flo_img.transpose((2,0,1))) 161 | mask = torch.from_numpy((flow_image[:,:,0] == 1).astype(np.float32)).type(torch.FloatTensor) 162 | 163 | return left_image, right_image, f.type(torch.FloatTensor), mask, h, w 164 | 165 | return left_image, right_image 166 | 167 | def __len__(self): 168 | return len(self.left) -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | import torchvision.transforms as transforms 4 | import torch.utils.data as data 5 | import numpy as np 6 | from torch.autograd import Variable 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.optim as optim 10 | import argparse 11 | import random 12 | from PIL import Image 13 | import matplotlib.pyplot as plt 14 | import cv2 15 | from models.networks.submodules import * 16 | from models.networks.resample2d_package.resample2d import Resample2d 17 | 18 | def gradient_x(img): 19 | gx = torch.add(img[:,:,:-1,:], -1, img[:,:,1:,:]) 20 | return gx 21 | 22 | def gradient_y(img): 23 | gy = torch.add(img[:,:,:,:-1], -1, img[:,:,:,1:]) 24 | return gy 25 | 26 | def get_disparity_smoothness(disp, pyramid): 27 | disp_gradients_x = [gradient_x(d) for d in disp] 28 | disp_gradients_y = [gradient_y(d) for d in disp] 29 | 30 | image_gradients_x = [gradient_x(img) for img in pyramid] 31 | image_gradients_y = [gradient_y(img) for img in pyramid] 32 | 33 | weights_x = [torch.exp(-torch.mean(torch.abs(g), 1, keepdim=True)) for g in image_gradients_x] 34 | weights_y = [torch.exp(-torch.mean(torch.abs(g), 1, keepdim=True)) for g in image_gradients_y] 35 | 36 | smoothness_x = [disp_gradients_x[i] * weights_x[i] for i in range(4)] 37 | smoothness_y = [disp_gradients_y[i] * weights_y[i] for i in range(4)] 38 | 39 | return smoothness_x + smoothness_y 40 | 41 | def SSIM(x, y): 42 | C1 = 0.01 ** 2 43 | C2 = 0.03 ** 2 44 | 45 | mu_x = F.avg_pool2d(x, 3, 1, 0) 46 | mu_y = F.avg_pool2d(y, 3, 1, 0) 47 | 48 | #(input, kernel, stride, padding) 49 | sigma_x = F.avg_pool2d(x ** 2, 3, 1, 0) - mu_x ** 2 50 | sigma_y = F.avg_pool2d(y ** 2, 3, 1, 0) - mu_y ** 2 51 | sigma_xy = F.avg_pool2d(x * y , 3, 1, 0) - mu_x * mu_y 52 | 53 | SSIM_n = (2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2) 54 | SSIM_d = (mu_x ** 2 + mu_y ** 2 + C1) * (sigma_x + sigma_y + C2) 55 | 56 | SSIM = SSIM_n / SSIM_d 57 | 58 | return torch.clamp((1 - SSIM) / 2, 0, 1) 59 | 60 | def cal_grad2_error(flo, image, beta): 61 | """ 62 | Calculate the image-edge-aware second-order smoothness loss for flo 63 | """ 64 | 65 | def gradient(pred): 66 | D_dy = pred[:, :, 1:, :] - pred[:, :, :-1, :] 67 | D_dx = pred[:, :, :, 1:] - pred[:, :, :, :-1] 68 | return D_dx, D_dy 69 | 70 | 71 | img_grad_x, img_grad_y = gradient(image) 72 | weights_x = torch.exp(-10.0 * torch.mean(torch.abs(img_grad_x), 1, keepdim=True)) 73 | weights_y = torch.exp(-10.0 * torch.mean(torch.abs(img_grad_y), 1, keepdim=True)) 74 | 75 | dx, dy = gradient(flo) 76 | dx2, dxdy = gradient(dx) 77 | dydx, dy2 = gradient(dy) 78 | 79 | return (torch.mean(beta*weights_x[:,:, :, 1:]*torch.abs(dx2)) + torch.mean(beta*weights_y[:, :, 1:, :]*torch.abs(dy2))) / 2.0 80 | 81 | def warp_2(est, img, occ_mask, args): 82 | l1_warp2 = torch.abs(est - img) * occ_mask 83 | l1_reconstruction_loss_warp2 = torch.mean(l1_warp2) / torch.mean(occ_mask) 84 | ssim_warp2 = SSIM(est * occ_mask, img * occ_mask) 85 | ssim_loss_warp2 = torch.mean(ssim_warp2) / torch.mean(occ_mask) 86 | image_loss_warp2 = args.alpha_image_loss * ssim_loss_warp2 + (1 - args.alpha_image_loss) * l1_reconstruction_loss_warp2 87 | return image_loss_warp2 88 | 89 | def create_mask(tensor, paddings): 90 | shape = tensor.shape 91 | inner_width = shape[3] - (paddings[1][0] + paddings[1][1]) 92 | inner_height = shape[2] - (paddings[0][0] + paddings[0][1]) 93 | inner = Variable(torch.ones((inner_height, inner_width)).cuda()) 94 | 95 | mask2d = nn.ZeroPad2d((paddings[1][0], paddings[1][1], paddings[0][0], paddings[0][1]))(inner) 96 | mask3d = mask2d.unsqueeze(0).repeat(shape[0], 1, 1) 97 | mask4d = mask3d.unsqueeze(1) 98 | return mask4d.detach() 99 | 100 | def create_border_mask(tensor, border_ratio = 0.1): 101 | num_batch, _, height, width = tensor.shape 102 | sz = np.ceil(height * border_ratio).astype(np.int).item(0) 103 | border_mask = create_mask(tensor, [[sz, sz], [sz, sz]]) 104 | return border_mask.detach() 105 | 106 | def length_sq(x): 107 | return torch.sum(x**2, 1, keepdim=True) 108 | 109 | def create_outgoing_mask(flow): 110 | num_batch, channel, height, width = flow.shape 111 | 112 | grid_x = torch.arange(width).view(1, 1, width) 113 | grid_x = grid_x.repeat(num_batch, height, 1) 114 | grid_y = torch.arange(height).view(1, height, 1) 115 | grid_y = grid_y.repeat(num_batch, 1, width) 116 | 117 | flow_u, flow_v = torch.unbind(flow, 1) 118 | pos_x = grid_x.type(torch.FloatTensor) + flow_u.data.cpu() 119 | pos_y = grid_y.type(torch.FloatTensor) + flow_v.data.cpu() 120 | inside_x = (pos_x <= (width - 1)) & (pos_x >= 0.0) 121 | inside_y = (pos_y <= (height - 1)) & (pos_y >= 0.0) 122 | inside = inside_x & inside_y 123 | return inside.type(torch.FloatTensor).unsqueeze(1) 124 | 125 | def get_mask(forward, backward, border_mask): 126 | flow_fw = forward 127 | flow_bw = backward 128 | mag_sq = length_sq(flow_fw) + length_sq(flow_bw) 129 | 130 | flow_bw_warped = Resample2d()(flow_bw, flow_fw) 131 | flow_fw_warped = Resample2d()(flow_fw, flow_bw) 132 | flow_diff_fw = flow_fw + flow_bw_warped 133 | flow_diff_bw = flow_bw + flow_fw_warped 134 | occ_thresh = 0.01 * mag_sq + 0.5 135 | fb_occ_fw = (length_sq(flow_diff_fw) > occ_thresh).type(torch.cuda.FloatTensor) 136 | fb_occ_bw = (length_sq(flow_diff_bw) > occ_thresh).type(torch.cuda.FloatTensor) 137 | 138 | if border_mask is None: 139 | mask_fw = create_outgoing_mask(flow_fw) 140 | mask_bw = create_outgoing_mask(flow_bw) 141 | else: 142 | mask_fw = border_mask 143 | mask_bw = border_mask 144 | fw = mask_fw * (1 - fb_occ_fw) 145 | bw = mask_bw * (1 - fb_occ_bw) 146 | 147 | return fw, bw, flow_diff_fw, flow_diff_bw 148 | 149 | def make_pyramid(image, num_scales): 150 | scale_image = [Variable(image.cuda())] 151 | height, width = image.shape[2:] 152 | 153 | for i in range(num_scales - 1): 154 | new = [] 155 | for j in range(image.shape[0]): 156 | ratio = 2 ** (i+1) 157 | nh = height // ratio 158 | nw = width // ratio 159 | tmp = transforms.ToPILImage()(image[j]).convert('RGB') 160 | tmp = transforms.Resize([nh, nw])(tmp) 161 | tmp = transforms.ToTensor()(tmp) 162 | new.append(tmp.unsqueeze(0)) 163 | this = torch.cat(new, 0) 164 | scale_image.append(Variable(this.cuda())) 165 | 166 | return scale_image 167 | 168 | def evaluate_flow(flow, flow_gt, valid_mask=None): 169 | 170 | if valid_mask is None: 171 | tmp = np.multiply(flow_gt[0,:,:], flow_gt[1,:,:]) 172 | valid_mask = np.ceil(np.clip(np.abs(tmp), 0, 1)) 173 | 174 | N = np.sum(valid_mask) 175 | 176 | u = flow[0, :, :] 177 | v = flow[1, :, :] 178 | 179 | u_gt = flow_gt[0, :, :] 180 | v_gt = flow_gt[1, :, :] 181 | 182 | ### compute_EPE 183 | du = u - u_gt 184 | dv = v - v_gt 185 | 186 | du2 = np.multiply(du, du) 187 | dv2 = np.multiply(dv, dv) 188 | 189 | EPE = np.multiply(np.sqrt(du2 + dv2), valid_mask) 190 | EPE_avg = np.sum(EPE) / N 191 | 192 | ### compute FL 193 | bad_pixels = np.logical_and( 194 | EPE > 3, 195 | (EPE / np.sqrt(np.sum(np.square(flow_gt), axis=0)) + 1e-5) > 0.05) 196 | FL_avg = bad_pixels.sum() / valid_mask.sum() 197 | 198 | return EPE_avg, FL_avg 199 | --------------------------------------------------------------------------------