├── .gitignore ├── AverageMeter.py ├── Colab_DAIN.ipynb ├── LICENSE ├── MegaDepth ├── LICENSE ├── MegaDepth_model.py ├── README.md ├── SDR_compute.py ├── __init__.py ├── data │ ├── __init__.py │ ├── aligned_data_loader.py │ ├── base_data_loader.py │ ├── data_loader.py │ └── image_folder.py ├── models │ ├── HG_model.py │ ├── __init__.py │ ├── base_model.py │ └── models.py ├── options │ ├── __init__.py │ ├── base_options.py │ ├── test_options.py │ └── train_options.py ├── pytorch_DIW_scratch.py ├── rmse_error_main.py └── util │ ├── __init__.py │ ├── html.py │ ├── image_pool.py │ ├── png.py │ ├── util.py │ └── visualizer.py ├── PWCNet ├── PWCNet.py ├── __init__.py ├── correlation_package_pytorch1_0 │ ├── __init__.py │ ├── build.sh │ ├── clean.sh │ ├── correlation.py │ ├── correlation_cuda.cc │ ├── correlation_cuda_kernel.cu │ ├── correlation_cuda_kernel.cuh │ └── setup.py └── models │ ├── PWCNet.py │ └── __init__.py ├── README.md ├── Resblock ├── BasicBlock.py └── __init__.py ├── S2D_models ├── S2DF.py └── __init__.py ├── Stack.py ├── balancedsampler.py ├── colab_interpolate.py ├── datasets ├── Vimeo_90K_interp.py ├── __init__.py └── listdatasets.py ├── demo_MiddleBury.py ├── demo_MiddleBury_slowmotion.py ├── environment.yaml ├── loss_function.py ├── lr_scheduler.py ├── my_args.py ├── my_package ├── DepthFlowProjection │ ├── DepthFlowProjectionLayer.py │ ├── DepthFlowProjectionModule.py │ ├── __init__.py │ ├── depthflowprojection_cuda.cc │ ├── depthflowprojection_cuda_kernel.cu │ ├── depthflowprojection_cuda_kernel.cuh │ └── setup.py ├── FilterInterpolation │ ├── FilterInterpolationLayer.py │ ├── FilterInterpolationModule.py │ ├── __init__.py │ ├── filterinterpolation_cuda.cc │ ├── filterinterpolation_cuda_kernel.cu │ ├── filterinterpolation_cuda_kernel.cuh │ └── setup.py ├── FlowProjection │ ├── FlowProjectionLayer.py │ ├── FlowProjectionModule.py │ ├── __init__.py │ ├── flowprojection_cuda.cc │ ├── flowprojection_cuda_kernel.cu │ ├── flowprojection_cuda_kernel.cuh │ └── setup.py ├── Interpolation │ ├── InterpolationLayer.py │ ├── InterpolationModule.py │ ├── __init__.py │ ├── interpolation_cuda.cc │ ├── interpolation_cuda_kernel.cu │ ├── interpolation_cuda_kernel.cuh │ └── setup.py ├── InterpolationCh │ ├── InterpolationChLayer.py │ ├── InterpolationChModule.py │ ├── __init__.py │ ├── interpolationch_cuda.cc │ ├── interpolationch_cuda_kernel.cu │ ├── interpolationch_cuda_kernel.cuh │ └── setup.py ├── MinDepthFlowProjection │ ├── __init__.py │ ├── minDepthFlowProjectionLayer.py │ ├── minDepthFlowProjectionModule.py │ ├── mindepthflowprojection_cuda.cc │ ├── mindepthflowprojection_cuda_kernel.cu │ ├── mindepthflowprojection_cuda_kernel.cuh │ └── setup.py ├── SeparableConv │ ├── SeparableConvLayer.py │ ├── SeparableConvModule.py │ ├── __init__.py │ ├── separableconv_cuda.cc │ ├── separableconv_cuda_kernel.cu │ ├── separableconv_cuda_kernel.cuh │ └── setup.py ├── SeparableConvFlow │ ├── SeparableConvFlowLayer.py │ ├── SeparableConvFlowModule.py │ ├── __init__.py │ ├── separableconvflow_cuda.cc │ ├── separableconvflow_cuda_kernel.cu │ ├── separableconvflow_cuda_kernel.cuh │ └── setup.py ├── build.sh ├── clean.sh ├── compiler_args.py └── test_module.py ├── networks ├── DAIN.py ├── DAIN_slowmotion.py └── __init__.py └── train.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore Git here 2 | .git 3 | 4 | # But not these files... 5 | # !.gitignore 6 | 7 | checkpoints/test_local/opt.txt 8 | PWCNet/pwc_net.pth.tar 9 | MegaDepth/checkpoints/* 10 | model_weights/* 11 | MiddleBurySet/* 12 | 13 | .nfs* 14 | 15 | # Created by .ignore support plugin (hsz.mobi) 16 | ### Python template 17 | # Byte-compiled / optimized / DLL files 18 | __pycache__/ 19 | *.py[cod] 20 | *$py.class 21 | 22 | # C extensions 23 | *.so 24 | 25 | # Distribution / packaging 26 | .Python 27 | env/ 28 | build/ 29 | develop-eggs/ 30 | dist/ 31 | downloads/ 32 | eggs/ 33 | .eggs/ 34 | lib/ 35 | lib64/ 36 | parts/ 37 | sdist/ 38 | var/ 39 | *.egg-info/ 40 | .installed.cfg 41 | *.egg 42 | 43 | # PyInstaller 44 | # Usually these files are written by a python script from a template 45 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 46 | *.manifest 47 | *.spec 48 | 49 | # Installer logs 50 | pip-log.txt 51 | pip-delete-this-directory.txt 52 | 53 | # Unit test / coverage reports 54 | htmlcov/ 55 | .tox/ 56 | .coverage 57 | .coverage.* 58 | .cache 59 | nosetests.xml 60 | coverage.xml 61 | *,cover 62 | .hypothesis/ 63 | 64 | # Translations 65 | *.mo 66 | *.pot 67 | 68 | # Django stuff: 69 | *.log 70 | local_settings.py 71 | 72 | # Flask stuff: 73 | instance/ 74 | .webassets-cache 75 | 76 | # Scrapy stuff: 77 | .scrapy 78 | 79 | # Sphinx documentation 80 | docs/_build/ 81 | 82 | # PyBuilder 83 | target/ 84 | 85 | # IPython Notebook 86 | .ipynb_checkpoints 87 | 88 | # pyenv 89 | .python-version 90 | 91 | # celery beat schedule file 92 | celerybeat-schedule 93 | 94 | # dotenv 95 | .env 96 | 97 | # virtualenv 98 | venv/ 99 | ENV/ 100 | 101 | # Spyder project settings 102 | .spyderproject 103 | 104 | # Rope project settings 105 | .ropeproject 106 | ### VirtualEnv template 107 | # Virtualenv 108 | # http://iamzed.com/2009/05/07/a-primer-on-virtualenv/ 109 | .Python 110 | [Bb]in 111 | [Ii]nclude 112 | [Ll]ib 113 | [Ll]ib64 114 | [Ll]ocal 115 | [Ss]cripts 116 | pyvenv.cfg 117 | .venv 118 | pip-selfcheck.json 119 | ### JetBrains template 120 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and Webstorm 121 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 122 | 123 | # User-specific stuff: 124 | .idea/workspace.xml 125 | .idea/tasks.xml 126 | .idea/dictionaries 127 | .idea/vcs.xml 128 | .idea/jsLibraryMappings.xml 129 | 130 | # Sensitive or high-churn files: 131 | .idea/dataSources.ids 132 | .idea/dataSources.xml 133 | .idea/dataSources.local.xml 134 | .idea/sqlDataSources.xml 135 | .idea/dynamic.xml 136 | .idea/uiDesigner.xml 137 | 138 | # Gradle: 139 | .idea/gradle.xml 140 | .idea/libraries 141 | 142 | # Mongo Explorer plugin: 143 | .idea/mongoSettings.xml 144 | 145 | .idea/ 146 | 147 | ## File-based project format: 148 | *.iws 149 | 150 | ## Plugin-specific files: 151 | 152 | # IntelliJ 153 | /out/ 154 | 155 | # mpeltonen/sbt-idea plugin 156 | .idea_modules/ 157 | 158 | # JIRA plugin 159 | atlassian-ide-plugin.xml 160 | 161 | # Crashlytics plugin (for Android Studio and IntelliJ) 162 | com_crashlytics_export_strings.xml 163 | crashlytics.properties 164 | crashlytics-build.properties 165 | fabric.properties 166 | 167 | -------------------------------------------------------------------------------- /AverageMeter.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | class AverageMeter(object): 4 | """Computes and stores the average and current value""" 5 | def __init__(self): 6 | self.reset() 7 | 8 | def reset(self): 9 | self.val = 0 10 | self.avg = 0 11 | self.sum = 0 12 | self.count = 0 13 | 14 | def update(self, val, n=1): 15 | self.val = val 16 | self.sum += val * n 17 | self.count += n 18 | self.avg = self.sum / self.count 19 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Wenbo Bao 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MegaDepth/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Zhengqi Li 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MegaDepth/MegaDepth_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | from torch.autograd import Variable 4 | import numpy as np 5 | from .options.train_options import TrainOptions 6 | from .models.models import create_model 7 | __all__ = ['HourGlass'] 8 | 9 | 10 | 11 | def HourGlass(pretrained=None): 12 | """Constructs a ResNet-18 model. 13 | 14 | Args: 15 | pretrained (bool): If True, returns a model pre-trained on ImageNet 16 | """ 17 | 18 | opt = TrainOptions().parse() # set CUDA_VISIBLE_DEVICES before import torch 19 | model = create_model(opt,pretrained) 20 | #netG is the real nn.Module 21 | return model.netG 22 | -------------------------------------------------------------------------------- /MegaDepth/README.md: -------------------------------------------------------------------------------- 1 | # MegaDepth: Learning Single-View Depth Prediction from Internet Photos 2 | 3 | This is a code of the algorithm described in "MegaDepth: Learning Single-View Depth Prediction from Internet Photos, Z. Li and N. Snavely, CVPR 2018". The code skeleton is based on "https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix". If you use our code or models for academic purposes, please consider citing: 4 | 5 | @inproceedings{MDLi18, 6 | title={MegaDepth: Learning Single-View Depth Prediction from Internet Photos}, 7 | author={Zhengqi Li and Noah Snavely}, 8 | booktitle={Computer Vision and Pattern Recognition (CVPR)}, 9 | year={2018} 10 | } 11 | 12 | #### Examples of single-view depth predictions on the photos we randomly downloaded from Internet: 13 | 14 | 15 | 16 | 17 | 18 | #### Dependencies: 19 | * The code was written in Pytorch 0.2 and Python 2.7, but it should be easy to adapt it to Python 3 and latest Pytorch version if needed. 20 | * You might need skimage, h5py libraries installed for python before running the code. 21 | 22 | #### Single-view depth prediction on any Internet photo: 23 | * Download pretrained model from: http://www.cs.cornell.edu/projects/megadepth/dataset/models/best_generalization_net_G.pth and put it in "checkpoints/test_local/best_generalization_net_G.pth 24 | * In python file "models/HG_model.py", in init function, change to "model_parameters = self.load_network(model, 'G', 'best_generalization')" 25 | * run demo code 26 | ```bash 27 | python demo.py 28 | ``` 29 | You should see an inverse depth prediction saved as demo.png from an original photo demo.jpg. If you want to use RGB maps for visualization, like the figures in our paper, you have to install/run semantic segmentation from https://github.com/kazuto1011/pspnet-pytorch trained on ADE20K to mask out sky, because inconsistent depth prediction of unmasked sky will not make RGB visualization resonable. 30 | 31 | 32 | #### Evaluation on the MegaDepth test splits: 33 | * Download MegaDepth V1 dataset from project website: http://www.cs.cornell.edu/projects/megadepth/. 34 | * Download pretrained model (specific for MD dataset) from http://www.cs.cornell.edu/projects/megadepth/dataset/models/best_vanila_net_G.pth and put it in "checkpoints/test_local/best_vanila_net_G.pth" 35 | * Download test list files from http://www.cs.cornell.edu/projects/megadepth/dataset/data_lists/test_lists.tar.gz, it should include two folders corresponding to images with landscape and portrait orientations. 36 | * To compute scale invarance RMSE on MD testset, change the variable "dataset_root" in python file "rmse_error_main.py" to the root directory of MegaDepth_v1 folder, and change variable "test_list_dir_l" and "test_list_dir_p" to corresponding folder paths of test lists, and run: 37 | ```bash 38 | python rmse_error_main.py 39 | ``` 40 | * To compute Structure from Motion Disagreement Rate (SDR), change the variable "dataset_root" in python file "rmse_error_main.py" to the root directory of MegaDepth_v1 folder, and change variable "test_list_dir_l" and "test_list_dir_p" to corresponding folder paths of test lists, and run: 41 | ```bash 42 | python SDR_compute.py 43 | ``` 44 | * If you want to run our model on arbitrary Internet photos, please download pretrained model from http://www.cs.cornell.edu/projects/megadepth/dataset/models/best_generalization_net_G.pth, which has much better generalization ability (qualitatively speaking) to completely unknown scenes. 45 | 46 | -------------------------------------------------------------------------------- /MegaDepth/SDR_compute.py: -------------------------------------------------------------------------------- 1 | import time 2 | import torch 3 | import sys 4 | 5 | from options.train_options import TrainOptions 6 | opt = TrainOptions().parse() # set CUDA_VISIBLE_DEVICES before import torch 7 | from data.data_loader import CreateDataLoader_TEST 8 | from models.models import create_model 9 | 10 | dataset_root = "/phoenix/S6/zl548/" 11 | test_list_dir_l = dataset_root + '/MegaDpeth_code/test_list/landscape/' 12 | input_height = 240 13 | input_width = 320 14 | test_data_loader_l = CreateDataLoader_TEST(dataset_root, test_list_dir_l, input_height, input_width) 15 | test_dataset_l = test_data_loader_l.load_data() 16 | test_dataset_size_l = len(test_data_loader_l) 17 | print('========================= test L images = %d' % test_dataset_size_l) 18 | 19 | test_list_dir_p = dataset_root + '/MegaDpeth_code/test_list/portrait/' 20 | input_height = 320 21 | input_width = 240 22 | test_data_loader_p = CreateDataLoader_TEST(dataset_root, test_list_dir_p, input_height, input_width) 23 | test_dataset_p = test_data_loader_p.load_data() 24 | test_dataset_size_p = len(test_data_loader_p) 25 | print('========================= test P images = %d' % test_dataset_size_p) 26 | 27 | model = create_model(opt) 28 | 29 | batch_size = 32 30 | diw_index = 0 31 | total_steps = 0 32 | best_loss = 100 33 | 34 | error_list = [0 , 0, 0] 35 | total_list = [0 , 0, 0] 36 | 37 | list_l = range(test_dataset_size_l) 38 | list_p = range(test_dataset_size_p) 39 | 40 | 41 | def test_SDR(model): 42 | total_loss =0 43 | # count = 0 44 | print("============================= TEST SDR============================") 45 | model.switch_to_eval() 46 | diw_index = 0 47 | 48 | for i, data in enumerate(test_dataset_l): 49 | stacked_img = data['img_1'] 50 | targets = data['target_1'] 51 | error, samples = model.evaluate_SDR(stacked_img, targets) 52 | 53 | for j in range(0,3): 54 | error_list[j] += error[j] 55 | total_list[j] += samples[j] 56 | 57 | print("EQUAL ", error_list[0]/float(total_list[0])) 58 | print("INEQUAL ", error_list[1]/float(total_list[1])) 59 | print("TOTAL ",error_list[2]/float(total_list[2])) 60 | 61 | for i, data in enumerate(test_dataset_p): 62 | stacked_img = data['img_1'] 63 | targets = data['target_1'] 64 | 65 | error, samples = model.evaluate_SDR(stacked_img, targets) 66 | 67 | for j in range(0,3): 68 | error_list[j] += error[j] 69 | total_list[j] += samples[j] 70 | 71 | print("EQUAL ", error_list[0]/float(total_list[0])) 72 | print("INEQUAL ", error_list[1]/float(total_list[1])) 73 | print("TOTAL ",error_list[2]/float(total_list[2])) 74 | 75 | 76 | print("=========================================================SDR Summary =====================") 77 | print("Equal SDR:\t" , float(error_list[0])/ float(total_list[0])) 78 | print("Unequal SDR:\t" , float(error_list[1])/ float(total_list[1])) 79 | print("SDR:\t" , float(error_list[2])/ float(total_list[2])) 80 | 81 | 82 | print("WE ARE TESTING SDR!!!!") 83 | test_SDR(model) 84 | -------------------------------------------------------------------------------- /MegaDepth/__init__.py: -------------------------------------------------------------------------------- 1 | from .MegaDepth_model import * 2 | -------------------------------------------------------------------------------- /MegaDepth/data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/baowenbo/DAIN/7c727aca56760f8f1a57ffb603c8c0b0fd1b1349/MegaDepth/data/__init__.py -------------------------------------------------------------------------------- /MegaDepth/data/aligned_data_loader.py: -------------------------------------------------------------------------------- 1 | import random 2 | import numpy as np 3 | import torch.utils.data 4 | from data.base_data_loader import BaseDataLoader 5 | from data.image_folder import ImageFolder 6 | from data.image_folder import ImageFolder_TEST 7 | from builtins import object 8 | import sys 9 | import h5py 10 | 11 | 12 | class PairedData(object): 13 | def __init__(self, data_loader, flip): 14 | self.data_loader = data_loader 15 | # self.fineSize = fineSize 16 | # self.max_dataset_size = max_dataset_size 17 | self.flip = flip 18 | self.data_loader_iter = iter(self.data_loader) 19 | self.iter = 0 20 | 21 | 22 | def __iter__(self): 23 | self.data_loader_iter = iter(self.data_loader) 24 | self.iter = 0 25 | return self 26 | 27 | def __next__(self): 28 | self.iter += 1 29 | 30 | final_img, target_1 = next(self.data_loader_iter) 31 | 32 | return {'img_1': final_img, 'target_1': target_1} 33 | 34 | 35 | class AlignedDataLoader(BaseDataLoader): 36 | def __init__(self,_root, _list_dir, _input_height, _input_width, _is_flip, _shuffle): 37 | transform = None 38 | dataset = ImageFolder(root=_root, \ 39 | list_dir =_list_dir, input_height = _input_height, input_width = _input_width, transform=transform, is_flip = _is_flip) 40 | 41 | data_loader = torch.utils.data.DataLoader(dataset, batch_size= 16, shuffle= _shuffle, num_workers=int(3)) 42 | 43 | self.dataset = dataset 44 | flip = False 45 | self.paired_data = PairedData(data_loader, flip) 46 | 47 | def name(self): 48 | return 'RMSEDataLoader' 49 | 50 | def load_data(self): 51 | return self.paired_data 52 | 53 | def __len__(self): 54 | return len(self.dataset) 55 | 56 | 57 | 58 | class AlignedDataLoader_TEST(BaseDataLoader): 59 | def __init__(self,_root, _list_dir, _input_height, _input_width): 60 | 61 | dataset = ImageFolder_TEST(root=_root, \ 62 | list_dir =_list_dir, _input_height = _input_height, _input_width = _input_width) 63 | 64 | data_loader = torch.utils.data.DataLoader(dataset, batch_size= 1, shuffle= False, num_workers=int(3)) 65 | self.dataset = dataset 66 | flip = False 67 | self.paired_data = PairedData(data_loader, flip) 68 | 69 | def name(self): 70 | return 'TestSDRDataLoader' 71 | 72 | def load_data(self): 73 | return self.paired_data 74 | 75 | 76 | def __len__(self): 77 | return len(self.dataset) 78 | -------------------------------------------------------------------------------- /MegaDepth/data/base_data_loader.py: -------------------------------------------------------------------------------- 1 | 2 | class BaseDataLoader(): 3 | def __init__(self): 4 | pass 5 | 6 | # def initialize(self): 7 | # # self.opt = opt 8 | # pass 9 | 10 | def load_data(): 11 | return None 12 | 13 | 14 | 15 | -------------------------------------------------------------------------------- /MegaDepth/data/data_loader.py: -------------------------------------------------------------------------------- 1 | 2 | def CreateDataLoader(_root, _list_dir, _input_height, _input_width, is_flip = True, shuffle = True): 3 | data_loader = None 4 | from data.aligned_data_loader import AlignedDataLoader 5 | data_loader = AlignedDataLoader(_root, _list_dir, _input_height, _input_width, is_flip, shuffle) 6 | return data_loader 7 | 8 | def CreateDataLoader_TEST(_root, _list_dir, _input_height, _input_width): 9 | data_loader = None 10 | from data.aligned_data_loader import AlignedDataLoader_TEST 11 | data_loader = AlignedDataLoader_TEST(_root, _list_dir, _input_height, _input_width) 12 | 13 | return data_loader 14 | -------------------------------------------------------------------------------- /MegaDepth/data/image_folder.py: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Code from 3 | # https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py 4 | # Modified the original code so that it also loads images from the current 5 | # directory as well as the subdirectories 6 | ################################################################################ 7 | import h5py 8 | import torch.utils.data as data 9 | import pickle 10 | import numpy as np 11 | import torch 12 | import os, os.path 13 | import math, random 14 | import sys 15 | from skimage.transform import resize 16 | from skimage import io 17 | 18 | 19 | 20 | def make_dataset(list_dir): 21 | # subgroup_name1 = "/dataset/image_list/" 22 | file_name = list_dir + "imgs_MD.p" 23 | file_name_1 = open( file_name, "rb" ) 24 | images_list = pickle.load( file_name_1) 25 | file_name_1.close() 26 | 27 | file_name_t= list_dir + "targets_MD.p" 28 | file_name_2 = open( file_name_t, "rb" ) 29 | targets_list = pickle.load(file_name_2) 30 | file_name_2.close() 31 | return images_list, targets_list 32 | 33 | # test for si-RMSE 34 | class ImageFolder(data.Dataset): 35 | 36 | def __init__(self, root, list_dir, input_height, input_width, transform=None, 37 | loader=None, is_flip = True): 38 | # load image list from hdf5 39 | img_list , targets_list = make_dataset(list_dir) 40 | if len(img_list) == 0: 41 | raise(RuntimeError("Found 0 images in: " + root + "\n" 42 | "Supported image extensions are: " + ",".join(IMG_EXTENSIONS))) 43 | # img_list_1, img_list_2 = selfshuffle_dataset(img_list) 44 | self.root = root 45 | self.list_dir = list_dir 46 | self.img_list = img_list 47 | self.targets_list = targets_list 48 | self.transform = transform 49 | # self.loader = loader 50 | self.input_height = input_height 51 | self.input_width = input_width 52 | self.is_flip = is_flip 53 | 54 | 55 | def load_MD(self, img_path, depth_path): 56 | 57 | MD_img = np.float32(io.imread(img_path))/255.0 58 | 59 | hdf5_file_read = h5py.File(depth_path,'r') 60 | gt = hdf5_file_read.get('/depth') 61 | gt = np.array(gt) 62 | 63 | assert(gt.shape[0] == MD_img.shape[0]) 64 | assert(gt.shape[1] == MD_img.shape[1]) 65 | 66 | color_rgb = np.zeros((self.input_height,self.input_width,3)) 67 | MD_img = resize(MD_img, (self.input_height, self.input_width), order = 1) 68 | 69 | if len(MD_img.shape) == 2: 70 | color_rgb[:,:,0] = MD_img.copy() 71 | color_rgb[:,:,1] = MD_img.copy() 72 | color_rgb[:,:,2] = MD_img.copy() 73 | else: 74 | color_rgb = MD_img.copy() 75 | 76 | if np.sum(gt > 1e-8) > 10: 77 | gt[ gt > np.percentile(gt[gt > 1e-8], 98)] = 0 78 | gt[ gt < np.percentile(gt[gt > 1e-8], 1)] = 0 79 | 80 | max_depth = np.max(gt) + 1e-9 81 | gt = gt/max_depth 82 | gt = resize(gt, (self.input_height, self.input_width), order = 0) 83 | gt = gt*max_depth 84 | 85 | mask = np.float32(gt > 1e-8) 86 | 87 | color_rgb = np.ascontiguousarray(color_rgb) 88 | gt = np.ascontiguousarray(gt) 89 | mask = np.ascontiguousarray(mask) 90 | 91 | hdf5_file_read.close() 92 | 93 | return color_rgb, gt, mask 94 | 95 | def __getitem__(self, index): 96 | # 00xx/1/ 97 | targets_1 = {} 98 | # targets_1['L'] = [] 99 | targets_1['path'] = [] 100 | 101 | img_path_suff = self.img_list[index] 102 | targets_path_suff = self.targets_list[index] 103 | 104 | img_path = self.root + "/MegaDepth_v1/" + img_path_suff 105 | depth_path = self.root + "/MegaDepth_v1/" + targets_path_suff 106 | 107 | img, gt, mask = self.load_MD(img_path, depth_path) 108 | 109 | gt[mask < 0.1] = 1.0 110 | 111 | targets_1['path'] = targets_path_suff 112 | targets_1['gt_0'] = torch.from_numpy(gt).float() 113 | targets_1['mask_0'] = torch.from_numpy(mask).float() 114 | 115 | final_img = torch.from_numpy( np.transpose(img, (2,0,1)) ).contiguous().float() 116 | 117 | return final_img, targets_1 118 | 119 | def __len__(self): 120 | return len(self.img_list) 121 | 122 | 123 | # Test for SDR 124 | class ImageFolder_TEST(data.Dataset): 125 | 126 | def __init__(self, root, list_dir, _input_height, _input_width): 127 | # load image list from hdf5 128 | img_list , targets_list = make_dataset(list_dir) 129 | if len(img_list) == 0: 130 | raise(RuntimeError("Found 0 images in: " + root + "\n" 131 | "Supported image extensions are: " + ",".join(IMG_EXTENSIONS))) 132 | self.root = root 133 | self.list_dir = list_dir 134 | self.img_list = img_list 135 | self.input_height = _input_height 136 | self.input_width = _input_width 137 | self.half_window = 1 138 | 139 | def load_SfM_ORD(self, img_path, targets_path): 140 | 141 | sfm_image = np.float32(io.imread(img_path))/255.0 142 | resized_sfm_img = resize(sfm_image, (self.input_height, self.input_width), order = 1) 143 | 144 | color_rgb = np.zeros((self.input_height, self.input_width,3)) 145 | 146 | if len(sfm_image.shape) == 2: 147 | color_rgb[:,:,0] = resized_sfm_img.copy() 148 | color_rgb[:,:,1] = resized_sfm_img.copy() 149 | color_rgb[:,:,2] = resized_sfm_img.copy() 150 | else: 151 | color_rgb = resized_sfm_img.copy() 152 | 153 | if color_rgb.shape[2] == 4: 154 | return color_rgb, 0, 0 ,0, 0, 0 155 | 156 | hdf5_file_read = h5py.File(targets_path,'r') 157 | gt = hdf5_file_read.get('/SfM_features') 158 | gt = np.array(gt) 159 | 160 | y_A = np.round( gt[0,:] * float(self.input_height) ) 161 | x_A = np.round( gt[1,:] * float(self.input_width) ) 162 | y_B = np.round( gt[2,:] * float(self.input_height) ) 163 | x_B = np.round( gt[3,:] * float(self.input_width) ) 164 | ord_ = gt[4,:] 165 | 166 | hdf5_file_read.close() 167 | 168 | return color_rgb, y_A, x_A ,y_B, x_B, ord_ 169 | 170 | def __getitem__(self, index): 171 | # 00xx/1/ 172 | targets_1 = {} 173 | # targets_1['L'] = [] 174 | targets_1['path'] = [] 175 | targets_1['sdr_xA'] = [] 176 | targets_1['sdr_yA'] = [] 177 | targets_1['sdr_xB'] = [] 178 | targets_1['sdr_yB'] = [] 179 | targets_1['sdr_gt'] = [] 180 | 181 | img_path_suff = self.img_list[index] 182 | img_path = self.root + "/MegaDepth_v1/" + img_path_suff 183 | folder_name = img_path_suff.split('/')[-4] 184 | img_name = img_path_suff.split('/')[-1] 185 | sparse_sift_path = self.root + "/sparse_features/" + folder_name + "/" + img_name + ".h5" 186 | 187 | # no sift features 188 | if not os.path.isfile(sparse_sift_path) or not os.path.isfile(img_path): 189 | 190 | img = np.zeros((self.input_height, self.input_width,3)) 191 | targets_1['has_SfM_feature'] = False 192 | 193 | else: 194 | 195 | img, y_A, x_A ,y_B, x_B, ordinal = self.load_SfM_ORD(img_path, sparse_sift_path) 196 | 197 | targets_1['sdr_xA'].append(torch.from_numpy(x_A).long()) 198 | targets_1['sdr_yA'].append(torch.from_numpy(y_A).long()) 199 | targets_1['sdr_xB'].append(torch.from_numpy(x_B).long()) 200 | targets_1['sdr_yB'].append(torch.from_numpy(y_B).long()) 201 | targets_1['sdr_gt'].append(torch.from_numpy(ordinal).float()) 202 | targets_1['has_SfM_feature'] = True 203 | 204 | final_img = torch.from_numpy( np.transpose(img, (2,0,1)) ).contiguous().float() 205 | 206 | 207 | return final_img, targets_1 208 | 209 | 210 | 211 | def __len__(self): 212 | return len(self.img_list) 213 | 214 | 215 | 216 | -------------------------------------------------------------------------------- /MegaDepth/models/HG_model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import os 4 | from torch.autograd import Variable 5 | from .base_model import BaseModel 6 | import sys 7 | # import pytorch_DIW_scratch 8 | import MegaDepth.pytorch_DIW_scratch as pytorch_DIW_scratch 9 | 10 | class HGModel(BaseModel): 11 | def name(self): 12 | return 'HGModel' 13 | 14 | def __init__(self, opt,pretrained=None): 15 | BaseModel.initialize(self, opt) 16 | 17 | # print("===========================================LOADING Hourglass NETWORK====================================================") 18 | model = pytorch_DIW_scratch.pytorch_DIW_scratch 19 | # model_temp = model 20 | # model= torch.nn.parallel.DataParallel(model, device_ids = [0,1]) 21 | # model_parameters = self.load_network(model, 'G', 'best_vanila') 22 | if pretrained is None: 23 | # model_parameters = self.load_network(model, 'G', 'best_generalization') 24 | # 25 | # model.load_state_dict(model_parameters) 26 | # self.netG = model.cuda() 27 | self.netG = model 28 | # print("No weights loaded for Hourglass Network") 29 | else: 30 | pretrained_dict = torch.load(pretrained) 31 | 32 | model_dict = model.state_dict() 33 | # print(len(pretrained_dict)) 34 | # print(len(model_dict)) 35 | # 1. filter out unnecessary keys 36 | # the saved model contains a 'module.' prefix for the data.parallel reason 37 | pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items()} # and not k[:10]== 'rectifyNet'} 38 | # print(str(len(pretrained_dict)) + " are updated") 39 | # 2. overwrite entries in the existing state dict 40 | model_dict.update(pretrained_dict) 41 | # 3. load the new state dict 42 | model.load_state_dict(model_dict) 43 | pretrained_dict = None 44 | self.netG = model 45 | 46 | 47 | 48 | def batch_classify(self, z_A_arr, z_B_arr, ground_truth ): 49 | threashold = 1.1 50 | depth_ratio = torch.div(z_A_arr, z_B_arr) 51 | 52 | depth_ratio = depth_ratio.cpu() 53 | 54 | estimated_labels = torch.zeros(depth_ratio.size(0)) 55 | 56 | estimated_labels[depth_ratio > (threashold)] = 1 57 | estimated_labels[depth_ratio < (1/threashold)] = -1 58 | 59 | diff = estimated_labels - ground_truth 60 | diff[diff != 0] = 1 61 | 62 | # error 63 | inequal_error_count = diff[ground_truth != 0] 64 | inequal_error_count = torch.sum(inequal_error_count) 65 | 66 | error_count = torch.sum(diff) #diff[diff !=0] 67 | # error_count = error_count.size(0) 68 | 69 | equal_error_count = error_count - inequal_error_count 70 | 71 | 72 | # total 73 | total_count = depth_ratio.size(0) 74 | ground_truth[ground_truth !=0 ] = 1 75 | 76 | inequal_count_total = torch.sum(ground_truth) 77 | equal_total_count = total_count - inequal_count_total 78 | 79 | 80 | error_list = [equal_error_count, inequal_error_count, error_count] 81 | count_list = [equal_total_count, inequal_count_total, total_count] 82 | 83 | return error_list, count_list 84 | 85 | 86 | def computeSDR(self, prediction_d, targets): 87 | # for each image 88 | total_error = [0,0,0] 89 | total_samples = [0,0,0] 90 | 91 | for i in range(0, prediction_d.size(0)): 92 | 93 | if targets['has_SfM_feature'][i] == False: 94 | continue 95 | 96 | x_A_arr = targets["sdr_xA"][i].squeeze(0) 97 | x_B_arr = targets["sdr_xB"][i].squeeze(0) 98 | y_A_arr = targets["sdr_yA"][i].squeeze(0) 99 | y_B_arr = targets["sdr_yB"][i].squeeze(0) 100 | 101 | predict_depth = torch.exp(prediction_d[i,:,:]) 102 | predict_depth = predict_depth.squeeze(0) 103 | ground_truth = targets["sdr_gt"][i] 104 | 105 | # print(x_A_arr.size()) 106 | # print(y_A_arr.size()) 107 | 108 | z_A_arr = torch.gather( torch.index_select(predict_depth, 1 ,x_A_arr.cuda()) , 0, y_A_arr.view(1, -1).cuda())# predict_depth:index(2, x_A_arr):gather(1, y_A_arr:view(1, -1)) 109 | z_B_arr = torch.gather( torch.index_select(predict_depth, 1 ,x_B_arr.cuda()) , 0, y_B_arr.view(1, -1).cuda()) 110 | 111 | z_A_arr = z_A_arr.squeeze(0) 112 | z_B_arr = z_B_arr.squeeze(0) 113 | 114 | error_list, count_list = self.batch_classify(z_A_arr, z_B_arr,ground_truth) 115 | 116 | for j in range(0,3): 117 | total_error[j] += error_list[j] 118 | total_samples[j] += count_list[j] 119 | 120 | return total_error, total_samples 121 | 122 | 123 | def evaluate_SDR(self, input_, targets): 124 | input_images = Variable(input_.cuda() ) 125 | prediction_d = self.netG.forward(input_images) 126 | 127 | total_error, total_samples = self.computeSDR(prediction_d.data, targets) 128 | 129 | return total_error, total_samples 130 | 131 | def rmse_Loss(self, log_prediction_d, mask, log_gt): 132 | N = torch.sum(mask) 133 | log_d_diff = log_prediction_d - log_gt 134 | log_d_diff = torch.mul(log_d_diff, mask) 135 | s1 = torch.sum( torch.pow(log_d_diff,2) )/N 136 | 137 | s2 = torch.pow(torch.sum(log_d_diff),2)/(N*N) 138 | data_loss = s1 - s2 139 | 140 | data_loss = torch.sqrt(data_loss) 141 | 142 | return data_loss 143 | 144 | def evaluate_RMSE(self, input_images, prediction_d, targets): 145 | count = 0 146 | total_loss = Variable(torch.cuda.FloatTensor(1)) 147 | total_loss[0] = 0 148 | mask_0 = Variable(targets['mask_0'].cuda(), requires_grad = False) 149 | d_gt_0 = torch.log(Variable(targets['gt_0'].cuda(), requires_grad = False)) 150 | 151 | for i in range(0, mask_0.size(0)): 152 | 153 | total_loss += self.rmse_Loss(prediction_d[i,:,:], mask_0[i,:,:], d_gt_0[i,:,:]) 154 | count += 1 155 | 156 | return total_loss.data[0], count 157 | 158 | 159 | def evaluate_sc_inv(self, input_, targets): 160 | input_images = Variable(input_.cuda() ) 161 | prediction_d = self.netG.forward(input_images) 162 | rmse_loss , count= self.evaluate_RMSE(input_images, prediction_d, targets) 163 | 164 | return rmse_loss, count 165 | 166 | 167 | def switch_to_train(self): 168 | self.netG.train() 169 | 170 | def switch_to_eval(self): 171 | self.netG.eval() 172 | 173 | -------------------------------------------------------------------------------- /MegaDepth/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/baowenbo/DAIN/7c727aca56760f8f1a57ffb603c8c0b0fd1b1349/MegaDepth/models/__init__.py -------------------------------------------------------------------------------- /MegaDepth/models/base_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | 4 | class BaseModel(): 5 | def name(self): 6 | return 'BaseModel' 7 | 8 | def initialize(self, opt): 9 | self.opt = opt 10 | self.gpu_ids = opt.gpu_ids 11 | self.isTrain = opt.isTrain 12 | self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor 13 | self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) 14 | 15 | def set_input(self, input): 16 | self.input = input 17 | 18 | def forward(self): 19 | pass 20 | 21 | # used in test time, no backprop 22 | def test(self): 23 | pass 24 | 25 | def get_image_paths(self): 26 | pass 27 | 28 | def optimize_parameters(self): 29 | pass 30 | 31 | def get_current_visuals(self): 32 | return self.input 33 | 34 | def get_current_errors(self): 35 | return {} 36 | 37 | def save(self, label): 38 | pass 39 | 40 | # helper saving function that can be used by subclasses 41 | def save_network(self, network, network_label, epoch_label, gpu_ids): 42 | save_filename = '_%s_net_%s.pth' % (epoch_label, network_label) 43 | save_path = os.path.join(self.save_dir, save_filename) 44 | torch.save(network.cpu().state_dict(), save_path) 45 | if len(gpu_ids) and torch.cuda.is_available(): 46 | network.cuda(device_id=gpu_ids[0]) 47 | 48 | # helper loading function that can be used by subclasses 49 | def load_network(self, network, network_label, epoch_label): 50 | save_filename = '%s_net_%s.pth' % (epoch_label, network_label) 51 | save_path = os.path.join(self.save_dir, save_filename) 52 | print(save_path) 53 | model = torch.load(save_path) 54 | return model 55 | # network.load_state_dict(torch.load(save_path)) 56 | 57 | def update_learning_rate(): 58 | pass 59 | -------------------------------------------------------------------------------- /MegaDepth/models/models.py: -------------------------------------------------------------------------------- 1 | 2 | def create_model(opt,pretrained=None): 3 | model = None 4 | from .HG_model import HGModel 5 | model = HGModel(opt,pretrained) 6 | # print("model [%s] was created" % (model.name())) 7 | return model 8 | -------------------------------------------------------------------------------- /MegaDepth/options/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/baowenbo/DAIN/7c727aca56760f8f1a57ffb603c8c0b0fd1b1349/MegaDepth/options/__init__.py -------------------------------------------------------------------------------- /MegaDepth/options/base_options.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from ..util import util 4 | 5 | class BaseOptions(): 6 | def __init__(self): 7 | self.parser = argparse.ArgumentParser() 8 | self.initialized = False 9 | 10 | def initialize(self): 11 | # self.parser.add_argument('--dataroot', required=True, help='path to images (should have subfolders trainA, trainB, valA, valB, etc)') 12 | self.parser.add_argument('--batchSize', type=int, default=1, help='input batch size') 13 | self.parser.add_argument('--loadSize', type=int, default=286, help='scale images to this size') 14 | self.parser.add_argument('--fineSize', type=int, default=256, help='then crop to this size') 15 | self.parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels') 16 | self.parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels') 17 | self.parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer') 18 | self.parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer') 19 | # self.parser.add_argument('--which_model_netD', type=str, default='basic', help='selects model to use for netD') 20 | self.parser.add_argument('--which_model_netG', type=str, default='unet_256', help='selects model to use for netG') 21 | # self.parser.add_argument('--n_layers_D', type=int, default=3, help='only used if which_model_netD==n_layers') 22 | self.parser.add_argument('--gpu_ids', type=str, default='0,1', help='gpu ids: e.g. 0 0,1,2, 0,2') 23 | self.parser.add_argument('--name', type=str, default='test_local', help='name of the experiment. It decides where to store samples and models') 24 | # self.parser.add_argument('--align_data', action='store_true', 25 | # help='if True, the datasets are loaded from "test" and "train" directories and the data pairs are aligned') 26 | self.parser.add_argument('--model', type=str, default='pix2pix', 27 | help='chooses which model to use. cycle_gan, one_direction_test, pix2pix, ...') 28 | # self.parser.add_argument('--which_direction', type=str, default='AtoB', help='AtoB or BtoA') 29 | self.parser.add_argument('--nThreads', default=2, type=int, help='# threads for loading data') 30 | self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints/', help='models are saved here') 31 | self.parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization') 32 | self.parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly') 33 | self.parser.add_argument('--display_winsize', type=int, default=256, help='display window size') 34 | self.parser.add_argument('--display_id', type=int, default=1, help='window id of the web display') 35 | self.parser.add_argument('--identity', type=float, default=0.0, help='use identity mapping. Setting identity other than 1 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set optidentity = 0.1') 36 | self.parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator') 37 | self.parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.') 38 | 39 | self.initialized = True 40 | 41 | def parse(self): 42 | if not self.initialized: 43 | self.initialize() 44 | self.opt = self.parser.parse_known_args()[0] #parse_args() 45 | self.opt.isTrain = self.isTrain # train or test 46 | 47 | str_ids = self.opt.gpu_ids.split(',') 48 | self.opt.gpu_ids = [] 49 | for str_id in str_ids: 50 | id = int(str_id) 51 | if id >= 0: 52 | self.opt.gpu_ids.append(id) 53 | 54 | args = vars(self.opt) 55 | 56 | # print('------------ Options -------------') 57 | # for k, v in sorted(args.items()): 58 | # print('%s: %s' % (str(k), str(v))) 59 | # print('-------------- End ----------------') 60 | 61 | # save to the disk 62 | expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name) 63 | util.mkdirs(expr_dir) 64 | file_name = os.path.join(expr_dir, 'opt.txt') 65 | with open(file_name, 'wt') as opt_file: 66 | opt_file.write('------------ Options -------------\n') 67 | for k, v in sorted(args.items()): 68 | opt_file.write('%s: %s\n' % (str(k), str(v))) 69 | opt_file.write('-------------- End ----------------\n') 70 | return self.opt 71 | -------------------------------------------------------------------------------- /MegaDepth/options/test_options.py: -------------------------------------------------------------------------------- 1 | from .base_options import BaseOptions 2 | 3 | class TestOptions(BaseOptions): 4 | def initialize(self): 5 | BaseOptions.initialize(self) 6 | self.parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.') 7 | self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.') 8 | self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images') 9 | self.parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc') 10 | self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') 11 | self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run') 12 | self.isTrain = False 13 | -------------------------------------------------------------------------------- /MegaDepth/options/train_options.py: -------------------------------------------------------------------------------- 1 | from .base_options import BaseOptions 2 | 3 | class TrainOptions(BaseOptions): 4 | def initialize(self): 5 | BaseOptions.initialize(self) 6 | self.parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen') 7 | self.parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console') 8 | self.parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results') 9 | self.parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs') 10 | self.parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model') 11 | self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc') 12 | self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') 13 | self.parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate') 14 | self.parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero') 15 | self.parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam') 16 | self.parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam') 17 | self.parser.add_argument('--no_lsgan', action='store_true', help='do *not* use least square GAN, if false, use vanilla GAN') 18 | self.parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)') 19 | self.parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)') 20 | self.parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images') 21 | self.parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/') 22 | self.parser.add_argument('--no_flip' , action='store_true', help='if specified, do not flip the images for data argumentation') 23 | 24 | # NOT-IMPLEMENTED self.parser.add_argument('--preprocessing', type=str, default='resize_and_crop', help='resizing/cropping strategy') 25 | self.isTrain = True 26 | -------------------------------------------------------------------------------- /MegaDepth/rmse_error_main.py: -------------------------------------------------------------------------------- 1 | import time 2 | import torch 3 | import sys 4 | 5 | from options.train_options import TrainOptions 6 | opt = TrainOptions().parse() # set CUDA_VISIBLE_DEVICES before import torch 7 | from data.data_loader import CreateDataLoader 8 | from models.models import create_model 9 | 10 | dataset_root = "/phoenix/S6/zl548/" 11 | test_list_dir_l = '/phoenix/S6/zl548/MegaDpeth_code/test_list/landscape/' 12 | input_height = 240 13 | input_width = 320 14 | is_flipped = False 15 | shuffle = False 16 | 17 | test_data_loader_l = CreateDataLoader(dataset_root, test_list_dir_l, input_height, input_width, is_flipped, shuffle) 18 | test_dataset_l = test_data_loader_l.load_data() 19 | test_dataset_size_l = len(test_data_loader_l) 20 | print('========================= test images = %d' % test_dataset_size_l) 21 | test_list_dir_p = '/phoenix/S6/zl548/MegaDpeth_code/test_list/portrait/' 22 | input_height = 320 23 | input_width = 240 24 | test_data_loader_p = CreateDataLoader(dataset_root, test_list_dir_p, input_height, input_width, is_flipped, shuffle) 25 | test_dataset_p = test_data_loader_p.load_data() 26 | test_dataset_size_p = len(test_data_loader_p) 27 | print('========================= test images = %d' % test_dataset_size_p) 28 | 29 | 30 | model = create_model(opt) 31 | 32 | 33 | def test(model): 34 | total_loss =0 35 | toal_count = 0 36 | print("============================= TEST ============================") 37 | model.switch_to_eval() 38 | for i, data in enumerate(test_dataset_l): 39 | stacked_img = data['img_1'] 40 | targets = data['target_1'] 41 | 42 | rmse_loss , count = model.evaluate_sc_inv(stacked_img, targets) 43 | 44 | total_loss += rmse_loss 45 | toal_count += count 46 | 47 | print('RMSE loss is', total_loss/float(toal_count)) 48 | 49 | for i, data in enumerate(test_dataset_p): 50 | stacked_img = data['img_1'] 51 | targets = data['target_1'] 52 | rmse_loss , count = model.evaluate_sc_inv(stacked_img, targets) 53 | 54 | total_loss += rmse_loss 55 | toal_count += count 56 | 57 | print('RMSE loss is', total_loss/float(toal_count)) 58 | 59 | 60 | print('average RMSE loss is', total_loss/float(toal_count)) 61 | 62 | print("WE ARE IN TESTING RMSE!!!!") 63 | test(model) 64 | print("WE ARE DONE TESTING!!!") 65 | 66 | 67 | print("We are done") 68 | -------------------------------------------------------------------------------- /MegaDepth/util/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/baowenbo/DAIN/7c727aca56760f8f1a57ffb603c8c0b0fd1b1349/MegaDepth/util/__init__.py -------------------------------------------------------------------------------- /MegaDepth/util/html.py: -------------------------------------------------------------------------------- 1 | import dominate 2 | from dominate.tags import * 3 | import os 4 | 5 | 6 | class HTML: 7 | def __init__(self, web_dir, title, reflesh=0): 8 | self.title = title 9 | self.web_dir = web_dir 10 | self.img_dir = os.path.join(self.web_dir, 'images') 11 | if not os.path.exists(self.web_dir): 12 | os.makedirs(self.web_dir) 13 | if not os.path.exists(self.img_dir): 14 | os.makedirs(self.img_dir) 15 | # print(self.img_dir) 16 | 17 | self.doc = dominate.document(title=title) 18 | if reflesh > 0: 19 | with self.doc.head: 20 | meta(http_equiv="reflesh", content=str(reflesh)) 21 | 22 | def get_image_dir(self): 23 | return self.img_dir 24 | 25 | def add_header(self, str): 26 | with self.doc: 27 | h3(str) 28 | 29 | def add_table(self, border=1): 30 | self.t = table(border=border, style="table-layout: fixed;") 31 | self.doc.add(self.t) 32 | 33 | def add_images(self, ims, txts, links, width=400): 34 | self.add_table() 35 | with self.t: 36 | with tr(): 37 | for im, txt, link in zip(ims, txts, links): 38 | with td(style="word-wrap: break-word;", halign="center", valign="top"): 39 | with p(): 40 | with a(href=os.path.join('images', link)): 41 | img(style="width:%dpx" % width, src=os.path.join('images', im)) 42 | br() 43 | p(txt) 44 | 45 | def save(self): 46 | html_file = '%s/index.html' % self.web_dir 47 | f = open(html_file, 'wt') 48 | f.write(self.doc.render()) 49 | f.close() 50 | 51 | 52 | if __name__ == '__main__': 53 | html = HTML('web/', 'test_html') 54 | html.add_header('hello world') 55 | 56 | ims = [] 57 | txts = [] 58 | links = [] 59 | for n in range(4): 60 | ims.append('image_%d.png' % n) 61 | txts.append('text_%d' % n) 62 | links.append('image_%d.png' % n) 63 | html.add_images(ims, txts, links) 64 | html.save() 65 | -------------------------------------------------------------------------------- /MegaDepth/util/image_pool.py: -------------------------------------------------------------------------------- 1 | import random 2 | import numpy as np 3 | import torch 4 | from pdb import set_trace as st 5 | from torch.autograd import Variable 6 | class ImagePool(): 7 | def __init__(self, pool_size): 8 | self.pool_size = pool_size 9 | if self.pool_size > 0: 10 | self.num_imgs = 0 11 | self.images = [] 12 | 13 | def query(self, images): 14 | if self.pool_size == 0: 15 | return images 16 | return_images = [] 17 | for image in images.data: 18 | image = torch.unsqueeze(image, 0) 19 | if self.num_imgs < self.pool_size: 20 | self.num_imgs = self.num_imgs + 1 21 | self.images.append(image) 22 | return_images.append(image) 23 | else: 24 | p = random.uniform(0, 1) 25 | if p > 0.5: 26 | random_id = random.randint(0, self.pool_size-1) 27 | tmp = self.images[random_id].clone() 28 | self.images[random_id] = image 29 | return_images.append(tmp) 30 | else: 31 | return_images.append(image) 32 | return_images = Variable(torch.cat(return_images, 0)) 33 | return return_images 34 | -------------------------------------------------------------------------------- /MegaDepth/util/png.py: -------------------------------------------------------------------------------- 1 | import struct 2 | import zlib 3 | 4 | def encode(buf, width, height): 5 | """ buf: must be bytes or a bytearray in py3, a regular string in py2. formatted RGBRGB... """ 6 | assert (width * height * 3 == len(buf)) 7 | bpp = 3 8 | 9 | def raw_data(): 10 | # reverse the vertical line order and add null bytes at the start 11 | row_bytes = width * bpp 12 | for row_start in range((height - 1) * width * bpp, -1, -row_bytes): 13 | yield b'\x00' 14 | yield buf[row_start:row_start + row_bytes] 15 | 16 | def chunk(tag, data): 17 | return [ 18 | struct.pack("!I", len(data)), 19 | tag, 20 | data, 21 | struct.pack("!I", 0xFFFFFFFF & zlib.crc32(data, zlib.crc32(tag))) 22 | ] 23 | 24 | SIGNATURE = b'\x89PNG\r\n\x1a\n' 25 | COLOR_TYPE_RGB = 2 26 | COLOR_TYPE_RGBA = 6 27 | bit_depth = 8 28 | return b''.join( 29 | [ SIGNATURE ] + 30 | chunk(b'IHDR', struct.pack("!2I5B", width, height, bit_depth, COLOR_TYPE_RGB, 0, 0, 0)) + 31 | chunk(b'IDAT', zlib.compress(b''.join(raw_data()), 9)) + 32 | chunk(b'IEND', b'') 33 | ) 34 | -------------------------------------------------------------------------------- /MegaDepth/util/util.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | import numpy as np 4 | from PIL import Image 5 | import inspect, re 6 | import numpy as np 7 | import os 8 | import collections 9 | 10 | # Converts a Tensor into a Numpy array 11 | # |imtype|: the desired type of the converted numpy array 12 | def tensor2im(image_tensor, imtype=np.uint8): 13 | image_numpy = image_tensor[0].cpu().float().numpy() 14 | image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 15 | return image_numpy.astype(imtype) 16 | 17 | 18 | def diagnose_network(net, name='network'): 19 | mean = 0.0 20 | count = 0 21 | for param in net.parameters(): 22 | if param.grad is not None: 23 | mean += torch.mean(torch.abs(param.grad.data)) 24 | count += 1 25 | if count > 0: 26 | mean = mean / count 27 | print(name) 28 | print(mean) 29 | 30 | 31 | def save_image(image_numpy, image_path): 32 | image_pil = Image.fromarray(image_numpy) 33 | image_pil.save(image_path) 34 | 35 | def info(object, spacing=10, collapse=1): 36 | """Print methods and doc strings. 37 | Takes module, class, list, dictionary, or string.""" 38 | methodList = [e for e in dir(object) if isinstance(getattr(object, e), collections.Callable)] 39 | processFunc = collapse and (lambda s: " ".join(s.split())) or (lambda s: s) 40 | print( "\n".join(["%s %s" % 41 | (method.ljust(spacing), 42 | processFunc(str(getattr(object, method).__doc__))) 43 | for method in methodList]) ) 44 | 45 | def varname(p): 46 | for line in inspect.getframeinfo(inspect.currentframe().f_back)[3]: 47 | m = re.search(r'\bvarname\s*\(\s*([A-Za-z_][A-Za-z0-9_]*)\s*\)', line) 48 | if m: 49 | return m.group(1) 50 | 51 | def print_numpy(x, val=True, shp=False): 52 | x = x.astype(np.float64) 53 | if shp: 54 | print('shape,', x.shape) 55 | if val: 56 | x = x.flatten() 57 | print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % ( 58 | np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x))) 59 | 60 | 61 | def mkdirs(paths): 62 | if isinstance(paths, list) and not isinstance(paths, str): 63 | for path in paths: 64 | mkdir(path) 65 | else: 66 | mkdir(paths) 67 | 68 | 69 | def mkdir(path): 70 | if not os.path.exists(path): 71 | os.makedirs(path) 72 | -------------------------------------------------------------------------------- /MegaDepth/util/visualizer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import ntpath 4 | import time 5 | from . import util 6 | from . import html 7 | 8 | class Visualizer(): 9 | def __init__(self, opt): 10 | # self.opt = opt 11 | self.display_id = opt.display_id 12 | self.use_html = opt.isTrain and not opt.no_html 13 | self.win_size = opt.display_winsize 14 | self.name = opt.name 15 | if self.display_id > 0: 16 | import visdom 17 | self.vis = visdom.Visdom() 18 | 19 | if self.use_html: 20 | self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web') 21 | self.img_dir = os.path.join(self.web_dir, 'images') 22 | print('create web directory %s...' % self.web_dir) 23 | util.mkdirs([self.web_dir, self.img_dir]) 24 | 25 | 26 | # |visuals|: dictionary of images to display or save 27 | def display_current_results(self, visuals, epoch): 28 | if self.display_id > 0: # show images in the browser 29 | idx = 1 30 | for label, image_numpy in visuals.items(): 31 | #image_numpy = np.flipud(image_numpy) 32 | self.vis.image(image_numpy.transpose([2,0,1]), opts=dict(title=label), 33 | win=self.display_id + idx) 34 | idx += 1 35 | 36 | if self.use_html: # save images to a html file 37 | for label, image_numpy in visuals.items(): 38 | img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label)) 39 | util.save_image(image_numpy, img_path) 40 | # update website 41 | webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1) 42 | for n in range(epoch, 0, -1): 43 | webpage.add_header('epoch [%d]' % n) 44 | ims = [] 45 | txts = [] 46 | links = [] 47 | 48 | for label, image_numpy in visuals.items(): 49 | img_path = 'epoch%.3d_%s.png' % (n, label) 50 | ims.append(img_path) 51 | txts.append(label) 52 | links.append(img_path) 53 | webpage.add_images(ims, txts, links, width=self.win_size) 54 | webpage.save() 55 | 56 | # errors: dictionary of error labels and values 57 | def plot_current_errors(self, epoch, counter_ratio, opt, errors): 58 | if not hasattr(self, 'plot_data'): 59 | self.plot_data = {'X':[],'Y':[], 'legend':list(errors.keys())} 60 | self.plot_data['X'].append(epoch + counter_ratio) 61 | self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']]) 62 | self.vis.line( 63 | X=np.stack([np.array(self.plot_data['X'])]*len(self.plot_data['legend']),1), 64 | Y=np.array(self.plot_data['Y']), 65 | opts={ 66 | 'title': self.name + ' loss over time', 67 | 'legend': self.plot_data['legend'], 68 | 'xlabel': 'epoch', 69 | 'ylabel': 'loss'}, 70 | win=self.display_id) 71 | 72 | # errors: same format as |errors| of plotCurrentErrors 73 | def print_current_errors(self, epoch, i, errors, t): 74 | message = '(epoch: %d, iters: %d, time: %.3f) ' % (epoch, i, t) 75 | for k, v in errors.items(): 76 | message += '%s: %.3f ' % (k, v) 77 | 78 | print(message) 79 | 80 | # save image to the disk 81 | def save_images(self, webpage, visuals, image_path): 82 | image_dir = webpage.get_image_dir() 83 | short_path = ntpath.basename(image_path[0]) 84 | name = os.path.splitext(short_path)[0] 85 | 86 | webpage.add_header(name) 87 | ims = [] 88 | txts = [] 89 | links = [] 90 | 91 | for label, image_numpy in visuals.items(): 92 | image_name = '%s_%s.png' % (name, label) 93 | save_path = os.path.join(image_dir, image_name) 94 | util.save_image(image_numpy, save_path) 95 | 96 | ims.append(image_name) 97 | txts.append(label) 98 | links.append(image_name) 99 | webpage.add_images(ims, txts, links, width=self.win_size) 100 | -------------------------------------------------------------------------------- /PWCNet/__init__.py: -------------------------------------------------------------------------------- 1 | from .PWCNet import * -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/baowenbo/DAIN/7c727aca56760f8f1a57ffb603c8c0b0fd1b1349/PWCNet/correlation_package_pytorch1_0/__init__.py -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | echo "Need pytorch>=1.0.0" 4 | source activate pytorch1.0.0 5 | 6 | export PYTHONPATH=$PYTHONPATH:$(pwd)/../../my_package 7 | 8 | rm -rf build *.egg-info dist 9 | python setup.py install 10 | -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/clean.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | echo "Need pytorch>=1.0.0" 4 | source activate pytorch1.0.0 5 | 6 | 7 | rm -rf build *.egg-info dist 8 | #python setup.py install 9 | -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/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 | -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/correlation_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | #include "correlation_cuda_kernel.cuh" 7 | 8 | int correlation_forward_cuda(at::Tensor& input1, at::Tensor& input2, at::Tensor& rInput1, at::Tensor& rInput2, at::Tensor& output, 9 | int pad_size, 10 | int kernel_size, 11 | int max_displacement, 12 | int stride1, 13 | int stride2, 14 | int corr_type_multiply) 15 | { 16 | 17 | int batchSize = input1.size(0); 18 | 19 | int nInputChannels = input1.size(1); 20 | int inputHeight = input1.size(2); 21 | int inputWidth = input1.size(3); 22 | 23 | int kernel_radius = (kernel_size - 1) / 2; 24 | int border_radius = kernel_radius + max_displacement; 25 | 26 | int paddedInputHeight = inputHeight + 2 * pad_size; 27 | int paddedInputWidth = inputWidth + 2 * pad_size; 28 | 29 | int nOutputChannels = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1); 30 | 31 | int outputHeight = ceil(static_cast(paddedInputHeight - 2 * border_radius) / static_cast(stride1)); 32 | int outputwidth = ceil(static_cast(paddedInputWidth - 2 * border_radius) / static_cast(stride1)); 33 | 34 | rInput1.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 35 | rInput2.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 36 | output.resize_({batchSize, nOutputChannels, outputHeight, outputwidth}); 37 | 38 | rInput1.fill_(0); 39 | rInput2.fill_(0); 40 | output.fill_(0); 41 | 42 | int success = correlation_forward_cuda_kernel( 43 | output, 44 | output.size(0), 45 | output.size(1), 46 | output.size(2), 47 | output.size(3), 48 | output.stride(0), 49 | output.stride(1), 50 | output.stride(2), 51 | output.stride(3), 52 | input1, 53 | input1.size(1), 54 | input1.size(2), 55 | input1.size(3), 56 | input1.stride(0), 57 | input1.stride(1), 58 | input1.stride(2), 59 | input1.stride(3), 60 | input2, 61 | input2.size(1), 62 | input2.stride(0), 63 | input2.stride(1), 64 | input2.stride(2), 65 | input2.stride(3), 66 | rInput1, 67 | rInput2, 68 | pad_size, 69 | kernel_size, 70 | max_displacement, 71 | stride1, 72 | stride2, 73 | corr_type_multiply, 74 | // at::globalContext().getCurrentCUDAStream() //works for 0.4.1 75 | at::cuda::getCurrentCUDAStream() //works for 1.0.0 76 | ); 77 | 78 | //check for errors 79 | if (!success) { 80 | AT_ERROR("CUDA call failed"); 81 | } 82 | 83 | return 1; 84 | 85 | } 86 | 87 | int correlation_backward_cuda(at::Tensor& input1, at::Tensor& input2, at::Tensor& rInput1, at::Tensor& rInput2, at::Tensor& gradOutput, 88 | at::Tensor& gradInput1, at::Tensor& gradInput2, 89 | int pad_size, 90 | int kernel_size, 91 | int max_displacement, 92 | int stride1, 93 | int stride2, 94 | int corr_type_multiply) 95 | { 96 | 97 | int batchSize = input1.size(0); 98 | int nInputChannels = input1.size(1); 99 | int paddedInputHeight = input1.size(2)+ 2 * pad_size; 100 | int paddedInputWidth = input1.size(3)+ 2 * pad_size; 101 | 102 | int height = input1.size(2); 103 | int width = input1.size(3); 104 | 105 | rInput1.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 106 | rInput2.resize_({batchSize, paddedInputHeight, paddedInputWidth, nInputChannels}); 107 | gradInput1.resize_({batchSize, nInputChannels, height, width}); 108 | gradInput2.resize_({batchSize, nInputChannels, height, width}); 109 | 110 | rInput1.fill_(0); 111 | rInput2.fill_(0); 112 | gradInput1.fill_(0); 113 | gradInput2.fill_(0); 114 | 115 | int success = correlation_backward_cuda_kernel(gradOutput, 116 | gradOutput.size(0), 117 | gradOutput.size(1), 118 | gradOutput.size(2), 119 | gradOutput.size(3), 120 | gradOutput.stride(0), 121 | gradOutput.stride(1), 122 | gradOutput.stride(2), 123 | gradOutput.stride(3), 124 | input1, 125 | input1.size(1), 126 | input1.size(2), 127 | input1.size(3), 128 | input1.stride(0), 129 | input1.stride(1), 130 | input1.stride(2), 131 | input1.stride(3), 132 | input2, 133 | input2.stride(0), 134 | input2.stride(1), 135 | input2.stride(2), 136 | input2.stride(3), 137 | gradInput1, 138 | gradInput1.stride(0), 139 | gradInput1.stride(1), 140 | gradInput1.stride(2), 141 | gradInput1.stride(3), 142 | gradInput2, 143 | gradInput2.size(1), 144 | gradInput2.stride(0), 145 | gradInput2.stride(1), 146 | gradInput2.stride(2), 147 | gradInput2.stride(3), 148 | rInput1, 149 | rInput2, 150 | pad_size, 151 | kernel_size, 152 | max_displacement, 153 | stride1, 154 | stride2, 155 | corr_type_multiply, 156 | // at::globalContext().getCurrentCUDAStream() //works for 0.4.1 157 | at::cuda::getCurrentCUDAStream() //works for 1.0.0 158 | ); 159 | 160 | if (!success) { 161 | AT_ERROR("CUDA call failed"); 162 | } 163 | 164 | return 1; 165 | } 166 | 167 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 168 | m.def("forward", &correlation_forward_cuda, "Correlation forward (CUDA)"); 169 | m.def("backward", &correlation_backward_cuda, "Correlation backward (CUDA)"); 170 | } 171 | 172 | -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/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 | -------------------------------------------------------------------------------- /PWCNet/correlation_package_pytorch1_0/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='correlation_cuda', 12 | ext_modules=[ 13 | CUDAExtension('correlation_cuda', [ 14 | 'correlation_cuda.cc', 15 | 'correlation_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /PWCNet/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .PWCNet import * 2 | -------------------------------------------------------------------------------- /Resblock/BasicBlock.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | import torch.nn.init as weight_init 5 | import torch 6 | __all__ = ['MultipleBasicBlock','MultipleBasicBlock_4'] 7 | def conv3x3(in_planes, out_planes, dilation = 1, stride=1): 8 | "3x3 convolution with padding" 9 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 10 | padding=int(dilation*(3-1)/2), dilation=dilation, bias=False) 11 | class BasicBlock(nn.Module): 12 | expansion = 1 13 | 14 | def __init__(self, inplanes, planes, dilation = 1, stride=1, downsample=None): 15 | super(BasicBlock, self).__init__() 16 | self.conv1 = conv3x3(inplanes, planes,dilation, stride) 17 | # self.bn1 = nn.BatchNorm2d(planes) 18 | self.relu = nn.ReLU(inplace=True) 19 | self.conv2 = conv3x3(planes, planes) 20 | # self.bn2 = nn.BatchNorm2d(planes) 21 | self.downsample = downsample 22 | self.stride = stride 23 | 24 | for m in self.modules(): 25 | if isinstance(m, nn.Conv2d): 26 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 27 | m.weight.data.normal_(0, math.sqrt(2. / n)) 28 | # weight_init.xavier_normal() 29 | elif isinstance(m, nn.BatchNorm2d): 30 | m.weight.data.fill_(1) 31 | m.bias.data.zero_() 32 | 33 | def forward(self, x): 34 | residual = x 35 | 36 | out = self.conv1(x) 37 | # out = self.bn1(out) 38 | out = self.relu(out) 39 | 40 | out = self.conv2(out) 41 | # out = self.bn2(out) 42 | 43 | if self.downsample is not None: 44 | residual = self.downsample(x) 45 | 46 | out += residual 47 | out = self.relu(out) 48 | 49 | return out 50 | class MultipleBasicBlock(nn.Module): 51 | 52 | def __init__(self,input_feature, 53 | block, num_blocks, 54 | intermediate_feature = 64, dense = True): 55 | super(MultipleBasicBlock, self).__init__() 56 | self.dense = dense 57 | self.num_block = num_blocks 58 | self.intermediate_feature = intermediate_feature 59 | 60 | self.block1= nn.Sequential(*[ 61 | nn.Conv2d(input_feature, intermediate_feature, 62 | kernel_size=7, stride=1, padding=3, bias=True), 63 | nn.ReLU(inplace=True) 64 | ]) 65 | 66 | # for i in range(1, num_blocks): 67 | self.block2 = block(intermediate_feature, intermediate_feature, dilation = 1) if num_blocks>=2 else None 68 | self.block3 = block(intermediate_feature, intermediate_feature, dilation = 1) if num_blocks>=3 else None 69 | self.block4 = block(intermediate_feature, intermediate_feature, dilation = 1) if num_blocks>=4 else None 70 | self.block5 = nn.Sequential(*[nn.Conv2d(intermediate_feature, 3 , (3, 3), 1, (1, 1))]) 71 | 72 | for m in self.modules(): 73 | if isinstance(m, nn.Conv2d): 74 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 75 | m.weight.data.normal_(0, math.sqrt(2. / n)) 76 | elif isinstance(m, nn.BatchNorm2d): 77 | m.weight.data.fill_(1) 78 | m.bias.data.zero_() 79 | 80 | def forward(self, x): 81 | x = self.block1(x) 82 | x = self.block2(x) if self.num_block>=2 else x 83 | x = self.block3(x) if self.num_block>=3 else x 84 | x = self.block4(x) if self.num_block== 4 else x 85 | x = self.block5(x) 86 | return x 87 | 88 | def MultipleBasicBlock_4(input_feature,intermediate_feature = 64): 89 | model = MultipleBasicBlock(input_feature, 90 | BasicBlock,4 , 91 | intermediate_feature) 92 | return model 93 | 94 | 95 | if __name__ == '__main__': 96 | 97 | # x= Variable(torch.randn(2,3,224,448)) 98 | # model = S2DF(BasicBlock,3,True) 99 | # y = model(x) 100 | model = MultipleBasicBlock(200, BasicBlock,4) 101 | model = BasicBlock(64,64,1) 102 | # y = model(x) 103 | exit(0) -------------------------------------------------------------------------------- /Resblock/__init__.py: -------------------------------------------------------------------------------- 1 | from .BasicBlock import * -------------------------------------------------------------------------------- /S2D_models/__init__.py: -------------------------------------------------------------------------------- 1 | from .S2DF import * -------------------------------------------------------------------------------- /Stack.py: -------------------------------------------------------------------------------- 1 | 2 | class Stack: 3 | def __init__(self): 4 | self.stack = [] 5 | def pop(self): 6 | if self.is_empty(): 7 | return None 8 | else: 9 | return self.stack.pop() 10 | def push(self,val): 11 | return self.stack.append(val) 12 | def peak(self): 13 | if self.is_empty(): 14 | return None 15 | else: 16 | return self.stack[-1] 17 | def size(self): 18 | return len(self.stack) 19 | def is_empty(self): 20 | return self.size() == 0 -------------------------------------------------------------------------------- /balancedsampler.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data.sampler import Sampler 2 | import torch 3 | 4 | class RandomBalancedSampler(Sampler): 5 | """Samples elements randomly, with an arbitrary size, independant from dataset length. 6 | this is a balanced sampling that will sample the whole dataset with a random permutation. 7 | 8 | Arguments: 9 | data_source (Dataset): dataset to sample from 10 | """ 11 | 12 | def __init__(self, data_source, epoch_size): 13 | self.data_size = len(data_source) 14 | self.epoch_size = epoch_size 15 | self.index = 0 16 | 17 | def __next__(self): 18 | if self.index == 0: 19 | #re-shuffle the sampler 20 | self.indices = torch.randperm(self.data_size) 21 | self.index = (self.index+1)%self.data_size 22 | return self.indices[self.index] 23 | 24 | def next(self): 25 | return self.__next__() 26 | 27 | def __iter__(self): 28 | return self 29 | 30 | def __len__(self): 31 | return min(self.data_size,self.epoch_size) if self.epoch_size>0 else self.data_size 32 | 33 | class SequentialBalancedSampler(Sampler): 34 | """Samples elements dequentially, with an arbitrary size, independant from dataset length. 35 | this is a balanced sampling that will sample the whole dataset before resetting it. 36 | 37 | Arguments: 38 | data_source (Dataset): dataset to sample from 39 | """ 40 | 41 | def __init__(self, data_source, epoch_size): 42 | self.data_size = len(data_source) 43 | self.epoch_size = epoch_size 44 | self.index = 0 45 | 46 | def __next__(self): 47 | self.index = (self.index+1)%self.data_size 48 | return self.index 49 | 50 | def next(self): 51 | return self.__next__() 52 | 53 | def __iter__(self): 54 | return self 55 | 56 | def __len__(self): 57 | return min(self.data_size,self.epoch_size) if self.epoch_size>0 else self.data_size 58 | -------------------------------------------------------------------------------- /colab_interpolate.py: -------------------------------------------------------------------------------- 1 | import time 2 | import os 3 | from torch.autograd import Variable 4 | import torch 5 | import numpy as np 6 | import numpy 7 | import networks 8 | from my_args import args 9 | from imageio import imread, imsave 10 | from AverageMeter import * 11 | import shutil 12 | import datetime 13 | torch.backends.cudnn.benchmark = True 14 | 15 | model = networks.__dict__[args.netName]( 16 | channel = args.channels, 17 | filter_size = args.filter_size, 18 | timestep = args.time_step, 19 | training = False) 20 | 21 | if args.use_cuda: 22 | model = model.cuda() 23 | 24 | model_path = './model_weights/best.pth' 25 | if not os.path.exists(model_path): 26 | print("*****************************************************************") 27 | print("**** We couldn't load any trained weights ***********************") 28 | print("*****************************************************************") 29 | exit(1) 30 | 31 | if args.use_cuda: 32 | pretrained_dict = torch.load(model_path) 33 | else: 34 | pretrained_dict = torch.load(model_path, map_location=lambda storage, loc: storage) 35 | 36 | model_dict = model.state_dict() 37 | # 1. filter out unnecessary keys 38 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 39 | # 2. overwrite entries in the existing state dict 40 | model_dict.update(pretrained_dict) 41 | # 3. load the new state dict 42 | model.load_state_dict(model_dict) 43 | # 4. release the pretrained dict for saving memory 44 | pretrained_dict = [] 45 | 46 | model = model.eval() # deploy mode 47 | 48 | frames_dir = args.frame_input_dir 49 | output_dir = args.frame_output_dir 50 | 51 | timestep = args.time_step 52 | time_offsets = [kk * timestep for kk in range(1, int(1.0 / timestep))] 53 | 54 | input_frame = args.start_frame - 1 55 | loop_timer = AverageMeter() 56 | 57 | final_frame = args.end_frame 58 | 59 | torch.set_grad_enabled(False) 60 | 61 | # we want to have input_frame between (start_frame-1) and (end_frame-2) 62 | # this is because at each step we read (frame) and (frame+1) 63 | # so the last iteration will actuall be (end_frame-1) and (end_frame) 64 | while input_frame < final_frame - 1: 65 | input_frame += 1 66 | 67 | start_time = time.time() 68 | 69 | filename_frame_1 = os.path.join(frames_dir, f'{input_frame:0>5d}.png') 70 | filename_frame_2 = os.path.join(frames_dir, f'{input_frame+1:0>5d}.png') 71 | 72 | X0 = torch.from_numpy(np.transpose(imread(filename_frame_1), (2,0,1)).astype("float32") / 255.0).type(args.dtype) 73 | X1 = torch.from_numpy(np.transpose(imread(filename_frame_2), (2,0,1)).astype("float32") / 255.0).type(args.dtype) 74 | 75 | assert (X0.size(1) == X1.size(1)) 76 | assert (X0.size(2) == X1.size(2)) 77 | 78 | intWidth = X0.size(2) 79 | intHeight = X0.size(1) 80 | channels = X0.size(0) 81 | if not channels == 3: 82 | print(f"Skipping {filename_frame_1}-{filename_frame_2} -- expected 3 color channels but found {channels}.") 83 | continue 84 | 85 | if intWidth != ((intWidth >> 7) << 7): 86 | intWidth_pad = (((intWidth >> 7) + 1) << 7) # more than necessary 87 | intPaddingLeft = int((intWidth_pad - intWidth) / 2) 88 | intPaddingRight = intWidth_pad - intWidth - intPaddingLeft 89 | else: 90 | intPaddingLeft = 32 91 | intPaddingRight= 32 92 | 93 | if intHeight != ((intHeight >> 7) << 7): 94 | intHeight_pad = (((intHeight >> 7) + 1) << 7) # more than necessary 95 | intPaddingTop = int((intHeight_pad - intHeight) / 2) 96 | intPaddingBottom = intHeight_pad - intHeight - intPaddingTop 97 | else: 98 | intPaddingTop = 32 99 | intPaddingBottom = 32 100 | 101 | pader = torch.nn.ReplicationPad2d([intPaddingLeft, intPaddingRight, intPaddingTop, intPaddingBottom]) 102 | 103 | X0 = Variable(torch.unsqueeze(X0,0)) 104 | X1 = Variable(torch.unsqueeze(X1,0)) 105 | X0 = pader(X0) 106 | X1 = pader(X1) 107 | 108 | if args.use_cuda: 109 | X0 = X0.cuda() 110 | X1 = X1.cuda() 111 | 112 | y_s, offset, filter = model(torch.stack((X0, X1),dim = 0)) 113 | y_ = y_s[args.save_which] 114 | 115 | if args.use_cuda: 116 | X0 = X0.data.cpu().numpy() 117 | if not isinstance(y_, list): 118 | y_ = y_.data.cpu().numpy() 119 | else: 120 | y_ = [item.data.cpu().numpy() for item in y_] 121 | offset = [offset_i.data.cpu().numpy() for offset_i in offset] 122 | filter = [filter_i.data.cpu().numpy() for filter_i in filter] if filter[0] is not None else None 123 | X1 = X1.data.cpu().numpy() 124 | else: 125 | X0 = X0.data.numpy() 126 | if not isinstance(y_, list): 127 | y_ = y_.data.numpy() 128 | else: 129 | y_ = [item.data.numpy() for item in y_] 130 | offset = [offset_i.data.numpy() for offset_i in offset] 131 | filter = [filter_i.data.numpy() for filter_i in filter] 132 | X1 = X1.data.numpy() 133 | 134 | X0 = np.transpose(255.0 * X0.clip(0,1.0)[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) 135 | y_ = [np.transpose(255.0 * item.clip(0,1.0)[0, :, intPaddingTop:intPaddingTop+intHeight, 136 | intPaddingLeft:intPaddingLeft+intWidth], (1, 2, 0)) for item in y_] 137 | offset = [np.transpose(offset_i[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) for offset_i in offset] 138 | filter = [np.transpose( 139 | filter_i[0, :, intPaddingTop:intPaddingTop + intHeight, intPaddingLeft: intPaddingLeft + intWidth], 140 | (1, 2, 0)) for filter_i in filter] if filter is not None else None 141 | X1 = np.transpose(255.0 * X1.clip(0,1.0)[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) 142 | 143 | interpolated_frame_number = 0 144 | shutil.copy(filename_frame_1, os.path.join(output_dir, f"{input_frame:0>5d}{interpolated_frame_number:0>3d}.png")) 145 | for item, time_offset in zip(y_, time_offsets): 146 | interpolated_frame_number += 1 147 | output_frame_file_path = os.path.join(output_dir, f"{input_frame:0>5d}{interpolated_frame_number:0>3d}.png") 148 | imsave(output_frame_file_path, np.round(item).astype(numpy.uint8)) 149 | 150 | end_time = time.time() 151 | loop_timer.update(end_time - start_time) 152 | 153 | frames_left = final_frame - input_frame 154 | estimated_seconds_left = frames_left * loop_timer.avg 155 | estimated_time_left = datetime.timedelta(seconds=estimated_seconds_left) 156 | print(f"****** Processed frame {input_frame} | Time per frame (avg): {loop_timer.avg:2.2f}s | Time left: {estimated_time_left} ******************" ) 157 | 158 | # Copying last frame 159 | last_frame_filename = os.path.join(frames_dir, str(str(final_frame).zfill(5))+'.png') 160 | shutil.copy(last_frame_filename, os.path.join(output_dir, f"{final_frame:0>5d}{0:0>3d}.png")) 161 | 162 | print("Finished processing images.") 163 | -------------------------------------------------------------------------------- /datasets/Vimeo_90K_interp.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | import random 3 | # import glob 4 | import math 5 | from .listdatasets import ListDataset,Vimeo_90K_loader 6 | 7 | 8 | def make_dataset(root, list_file): 9 | raw_im_list = open(os.path.join(root, list_file)).read().splitlines() 10 | # the last line is invalid in test set. 11 | # print("The last sample is : " + raw_im_list[-1]) 12 | raw_im_list = raw_im_list[:-1] 13 | assert len(raw_im_list) > 0 14 | random.shuffle(raw_im_list) 15 | 16 | return raw_im_list 17 | 18 | def Vimeo_90K_interp(root, split=1.0, single=False, task = 'interp' ): 19 | train_list = make_dataset(root,"tri_trainlist.txt") 20 | test_list = make_dataset(root,"tri_testlist.txt") 21 | train_dataset = ListDataset(root, train_list, loader=Vimeo_90K_loader) 22 | test_dataset = ListDataset(root, test_list, loader=Vimeo_90K_loader) 23 | return train_dataset, test_dataset -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .Vimeo_90K_interp import Vimeo_90K_interp 2 | 3 | __all__ = ( 4 | 'Vimeo_90K_interp', 5 | ) 6 | 7 | # Vimeo_90K = "/tmp4/wenbobao_data/vimeo_triplet" 8 | -------------------------------------------------------------------------------- /datasets/listdatasets.py: -------------------------------------------------------------------------------- 1 | import torch.utils.data as data 2 | import os 3 | import os.path 4 | from scipy.ndimage import imread 5 | import numpy as np 6 | import random 7 | 8 | def Vimeo_90K_loader(root, im_path, input_frame_size = (3, 256, 448), output_frame_size = (3, 256, 448), data_aug = True): 9 | 10 | 11 | root = os.path.join(root,'sequences',im_path) 12 | 13 | if data_aug and random.randint(0, 1): 14 | path_pre2 = os.path.join(root, "im1.png") 15 | path_mid = os.path.join(root, "im2.png") 16 | path_pre1 = os.path.join(root, "im3.png") 17 | else: 18 | path_pre1 = os.path.join(root, "im1.png") 19 | path_mid = os.path.join(root, "im2.png") 20 | path_pre2 = os.path.join(root, "im3.png") 21 | 22 | im_pre2 = imread(path_pre2) 23 | im_pre1 = imread(path_pre1) 24 | im_mid = imread(path_mid) 25 | 26 | h_offset = random.choice(range(256 - input_frame_size[1] + 1)) 27 | w_offset = random.choice(range(448 - input_frame_size[2] + 1)) 28 | 29 | im_pre2 = im_pre2[h_offset:h_offset + input_frame_size[1], w_offset: w_offset + input_frame_size[2], :] 30 | im_pre1 = im_pre1[h_offset:h_offset + input_frame_size[1], w_offset: w_offset + input_frame_size[2], :] 31 | im_mid = im_mid[h_offset:h_offset + input_frame_size[1], w_offset: w_offset + input_frame_size[2], :] 32 | 33 | if data_aug: 34 | if random.randint(0, 1): 35 | im_pre2 = np.fliplr(im_pre2) 36 | im_mid = np.fliplr(im_mid) 37 | im_pre1 = np.fliplr(im_pre1) 38 | if random.randint(0, 1): 39 | im_pre2 = np.flipud(im_pre2) 40 | im_mid = np.flipud(im_mid) 41 | im_pre1 = np.flipud(im_pre1) 42 | 43 | X0 = np.transpose(im_pre1,(2,0,1)) 44 | X2 = np.transpose(im_pre2, (2, 0, 1)) 45 | 46 | y = np.transpose(im_mid, (2, 0, 1)) 47 | return X0.astype("float32")/ 255.0, \ 48 | X2.astype("float32")/ 255.0,\ 49 | y.astype("float32")/ 255.0 50 | 51 | 52 | 53 | class ListDataset(data.Dataset): 54 | def __init__(self, root, path_list, loader=Vimeo_90K_loader): 55 | 56 | self.root = root 57 | self.path_list = path_list 58 | self.loader = loader 59 | 60 | def __getitem__(self, index): 61 | path = self.path_list[index] 62 | # print(path) 63 | image_0,image_2,image_1 = self.loader(self.root, path) 64 | return image_0,image_2,image_1 65 | 66 | def __len__(self): 67 | return len(self.path_list) 68 | -------------------------------------------------------------------------------- /demo_MiddleBury.py: -------------------------------------------------------------------------------- 1 | import time 2 | import os 3 | from torch.autograd import Variable 4 | import math 5 | import torch 6 | 7 | import random 8 | import numpy as np 9 | import numpy 10 | import networks 11 | from my_args import args 12 | 13 | from scipy.misc import imread, imsave 14 | from AverageMeter import * 15 | 16 | torch.backends.cudnn.benchmark = True # to speed up the 17 | 18 | 19 | DO_MiddleBurryOther = True 20 | MB_Other_DATA = "./MiddleBurySet/other-data/" 21 | MB_Other_RESULT = "./MiddleBurySet/other-result-author/" 22 | MB_Other_GT = "./MiddleBurySet/other-gt-interp/" 23 | if not os.path.exists(MB_Other_RESULT): 24 | os.mkdir(MB_Other_RESULT) 25 | 26 | 27 | 28 | model = networks.__dict__[args.netName](channel=args.channels, 29 | filter_size = args.filter_size , 30 | timestep=args.time_step, 31 | training=False) 32 | 33 | if args.use_cuda: 34 | model = model.cuda() 35 | 36 | args.SAVED_MODEL = './model_weights/best.pth' 37 | if os.path.exists(args.SAVED_MODEL): 38 | print("The testing model weight is: " + args.SAVED_MODEL) 39 | if not args.use_cuda: 40 | pretrained_dict = torch.load(args.SAVED_MODEL, map_location=lambda storage, loc: storage) 41 | # model.load_state_dict(torch.load(args.SAVED_MODEL, map_location=lambda storage, loc: storage)) 42 | else: 43 | pretrained_dict = torch.load(args.SAVED_MODEL) 44 | # model.load_state_dict(torch.load(args.SAVED_MODEL)) 45 | 46 | model_dict = model.state_dict() 47 | # 1. filter out unnecessary keys 48 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 49 | # 2. overwrite entries in the existing state dict 50 | model_dict.update(pretrained_dict) 51 | # 3. load the new state dict 52 | model.load_state_dict(model_dict) 53 | # 4. release the pretrained dict for saving memory 54 | pretrained_dict = [] 55 | else: 56 | print("*****************************************************************") 57 | print("**** We don't load any trained weights **************************") 58 | print("*****************************************************************") 59 | 60 | model = model.eval() # deploy mode 61 | 62 | 63 | use_cuda=args.use_cuda 64 | save_which=args.save_which 65 | dtype = args.dtype 66 | unique_id =str(random.randint(0, 100000)) 67 | print("The unique id for current testing is: " + str(unique_id)) 68 | 69 | interp_error = AverageMeter() 70 | if DO_MiddleBurryOther: 71 | subdir = os.listdir(MB_Other_DATA) 72 | gen_dir = os.path.join(MB_Other_RESULT, unique_id) 73 | os.mkdir(gen_dir) 74 | 75 | tot_timer = AverageMeter() 76 | proc_timer = AverageMeter() 77 | end = time.time() 78 | for dir in subdir: 79 | print(dir) 80 | os.mkdir(os.path.join(gen_dir, dir)) 81 | arguments_strFirst = os.path.join(MB_Other_DATA, dir, "frame10.png") 82 | arguments_strSecond = os.path.join(MB_Other_DATA, dir, "frame11.png") 83 | arguments_strOut = os.path.join(gen_dir, dir, "frame10i11.png") 84 | gt_path = os.path.join(MB_Other_GT, dir, "frame10i11.png") 85 | 86 | X0 = torch.from_numpy( np.transpose(imread(arguments_strFirst) , (2,0,1)).astype("float32")/ 255.0).type(dtype) 87 | X1 = torch.from_numpy( np.transpose(imread(arguments_strSecond) , (2,0,1)).astype("float32")/ 255.0).type(dtype) 88 | 89 | 90 | y_ = torch.FloatTensor() 91 | 92 | assert (X0.size(1) == X1.size(1)) 93 | assert (X0.size(2) == X1.size(2)) 94 | 95 | intWidth = X0.size(2) 96 | intHeight = X0.size(1) 97 | channel = X0.size(0) 98 | if not channel == 3: 99 | continue 100 | 101 | if intWidth != ((intWidth >> 7) << 7): 102 | intWidth_pad = (((intWidth >> 7) + 1) << 7) # more than necessary 103 | intPaddingLeft =int(( intWidth_pad - intWidth)/2) 104 | intPaddingRight = intWidth_pad - intWidth - intPaddingLeft 105 | else: 106 | intWidth_pad = intWidth 107 | intPaddingLeft = 32 108 | intPaddingRight= 32 109 | 110 | if intHeight != ((intHeight >> 7) << 7): 111 | intHeight_pad = (((intHeight >> 7) + 1) << 7) # more than necessary 112 | intPaddingTop = int((intHeight_pad - intHeight) / 2) 113 | intPaddingBottom = intHeight_pad - intHeight - intPaddingTop 114 | else: 115 | intHeight_pad = intHeight 116 | intPaddingTop = 32 117 | intPaddingBottom = 32 118 | 119 | pader = torch.nn.ReplicationPad2d([intPaddingLeft, intPaddingRight , intPaddingTop, intPaddingBottom]) 120 | 121 | torch.set_grad_enabled(False) 122 | X0 = Variable(torch.unsqueeze(X0,0)) 123 | X1 = Variable(torch.unsqueeze(X1,0)) 124 | X0 = pader(X0) 125 | X1 = pader(X1) 126 | 127 | if use_cuda: 128 | X0 = X0.cuda() 129 | X1 = X1.cuda() 130 | proc_end = time.time() 131 | y_s,offset,filter = model(torch.stack((X0, X1),dim = 0)) 132 | y_ = y_s[save_which] 133 | 134 | proc_timer.update(time.time() -proc_end) 135 | tot_timer.update(time.time() - end) 136 | end = time.time() 137 | print("*****************current image process time \t " + str(time.time()-proc_end )+"s ******************" ) 138 | if use_cuda: 139 | X0 = X0.data.cpu().numpy() 140 | y_ = y_.data.cpu().numpy() 141 | offset = [offset_i.data.cpu().numpy() for offset_i in offset] 142 | filter = [filter_i.data.cpu().numpy() for filter_i in filter] if filter[0] is not None else None 143 | X1 = X1.data.cpu().numpy() 144 | else: 145 | X0 = X0.data.numpy() 146 | y_ = y_.data.numpy() 147 | offset = [offset_i.data.numpy() for offset_i in offset] 148 | filter = [filter_i.data.numpy() for filter_i in filter] 149 | X1 = X1.data.numpy() 150 | 151 | 152 | 153 | X0 = np.transpose(255.0 * X0.clip(0,1.0)[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) 154 | y_ = np.transpose(255.0 * y_.clip(0,1.0)[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) 155 | offset = [np.transpose(offset_i[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) for offset_i in offset] 156 | filter = [np.transpose( 157 | filter_i[0, :, intPaddingTop:intPaddingTop + intHeight, intPaddingLeft: intPaddingLeft + intWidth], 158 | (1, 2, 0)) for filter_i in filter] if filter is not None else None 159 | X1 = np.transpose(255.0 * X1.clip(0,1.0)[0, :, intPaddingTop:intPaddingTop+intHeight, intPaddingLeft: intPaddingLeft+intWidth], (1, 2, 0)) 160 | 161 | 162 | imsave(arguments_strOut, np.round(y_).astype(numpy.uint8)) 163 | 164 | 165 | rec_rgb = imread(arguments_strOut) 166 | gt_rgb = imread(gt_path) 167 | 168 | diff_rgb = 128.0 + rec_rgb - gt_rgb 169 | avg_interp_error_abs = np.mean(np.abs(diff_rgb - 128.0)) 170 | 171 | interp_error.update(avg_interp_error_abs, 1) 172 | 173 | mse = numpy.mean((diff_rgb - 128.0) ** 2) 174 | 175 | PIXEL_MAX = 255.0 176 | psnr = 20 * math.log10(PIXEL_MAX / math.sqrt(mse)) 177 | 178 | print("interpolation error / PSNR : " + str(round(avg_interp_error_abs,4)) + " / " + str(round(psnr,4))) 179 | metrics = "The average interpolation error / PSNR for all images are : " + str(round(interp_error.avg, 4)) 180 | print(metrics) 181 | 182 | -------------------------------------------------------------------------------- /environment.yaml: -------------------------------------------------------------------------------- 1 | name: pytorch1.0.0 2 | channels: 3 | - pytorch 4 | - serge-sans-paille 5 | - anaconda 6 | - conda-forge 7 | - defaults 8 | dependencies: 9 | - ca-certificates=2019.1.23=0 10 | - certifi=2018.11.29=py36_0 11 | - cloudpickle=0.7.0=py_0 12 | - cytoolz=0.9.0.1=py36h14c3975_1 13 | - dask-core=1.1.1=py_0 14 | - decorator=4.3.2=py36_0 15 | - imageio=2.4.1=py36_0 16 | - networkx=2.2=py36_1 17 | - openssl=1.1.1=h7b6447c_0 18 | - pywavelets=1.0.1=py36hdd07704_0 19 | - scikit-image=0.14.1=py36he6710b0_0 20 | - scipy=1.1.0=py36h7c811a0_0 21 | - toolz=0.9.0=py36_0 22 | - cycler=0.10.0=py_1 23 | - expat=2.2.5=hf484d3e_1002 24 | - fontconfig=2.13.1=h2176d3f_1000 25 | - gettext=0.19.8.1=h9745a5d_1001 26 | - glib=2.56.2=had28632_1001 27 | - icu=58.2=hf484d3e_1000 28 | - kiwisolver=1.0.1=py36h6bb024c_1002 29 | - libiconv=1.15=h14c3975_1004 30 | - libprotobuf=3.6.1=hdbcaa40_1000 31 | - libuuid=2.32.1=h14c3975_1000 32 | - libxcb=1.13=h14c3975_1002 33 | - libxml2=2.9.8=h143f9aa_1005 34 | - matplotlib=3.0.2=py36_1002 35 | - matplotlib-base=3.0.2=py36h167e16e_1002 36 | - protobuf=3.6.1=py36hf484d3e_1001 37 | - pthread-stubs=0.4=h14c3975_1001 38 | - pyparsing=2.3.1=py_0 39 | - pyqt=5.6.0=py36h13b7fb3_1008 40 | - python-dateutil=2.8.0=py_0 41 | - sip=4.18.1=py36hf484d3e_1000 42 | - tensorboardx=1.6=py_0 43 | - tk=8.6.9=h84994c4_1000 44 | - tornado=5.1.1=py36h14c3975_1000 45 | - xorg-libxau=1.0.9=h14c3975_0 46 | - xorg-libxdmcp=1.1.2=h14c3975_1007 47 | - blas=1.0=mkl 48 | - cffi=1.11.5=py36he75722e_1 49 | - cudatoolkit=9.0=h13b8566_0 50 | - dbus=1.13.2=h714fa37_1 51 | - freetype=2.9.1=h8a8886c_1 52 | - gst-plugins-base=1.14.0=hbbd80ab_1 53 | - gstreamer=1.14.0=hb453b48_1 54 | - intel-openmp=2019.1=144 55 | - isl=0.12.2=0 56 | - jpeg=9b=h024ee3a_2 57 | - libedit=3.1.20181209=hc058e9b_0 58 | - libffi=3.2.1=hd88cf55_4 59 | - libgcc-ng=8.2.0=hdf63c60_1 60 | - libgfortran-ng=7.3.0=hdf63c60_0 61 | - libpng=1.6.36=hbc83047_0 62 | - libstdcxx-ng=8.2.0=hdf63c60_1 63 | - libtiff=4.0.10=h2733197_2 64 | - mkl=2019.1=144 65 | - mkl_fft=1.0.10=py36ha843d7b_0 66 | - mkl_random=1.0.2=py36hd81dba3_0 67 | - mpc=1.0.3=hf803216_4 68 | - mpfr=3.1.5=h12ff648_1 69 | - ncurses=6.1=he6710b0_1 70 | - ninja=1.8.2=py36h6bb024c_1 71 | - numpy=1.15.4=py36h7e9f1db_0 72 | - numpy-base=1.15.4=py36hde5b4d6_0 73 | - olefile=0.46=py36_0 74 | - pcre=8.42=h439df22_0 75 | - pillow=5.4.1=py36h34e0f95_0 76 | - pip=19.0.1=py36_0 77 | - pycparser=2.19=py36_0 78 | - python=3.6.8=h0371630_0 79 | - qt=5.6.3=h8bf5577_3 80 | - readline=7.0=h7b6447c_5 81 | - setuptools=40.8.0=py36_0 82 | - six=1.12.0=py36_0 83 | - sqlite=3.26.0=h7b6447c_0 84 | - wheel=0.32.3=py36_0 85 | - xz=5.2.4=h14c3975_4 86 | - zlib=1.2.11=h7b6447c_3 87 | - zstd=1.3.7=h0b5b093_0 88 | - pytorch=1.0.1=py3.6_cuda9.0.176_cudnn7.4.2_2 89 | - torchvision=0.2.1=py_2 90 | - cloog=0.18.1=1 91 | - gcc_49=4.9.1=6 92 | - gmp=5.1.3=0 93 | - pip: 94 | - correlation-cuda==0.0.0 95 | - dask==1.1.1 96 | - depthflowprojection-cuda==0.0.0 97 | - filterinterpolation-cuda==0.0.0 98 | - flowprojection-cuda==0.0.0 99 | - interpolation-cuda==0.0.0 100 | - interpolationch-cuda==0.0.0 101 | - mindepthflowprojection-cuda==0.0.0 102 | - separableconv-cuda==0.0.0 103 | - separableconvflow-cuda==0.0.0 104 | - torch==1.0.1.post2 105 | prefix: /home/wenbobao/anaconda3_new/envs/pytorch1.0.0 106 | 107 | -------------------------------------------------------------------------------- /loss_function.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | import sys 5 | import threading 6 | import torch 7 | from torch.autograd import Variable 8 | from lr_scheduler import * 9 | from torch.autograd import gradcheck 10 | 11 | import numpy 12 | 13 | 14 | 15 | 16 | def charbonier_loss(x,epsilon): 17 | loss = torch.mean(torch.sqrt(x * x + epsilon * epsilon)) 18 | return loss 19 | def negPSNR_loss(x,epsilon): 20 | loss = torch.mean(torch.mean(torch.mean(torch.sqrt(x * x + epsilon * epsilon),dim=1),dim=1),dim=1) 21 | return torch.mean(-torch.log(1.0/loss) /100.0) 22 | 23 | def tv_loss(x,epsilon): 24 | loss = torch.mean( torch.sqrt( 25 | (x[:, :, :-1, :-1] - x[:, :, 1:, :-1]) ** 2 + 26 | (x[:, :, :-1, :-1] - x[:, :, :-1, 1:]) ** 2 + epsilon *epsilon 27 | ) 28 | ) 29 | return loss 30 | 31 | 32 | def gra_adap_tv_loss(flow, image, epsilon): 33 | w = torch.exp( - torch.sum( torch.abs(image[:,:,:-1, :-1] - image[:,:,1:, :-1]) + 34 | torch.abs(image[:,:,:-1, :-1] - image[:,:,:-1, 1:]), dim = 1)) 35 | tv = torch.sum(torch.sqrt((flow[:, :, :-1, :-1] - flow[:, :, 1:, :-1]) ** 2 + (flow[:, :, :-1, :-1] - flow[:, :, :-1, 1:]) ** 2 + epsilon *epsilon) ,dim=1) 36 | loss = torch.mean( w * tv ) 37 | return loss 38 | 39 | def smooth_loss(x,epsilon): 40 | loss = torch.mean( 41 | torch.sqrt( 42 | (x[:,:,:-1,:-1] - x[:,:,1:,:-1]) **2 + 43 | (x[:,:,:-1,:-1] - x[:,:,:-1,1:]) **2+ epsilon**2 44 | ) 45 | ) 46 | return loss 47 | 48 | 49 | def motion_sym_loss(offset, epsilon, occlusion = None): 50 | if occlusion == None: 51 | # return torch.mean(torch.sqrt( (offset[:,:2,...] + offset[:,2:,...])**2 + epsilon **2)) 52 | return torch.mean(torch.sqrt( (offset[0] + offset[1])**2 + epsilon **2)) 53 | else: 54 | # TODO: how to design the occlusion aware offset symmetric loss? 55 | # return torch.mean(torch.sqrt((offset[:,:2,...] + offset[:,2:,...])**2 + epsilon **2)) 56 | return torch.mean(torch.sqrt((offset[0] + offset[1])**2 + epsilon **2)) 57 | 58 | 59 | 60 | 61 | def part_loss(diffs, offsets, occlusions, images, epsilon, use_negPSNR=False): 62 | if use_negPSNR: 63 | pixel_loss = [negPSNR_loss(diff, epsilon) for diff in diffs] 64 | else: 65 | pixel_loss = [charbonier_loss(diff, epsilon) for diff in diffs] 66 | #offset_loss = [tv_loss(offset[0], epsilon) + tv_loss(offset[1], epsilon) for offset in 67 | # offsets] 68 | 69 | if offsets[0][0] is not None: 70 | offset_loss = [gra_adap_tv_loss(offset[0],images[0], epsilon) + gra_adap_tv_loss(offset[1], images[1], epsilon) for offset in 71 | offsets] 72 | else: 73 | offset_loss = [Variable(torch.zeros(1).cuda())] 74 | # print(torch.max(occlusions[0])) 75 | # print(torch.min(occlusions[0])) 76 | # print(torch.mean(occlusions[0])) 77 | 78 | # occlusion_loss = [smooth_loss(occlusion, epsilon) + charbonier_loss(occlusion - 0.5, epsilon) for occlusion in occlusions] 79 | # occlusion_loss = [smooth_loss(occlusion, epsilon) + charbonier_loss(occlusion[:, 0, ...] - occlusion[:, 1, ...], epsilon) for occlusion in occlusions] 80 | 81 | 82 | 83 | sym_loss = [motion_sym_loss(offset,epsilon=epsilon) for offset in offsets] 84 | # sym_loss = [ motion_sym_loss(offset,occlusion) for offset,occlusion in zip(offsets,occlusions)] 85 | return pixel_loss, offset_loss, sym_loss 86 | 87 | -------------------------------------------------------------------------------- /my_args.py: -------------------------------------------------------------------------------- 1 | import os 2 | import datetime 3 | import argparse 4 | import numpy 5 | import networks 6 | import torch 7 | modelnames = networks.__all__ 8 | # import datasets 9 | datasetNames = ('Vimeo_90K_interp') #datasets.__all__ 10 | 11 | parser = argparse.ArgumentParser(description='DAIN') 12 | 13 | parser.add_argument('--debug',action = 'store_true', help='Enable debug mode') 14 | parser.add_argument('--netName', type=str, default='DAIN', 15 | choices = modelnames,help = 'model architecture: ' + 16 | ' | '.join(modelnames) + 17 | ' (default: DAIN)') 18 | 19 | parser.add_argument('--datasetName', default='Vimeo_90K_interp', 20 | choices= datasetNames,nargs='+', 21 | help='dataset type : ' + 22 | ' | '.join(datasetNames) + 23 | ' (default: Vimeo_90K_interp)') 24 | parser.add_argument('--datasetPath',default='',help = 'the path of selected datasets') 25 | parser.add_argument('--dataset_split', type = int, default=97, help = 'Split a dataset into trainining and validation by percentage (default: 97)') 26 | 27 | parser.add_argument('--seed', type=int, default=1, help='random seed (default: 1)') 28 | 29 | parser.add_argument('--numEpoch', '-e', type = int, default=100, help= 'Number of epochs to train(default:150)') 30 | 31 | parser.add_argument('--batch_size', '-b',type = int ,default=1, help = 'batch size (default:1)' ) 32 | parser.add_argument('--workers', '-w', type =int,default=8, help = 'parallel workers for loading training samples (default : 1.6*10 = 16)') 33 | parser.add_argument('--channels', '-c', type=int,default=3,choices = [1,3], help ='channels of images (default:3)') 34 | parser.add_argument('--filter_size', '-f', type=int, default=4, help = 'the size of filters used (default: 4)', 35 | choices=[2,4,6, 5,51] 36 | ) 37 | 38 | 39 | parser.add_argument('--lr', type =float, default= 0.002, help= 'the basic learning rate for three subnetworks (default: 0.002)') 40 | parser.add_argument('--rectify_lr', type=float, default=0.001, help = 'the learning rate for rectify/refine subnetworks (default: 0.001)') 41 | 42 | parser.add_argument('--save_which', '-s', type=int, default=1, choices=[0,1], help='choose which result to save: 0 ==> interpolated, 1==> rectified') 43 | parser.add_argument('--time_step', type=float, default=0.5, help='choose the time steps') 44 | parser.add_argument('--flow_lr_coe', type = float, default=0.01, help = 'relative learning rate w.r.t basic learning rate (default: 0.01)') 45 | parser.add_argument('--occ_lr_coe', type = float, default=1.0, help = 'relative learning rate w.r.t basic learning rate (default: 1.0)') 46 | parser.add_argument('--filter_lr_coe', type = float, default=1.0, help = 'relative learning rate w.r.t basic learning rate (default: 1.0)') 47 | parser.add_argument('--ctx_lr_coe', type = float, default=1.0, help = 'relative learning rate w.r.t basic learning rate (default: 1.0)') 48 | parser.add_argument('--depth_lr_coe', type = float, default=0.001, help = 'relative learning rate w.r.t basic learning rate (default: 0.01)') 49 | # parser.add_argument('--deblur_lr_coe', type = float, default=0.01, help = 'relative learning rate w.r.t basic learning rate (default: 0.01)') 50 | 51 | parser.add_argument('--alpha', type=float,nargs='+', default=[0.0, 1.0], help= 'the ration of loss for interpolated and rectified result (default: [0.0, 1.0])') 52 | 53 | parser.add_argument('--epsilon', type = float, default=1e-6, help = 'the epsilon for charbonier loss,etc (default: 1e-6)') 54 | parser.add_argument('--weight_decay', type = float, default=0, help = 'the weight decay for whole network ' ) 55 | parser.add_argument('--patience', type=int, default=5, help = 'the patience of reduce on plateou') 56 | parser.add_argument('--factor', type = float, default=0.2, help = 'the factor of reduce on plateou') 57 | # 58 | parser.add_argument('--pretrained', dest='SAVED_MODEL', default=None, help ='path to the pretrained model weights') 59 | parser.add_argument('--no-date', action='store_true', help='don\'t append date timestamp to folder' ) 60 | parser.add_argument('--use_cuda', default= True, type = bool, help='use cuda or not') 61 | parser.add_argument('--use_cudnn',default=1,type=int, help = 'use cudnn or not') 62 | parser.add_argument('--dtype', default=torch.cuda.FloatTensor, choices = [torch.cuda.FloatTensor,torch.FloatTensor],help = 'tensor data type ') 63 | # parser.add_argument('--resume', default='', type=str, help='path to latest checkpoint (default: none)') 64 | 65 | 66 | parser.add_argument('--uid', type=str, default= None, help='unique id for the training') 67 | parser.add_argument('--force', action='store_true', help='force to override the given uid') 68 | 69 | # Colab version 70 | parser.add_argument('--start_frame', type = int, default = 1, help='first frame number to process') 71 | parser.add_argument('--end_frame', type = int, default = 100, help='last frame number to process') 72 | parser.add_argument('--frame_input_dir', type = str, default = '/content/DAIN/input_frames', help='frame input directory') 73 | parser.add_argument('--frame_output_dir', type = str, default = '/content/DAIN/output_frames', help='frame output directory') 74 | 75 | args = parser.parse_args() 76 | 77 | import shutil 78 | 79 | if args.uid == None: 80 | unique_id = str(numpy.random.randint(0, 100000)) 81 | print("revise the unique id to a random numer " + str(unique_id)) 82 | args.uid = unique_id 83 | timestamp = datetime.datetime.now().strftime("%a-%b-%d-%H-%M") 84 | save_path = './model_weights/'+ args.uid +'-' + timestamp 85 | else: 86 | save_path = './model_weights/'+ str(args.uid) 87 | 88 | # print("no pth here : " + save_path + "/best"+".pth") 89 | if not os.path.exists(save_path + "/best"+".pth"): 90 | # print("no pth here : " + save_path + "/best" + ".pth") 91 | os.makedirs(save_path,exist_ok=True) 92 | else: 93 | if not args.force: 94 | raise("please use another uid ") 95 | else: 96 | print("override this uid" + args.uid) 97 | for m in range(1,10): 98 | if not os.path.exists(save_path+"/log.txt.bk" + str(m)): 99 | shutil.copy(save_path+"/log.txt", save_path+"/log.txt.bk"+str(m)) 100 | shutil.copy(save_path+"/args.txt", save_path+"/args.txt.bk"+str(m)) 101 | break 102 | 103 | 104 | 105 | parser.add_argument('--save_path',default=save_path,help = 'the output dir of weights') 106 | parser.add_argument('--log', default = save_path+'/log.txt', help = 'the log file in training') 107 | parser.add_argument('--arg', default = save_path+'/args.txt', help = 'the args used') 108 | 109 | args = parser.parse_args() 110 | 111 | 112 | with open(args.log, 'w') as f: 113 | f.close() 114 | with open(args.arg, 'w') as f: 115 | print(args) 116 | print(args,file=f) 117 | f.close() 118 | if args.use_cudnn: 119 | print("cudnn is used") 120 | torch.backends.cudnn.benchmark = True # to speed up the 121 | else: 122 | print("cudnn is not used") 123 | torch.backends.cudnn.benchmark = False # to speed up the 124 | 125 | -------------------------------------------------------------------------------- /my_package/DepthFlowProjection/DepthFlowProjectionLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | #import _ext.my_lib as my_lib 5 | import depthflowprojection_cuda as my_lib 6 | 7 | class DepthFlowProjectionLayer(Function): 8 | def __init__(self,requires_grad): 9 | super(DepthFlowProjectionLayer,self).__init__() 10 | # self.requires_grad = requires_grad 11 | 12 | @staticmethod 13 | def forward(ctx, input1, input2, requires_grad): 14 | # print("Depth Aware Flow Projection") 15 | assert(input1.is_contiguous()) 16 | assert(input2.is_contiguous()) 17 | # self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 18 | # self.input2 = input2.contiguous() 19 | fillhole = 1 if requires_grad == False else 0 20 | # if input1.is_cuda: 21 | # self.device = torch.cuda.current_device() 22 | # else: 23 | # self.device = -1 24 | 25 | # count = torch.zeros(input1.size(0),1,input1.size(2),input1.size(3)) # for accumulating the homography projections 26 | # output = torch.zeros(input1.size()) 27 | 28 | if input1.is_cuda: 29 | # output = output.cuda() 30 | # count = count.cuda() 31 | # print("correct") 32 | count = torch.cuda.FloatTensor().resize_(input1.size(0), 1, input1.size(2), input1.size(3)).zero_() 33 | output = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 34 | err = my_lib.DepthFlowProjectionLayer_gpu_forward(input1,input2, count,output, fillhole) 35 | else: 36 | # output = torch.cuda.FloatTensor(input1.data.size()) 37 | count = torch.FloatTensor().resize_(input1.size(0), 1, input1.size(2), input1.size(3)).zero_() 38 | output = torch.FloatTensor().resize_(input1.size()).zero_() 39 | err = my_lib.DepthFlowProjectionLayer_cpu_forward(input1,input2, count, output,fillhole) 40 | if err != 0: 41 | print(err) 42 | # output = output/count # to divide the counter 43 | 44 | # self.count = count #to keep this 45 | # self.output = output 46 | 47 | ctx.save_for_backward(input1, input2,count,output) 48 | ctx.fillhole = fillhole 49 | 50 | # print(self.input1[0, 0, :10, :10]) 51 | # print(self.count[0, 0, :10, :10]) 52 | # print(self.input1[0, 0, -10:, -10:]) 53 | # print(self.count[0, 0, -10:, -10:]) 54 | 55 | # the function returns the output to its caller 56 | return output 57 | 58 | @staticmethod 59 | def backward(ctx, gradoutput): 60 | # print("Backward of Filter Interpolation Layer") 61 | # gradinput1 = input1.new().zero_() 62 | # gradinput2 = input2.new().zero_() 63 | # gradinput1 = torch.zeros(self.input1.size()) 64 | 65 | input1, input2, count, output = ctx.saved_tensors 66 | # fillhole = ctx.fillhole 67 | 68 | if input1.is_cuda: 69 | # print("CUDA backward") 70 | # gradinput1 = gradinput1.cuda(self.device) 71 | gradinput1 = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 72 | gradinput2 = torch.cuda.FloatTensor().resize_(input2.size()).zero_() 73 | 74 | err = my_lib.DepthFlowProjectionLayer_gpu_backward(input1,input2, 75 | count, output, 76 | gradoutput, gradinput1,gradinput2) 77 | # print(err) 78 | if err != 0 : 79 | print(err) 80 | 81 | else: 82 | # print("CPU backward") 83 | # print(gradoutput) 84 | gradinput1 = torch.FloatTensor().resize_(input1.size()).zero_() 85 | gradinput2 = torch.FloatTensor().resize_(input2.size()).zero_() 86 | err = my_lib.DepthFlowProjectionLayer_cpu_backward(input1, input2, 87 | count, output, 88 | gradoutput, gradinput1,gradinput2) 89 | # print(err) 90 | if err != 0: 91 | print(err) 92 | # print(gradinput1) 93 | # print(gradinput2) 94 | 95 | # print(gradinput1) 96 | 97 | return gradinput1,gradinput2,None 98 | -------------------------------------------------------------------------------- /my_package/DepthFlowProjection/DepthFlowProjectionModule.py: -------------------------------------------------------------------------------- 1 | # modules/FlowProjectionModule.py 2 | from torch.nn.modules.module import Module 3 | from .DepthFlowProjectionLayer import DepthFlowProjectionLayer #, FlowFillholeLayer 4 | 5 | __all__ =['DepthFlowProjectionModule'] 6 | 7 | class DepthFlowProjectionModule(Module): 8 | def __init__(self, requires_grad = True): 9 | super(DepthFlowProjectionModule, self).__init__() 10 | self.requires_grad = requires_grad 11 | # self.f = DepthFlowProjectionLayer(requires_grad) 12 | 13 | def forward(self, input1, input2): 14 | return DepthFlowProjectionLayer.apply(input1, input2,self.requires_grad) 15 | 16 | # class FlowFillholeModule(Module): 17 | # def __init__(self,hole_value = -10000.0): 18 | # super(FlowFillholeModule, self).__init__() 19 | # self.f = FlowFillholeLayer() 20 | # 21 | # def forward(self, input1): 22 | # return self.f(input1) 23 | 24 | #we actually dont need to write the backward code for a module, since we have 25 | 26 | -------------------------------------------------------------------------------- /my_package/DepthFlowProjection/__init__.py: -------------------------------------------------------------------------------- 1 | from .DepthFlowProjectionModule import * 2 | -------------------------------------------------------------------------------- /my_package/DepthFlowProjection/depthflowprojection_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "depthflowprojection_cuda_kernel.cuh" 8 | 9 | 10 | int DepthFlowProjectionLayer_gpu_forward( 11 | at::Tensor& input1, 12 | at::Tensor& input2, 13 | at::Tensor& count, 14 | at::Tensor& output, 15 | int fillhole 16 | ) 17 | { 18 | 19 | int error = 1 ; 20 | 21 | int channel = input1.size( 1); 22 | if(channel!= 2) return error; 23 | int batch = input1.size(0); 24 | 25 | int h = input1.size(2); 26 | int w = input1.size(3); 27 | 28 | if(input2.size(1) !=1 ) return error; 29 | 30 | int input1_b_stride = input1.stride(0); 31 | int input1_c_stride = input1.stride(1); 32 | int input1_h_stride = input1.stride(2); 33 | int input1_w_stride = input1.stride(3); 34 | 35 | int input2_b_stride = input2.stride(0); 36 | int input2_c_stride = input2.stride(1); 37 | int input2_h_stride = input2.stride(2); 38 | int input2_w_stride = input2.stride(3); 39 | 40 | int count_b_stride = count.stride(0); 41 | int count_c_stride = count.stride(1); 42 | int count_h_stride = count.stride(2); 43 | int count_w_stride = count.stride(3); 44 | //TODO: do we need to assert the w_stride to be 1 45 | //if(w_stride !=1) return error; 46 | if(input1_b_stride != output.stride(0)) return error; 47 | if(input1_c_stride != output.stride(1)) return error; 48 | 49 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 50 | // printf("In gpu forward\n"); 51 | error = DepthFlowProjection_gpu_forward_kernel( 52 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 53 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 54 | nElement,w,h,channel,batch,fillhole, 55 | 56 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 57 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 58 | count_b_stride,count_c_stride,count_h_stride,count_w_stride, 59 | 60 | input1, 61 | input2, 62 | count, 63 | output); 64 | if (error) {AT_ERROR("CUDA call failed");} 65 | 66 | return error; 67 | 68 | } 69 | 70 | int DepthFlowProjectionLayer_gpu_backward( 71 | at::Tensor& input1, 72 | at::Tensor& input2, 73 | at::Tensor& count, 74 | at::Tensor& output, 75 | at::Tensor& gradoutput, 76 | at::Tensor& gradinput1, 77 | at::Tensor& gradinput2 78 | ) 79 | { 80 | int error = 1 ; 81 | int channel = input1.size( 1); 82 | if(channel!=2) return error; 83 | int batch = input1.size(0); 84 | if(count.size( 0) != batch) return error; 85 | if(count.size(1) != 1) return error; 86 | 87 | int h = input1.size(2); 88 | int w = input1.size(3); 89 | if(input2.size(1) !=1 ) return error; 90 | if(count.size(2) != h) return error;// to add some checkpoint 91 | if(count.size(3) != w) return error; 92 | 93 | int input1_b_stride = input1.stride(0); 94 | int input1_c_stride = input1.stride(1); 95 | int input1_h_stride = input1.stride(2); 96 | int input1_w_stride = input1.stride(3); 97 | 98 | int input2_b_stride = input2.stride(0); 99 | int input2_c_stride = input2.stride(1); 100 | int input2_h_stride = input2.stride(2); 101 | int input2_w_stride = input2.stride(3); 102 | 103 | int count_b_stride = count.stride(0); 104 | int count_c_stride = count.stride(1); 105 | int count_h_stride = count.stride(2); 106 | int count_w_stride = count.stride(3); 107 | //TODO: do we need to assert the w_stride to be 1 108 | //if(w_stride !=1) return error; 109 | if(input1_b_stride != gradinput1.stride(0)) return error; 110 | if(input1_c_stride != gradinput1.stride(1)) return error; 111 | 112 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 113 | // printf("GPU backward: %d,%d,%d,%d\n", count_b_stride,count_c_stride,count_h_stride,count_w_stride); 114 | 115 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 116 | 117 | error = DepthFlowProjection_gpu_backward_kernel( 118 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 119 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 120 | nElement, //to let the nummous 121 | w,h,channel,batch, 122 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 123 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 124 | count_b_stride,count_c_stride,count_h_stride,count_w_stride, 125 | 126 | input1, 127 | input2, 128 | count, 129 | output, 130 | gradoutput, 131 | gradinput1, 132 | gradinput2 133 | ); 134 | if (error) {AT_ERROR("CUDA call failed");} 135 | //printf("Am I good in backward function %d",error); 136 | 137 | return error; 138 | 139 | } 140 | 141 | 142 | 143 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 144 | m.def("DepthFlowProjectionLayer_gpu_forward", &DepthFlowProjectionLayer_gpu_forward, "DepthFlowProjection forward (CUDA)"); 145 | m.def("DepthFlowProjectionLayer_gpu_backward", &DepthFlowProjectionLayer_gpu_backward, "DepthFlowProjection backward (CUDA)"); 146 | } 147 | -------------------------------------------------------------------------------- /my_package/DepthFlowProjection/depthflowprojection_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int DepthFlowProjection_gpu_forward_kernel( 8 | cudaStream_t stream, const int nElement, 9 | const int w, const int h, const int channel, const int batch, const int fillhole, 10 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 11 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 12 | const int count_b_stride, const int count_c_stride, const int count_h_stride, const int count_w_stride, 13 | 14 | at::Tensor& input1, at::Tensor& input2, 15 | at::Tensor& count, 16 | at::Tensor& output 17 | 18 | ); 19 | 20 | int DepthFlowProjection_gpu_backward_kernel( 21 | cudaStream_t stream, 22 | const int nElement, 23 | const int w, 24 | const int h, 25 | const int channel, 26 | const int batch, 27 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 28 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 29 | const int count_b_stride, const int count_c_stride, const int count_h_stride, const int count_w_stride, 30 | 31 | at::Tensor& input1, 32 | at::Tensor& input2, 33 | at::Tensor& count, 34 | at::Tensor& output, 35 | at::Tensor& gradoutput, 36 | at::Tensor& gradinput1, 37 | at::Tensor& gradinput2 38 | ); 39 | -------------------------------------------------------------------------------- /my_package/DepthFlowProjection/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='depthflowprojection_cuda', 12 | ext_modules=[ 13 | CUDAExtension('depthflowprojection_cuda', [ 14 | 'depthflowprojection_cuda.cc', 15 | 'depthflowprojection_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/FilterInterpolation/FilterInterpolationModule.py: -------------------------------------------------------------------------------- 1 | # modules/AdaptiveInterpolationLayer.py 2 | from torch.nn import Module 3 | import torch 4 | from torch.autograd import Variable 5 | from torch.autograd import gradcheck 6 | from .FilterInterpolationLayer import FilterInterpolationLayer,WeightLayer, PixelValueLayer,PixelWeightLayer,ReliableWeightLayer 7 | 8 | class FilterInterpolationModule(Module): 9 | def __init__(self): 10 | super(FilterInterpolationModule, self).__init__() 11 | # self.f = FilterInterpolationLayer() 12 | 13 | def forward(self, input1, input2, input3): 14 | return FilterInterpolationLayer.apply(input1, input2, input3) 15 | 16 | #we actually dont need to write the backward code for a module, since we have 17 | 18 | #class WeightModule(Module): 19 | # def __init__(self): 20 | # super(WeightModule, self).__init__() 21 | # self.f = WeightLayer() 22 | # 23 | # def forward(self, input1, input2, input3): 24 | # return self.f(input1, input2, input3) 25 | class AdaptiveWeightInterpolationModule(Module): 26 | def __init__(self, training = False, threshhold = 1e-6, 27 | lambda_e = 30.0/255.0, lambda_v = 1.0, Nw = 3.0, 28 | sigma_d =1.5, tao_r = 0.05, Prowindow = 2 ): 29 | super(AdaptiveWeightInterpolationModule, self).__init__() 30 | 31 | self.calc_weight1 = WeightLayer(lambda_e, lambda_v, Nw ) 32 | self.padder1 = torch.nn.ReplicationPad2d([0, 1 , 0, 1]) 33 | self.interpolate1 = PixelValueLayer(sigma_d, tao_r , Prowindow) 34 | self.interpolate1_1 = PixelWeightLayer(101* threshhold, sigma_d,tao_r, Prowindow) 35 | # self.interpolate_R1 = ReliableValueLayer(Nw, tao_r , Prowindow) 36 | self.interpolate_R1_1 = ReliableWeightLayer(101* threshhold, sigma_d,tao_r, Prowindow) 37 | 38 | self.calc_weight2 = WeightLayer(lambda_e, lambda_v,Nw) 39 | self.padder2 = torch.nn.ReplicationPad2d([0, 1 , 0, 1]) 40 | self.interpolate2 = PixelValueLayer(sigma_d, tao_r , Prowindow ) 41 | self.interpolate2_1 = PixelWeightLayer(101*threshhold,sigma_d,tao_r, Prowindow) 42 | #self.interpolate_R2 = ReliableValueLayer(Nw, tao_r , Prowindow) 43 | self.interpolate_R2_1 = ReliableWeightLayer(101*threshhold, sigma_d,tao_r, Prowindow) 44 | 45 | self.training = training 46 | self.threshold = threshhold 47 | return 48 | #self.lambda_e = lambda_e 49 | #self.lambda_v = lambda_v 50 | #self.sigma_d = sigma_d 51 | #self.Nw = Nw 52 | #self.tao_r = tao_r #maybe not useable 53 | #self.Prowindow = Prowindow 54 | # lambda_e = self.lambda_e , lambda_v = self.lambda_v,Nw = self.Nw 55 | # sigma_d = self.sigma_d, tao_r = self.tao_r , Prowindow = self.Prowindow 56 | #self.sigma_d, self.tao_r , self.Prowindow 57 | 58 | 59 | # input1 ==> ref1 image 60 | # #input2 ==> ref2 image 61 | # input3 ==> ref1 flow 62 | # input4 ==> ref2 flow 63 | def forward(self, input1, input2, input3, input4): 64 | epsilon = 1e-6 65 | #flow1_grad = torch.sum(torch.sqrt( 66 | # (input3[:, :, :-1, :-1] - input3[:, :, 1:, :-1]) ** 2 + 67 | # (input3[:, :, :-1, :-1] - input3[:, :, :-1, 1:]) ** 2 + epsilon * epsilon 68 | # ), dim = 1,keepdim =True) 69 | #flow1_grad = self.padder1(flow1_grad) 70 | # if input1.is_cuda: 71 | # err = gradcheck(self.calc_weight1,(Variable(input1.data,requires_grad=True), 72 | # Variable(input2 .data,requires_grad=True), 73 | # Variable(input3.data,requires_grad= True), 74 | # # Variable(flow1_grad.data,requires_grad=True) 75 | # ), eps=1e-3) 76 | # print(err) 77 | # pass 78 | #input1.requires_grad = True 79 | #input2.requires_grad = True 80 | 81 | flow_weight1 = self.calc_weight1(input1,input2,input3 ) 82 | # if flow1_grad.is_cuda: 83 | # err = gradcheck(self.interpolate1,(Variable(input1.data,requires_grad=True), 84 | # Variable(input3.data,requires_grad= True), 85 | # Variable(flow_weight1.data,requires_grad=True)), eps=1e-3) 86 | # err = gradcheck(self.interpolate1_1, (Variable(input3.data,requires_grad=True), 87 | # Variable(flow_weight1.data, requires_grad =True)),eps=1e-3) 88 | # err = gradcheck(self.interpolate_R1_1,(input3,),eps=1e-3) 89 | # print(err) 90 | # print(flow_weight1[0,:,50:100,50:100]) 91 | p1 = self.interpolate1(input1, input3, flow_weight1) 92 | p1_r,p1_g,p1_b = torch.split(p1,1,dim=1) 93 | pw1 = self.interpolate1_1(input3, flow_weight1) 94 | i1_r,i1_g,i1_b = (p1_r)/(pw1+self.threshold),\ 95 | (p1_g)/(pw1+self.threshold), \ 96 | (p1_b)/(pw1+self.threshold) 97 | #if not self.training: 98 | # i1_r[pw1<=10*self.threshold], i1_g[pw1<=10*self.threshold], i1_b[pw1<=10*self.threshold] = 0,0,0 99 | #i1 = torch.cat((i1_r,i1_g,i1_b),dim=1 100 | #r1 = self.interpolate_R1(input3, flow_weight1) 101 | r1 = pw1 102 | rw1 = self.interpolate_R1_1(input3) 103 | w1 = (r1)/(rw1+self.threshold) 104 | # if torch.sum(w1 <= 0).cpu().data.numpy()[0] > 0: 105 | # pass 106 | # print("there are holes in i1 :" ) 107 | # print(torch.sum(w1 <= 0)) 108 | #if not self.training: 109 | # w1[rw1 <=10*self.threshold] = 0 110 | 111 | # flow2_grad = torch.sum(torch.sqrt( 112 | # (input4[:, :, :-1, :-1] - input4[:, :, 1:, :-1]) ** 2 + 113 | # (input4[:, :, :-1, :-1] - input4[:, :, :-1, 1:]) ** 2 + epsilon * epsilon 114 | # ), dim = 1,keepdim=True) 115 | # flow2_grad = self.padder2(flow2_grad) 116 | 117 | flow_weight2 = self.calc_weight2(input2,input1,input4) 118 | p2 = self.interpolate2(input2, input4, flow_weight2) 119 | p2_r,p2_g,p2_b = torch.split(p2,1,dim=1) 120 | pw2 = self.interpolate2_1(input4, flow_weight2) 121 | i2_r,i2_g,i2_b = (p2_r)/(pw2+self.threshold),\ 122 | (p2_g)/(pw2+self.threshold), \ 123 | (p2_b)/(pw2+self.threshold) 124 | #if not self.training: 125 | # i2_r[pw2<=10*self.threshold], i2_g[pw2<=10*self.threshold], i2_b[pw2<=10*self.threshold] = 0,0,0 126 | #i2 = torch.cat((p2[:,0,...] /pw2, p2[:,1,...] /pw2, p2[:,2,...]/pw2),dim=1) 127 | #r2 = self.interpolate_R2(input4, flow_weight2) 128 | r2 = pw2 129 | rw2 = self.interpolate_R2_1(input4) 130 | w2 = (r2)/(rw2+self.threshold) 131 | #if torch.sum(w2 <= 0).cpu().data.numpy()[0] > 0: 132 | # pass 133 | # print("there are holes in i2 :" ) 134 | # print(torch.sum(w2 <= 0)) 135 | #if not self.training: 136 | # w2[rw2 <= 10*self.threshold] = 0 137 | # i = (i1 * w1 + i2 * w2 )/ (w1 + w2) 138 | 139 | w = w1+w2 140 | i_r = (i1_r * w1 + i2_r * w2)/ (w + self.threshold) #(w1 + w2) 141 | i_g = (i1_g * w1 + i2_g * w2)/ (w + self.threshold) #(w1 + w2) 142 | i_b = (i1_b * w1 + i2_b * w2)/ (w + self.threshold) #(w1 + w2) 143 | #if torch.sum(w <= 0).cpu().data.numpy()[0] > 0: 144 | # print("there are holes in i :") 145 | # print(torch.sum(w <= 0)) 146 | if not self.training: 147 | i_r[w<= 10*self.threshold], i_g[w<=10*self.threshold], i_b[w<=10*self.threshold] = 0,0,0 148 | w[w <= 10 *self.threshold] = 0 149 | i = torch.cat((i_r,i_g,i_b),dim=1) 150 | return i 151 | -------------------------------------------------------------------------------- /my_package/FilterInterpolation/__init__.py: -------------------------------------------------------------------------------- 1 | from .FilterInterpolationModule import * 2 | -------------------------------------------------------------------------------- /my_package/FilterInterpolation/filterinterpolation_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "filterinterpolation_cuda_kernel.cuh" 8 | 9 | 10 | 11 | int FilterInterpolationLayer_gpu_forward( 12 | at::Tensor& input1, 13 | at::Tensor& input2, 14 | at::Tensor& input3, 15 | at::Tensor& output 16 | 17 | ) 18 | { 19 | int error = 1 ; 20 | 21 | int channel = input1.size( 1); 22 | //if(channel!=3) return error; 23 | int batch = input1.size(0); 24 | if(input2.size( 0) != batch) return error; 25 | if(input2.size(1) != 2) return error; 26 | 27 | int h = input1.size(2); 28 | int w = input1.size(3); 29 | if(input2.size(2) != h) return error;// to add some checkpoint 30 | if(input2.size(3) != w) return error; 31 | 32 | int filter_size2 = input3.size( 1); 33 | int filter_size = (int) sqrt((float) filter_size2); 34 | // printf("filter size is: %d,or %f", filter_size, sqrt((float)filter_size2)); 35 | 36 | 37 | int input1_b_stride = input1.stride(0); 38 | int input1_c_stride = input1.stride(1); 39 | int input1_h_stride = input1.stride(2); 40 | int input1_w_stride = input1.stride(3); 41 | 42 | int input2_b_stride = input2.stride(0); 43 | int input2_c_stride = input2.stride(1); 44 | int input2_h_stride = input2.stride(2); 45 | int input2_w_stride = input2.stride(3); 46 | 47 | int input3_b_stride = input3.stride(0); 48 | int input3_c_stride = input3.stride(1); 49 | int input3_h_stride = input3.stride(2); 50 | int input3_w_stride = input3.stride(3); 51 | // printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 52 | 53 | 54 | //TODO: do we need to assert the w_stride to be 1 55 | if(input1_w_stride !=1) return error; 56 | if(input2_w_stride !=1) return error; 57 | if(input3_w_stride !=1) return error; 58 | if(input1_b_stride != output.stride(0)) return error; 59 | if(input1_c_stride != output.stride(1)) return error; 60 | 61 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 62 | 63 | 64 | error = FilterInterpolationLayer_gpu_forward_kernel( 65 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 66 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 67 | nElement,w,h,channel,batch, filter_size, 68 | 69 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 70 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 71 | input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride, 72 | 73 | 74 | input1, 75 | input2, 76 | input3, 77 | output); 78 | if (error) {AT_ERROR("CUDA call failed");} 79 | 80 | return error; 81 | 82 | } 83 | int FilterInterpolationLayer_gpu_backward( 84 | at::Tensor& input1, 85 | at::Tensor& input2, 86 | at::Tensor& input3, 87 | at::Tensor& gradoutput, 88 | at::Tensor& gradinput1, 89 | at::Tensor& gradinput2, 90 | at::Tensor& gradinput3 91 | ) 92 | { 93 | 94 | 95 | int error = 1 ; 96 | int channel = input1.size( 1); 97 | //if(channel!=3) return error; 98 | int batch = input1.size(0); 99 | if(input2.size( 0) != batch) return error; 100 | if(input2.size(1) != 2) return error; 101 | 102 | int h = input1.size(2); 103 | int w = input1.size(3); 104 | if(input2.size(2) != h) return error;// to add some checkpoint 105 | if(input2.size(3) != w) return error; 106 | 107 | 108 | int filter_size2 = input3.size( 1); 109 | int filter_size = (int) sqrt((float) filter_size2); 110 | // printf("filter size is: %d,or %f", filter_size, sqrt((float)filter_size2)); 111 | 112 | int input1_b_stride = input1.stride(0); 113 | int input1_c_stride = input1.stride(1); 114 | int input1_h_stride = input1.stride(2); 115 | int input1_w_stride = input1.stride(3); 116 | 117 | int input2_b_stride = input2.stride(0); 118 | int input2_c_stride = input2.stride(1); 119 | int input2_h_stride = input2.stride(2); 120 | int input2_w_stride = input2.stride(3); 121 | 122 | int input3_b_stride = input3.stride(0); 123 | int input3_c_stride = input3.stride(1); 124 | int input3_h_stride = input3.stride(2); 125 | int input3_w_stride = input3.stride(3); 126 | // printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 127 | 128 | 129 | //TODO: do we need to assert the w_stride to be 1 130 | if(input1_w_stride !=1) return error; 131 | if(input2_w_stride !=1) return error; 132 | if(input3_w_stride !=1) return error; 133 | if(input1_b_stride != gradinput1.stride(0)) return error; 134 | if(input2_b_stride != gradinput2.stride(0)) return error; 135 | if(input1_c_stride != gradinput1.stride(1)) return error; 136 | if(input2_c_stride != gradinput2.stride(1)) return error; 137 | if(input3_c_stride != gradinput3.stride(1)) return error; 138 | 139 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 140 | 141 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 142 | 143 | error = FilterInterpolationLayer_gpu_backward_kernel( 144 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 145 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 146 | nElement, //to let the nummous 147 | w,h,channel,batch, filter_size, 148 | 149 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 150 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 151 | input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride, 152 | 153 | input1, 154 | input2, 155 | input3, 156 | gradoutput, 157 | gradinput1, 158 | gradinput2, 159 | gradinput3 160 | ); 161 | if (error) {AT_ERROR("CUDA call failed");} 162 | 163 | return error; 164 | } 165 | 166 | 167 | 168 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 169 | m.def("FilterInterpolationLayer_gpu_forward", &FilterInterpolationLayer_gpu_forward, "FilterInterpolation forward (CUDA)"); 170 | m.def("FilterInterpolationLayer_gpu_backward", &FilterInterpolationLayer_gpu_backward, "FilterInterpolation backward (CUDA)"); 171 | } 172 | -------------------------------------------------------------------------------- /my_package/FilterInterpolation/filterinterpolation_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int FilterInterpolationLayer_gpu_forward_kernel( 8 | cudaStream_t stream, 9 | const int nElement, 10 | const int w, const int h, const int channel, const int batch, const int filter_size, 11 | 12 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 13 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 14 | const int input3_b_stride, const int input3_c_stride, const int input3_h_stride, const int input3_w_stride, 15 | 16 | at::Tensor& input1, at::Tensor& input2, at::Tensor& input3, at::Tensor& output 17 | 18 | ); 19 | 20 | int FilterInterpolationLayer_gpu_backward_kernel( 21 | cudaStream_t stream, 22 | const int nElement, 23 | const int w, const int h, const int channel, const int batch, const int filter_size, 24 | 25 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 26 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 27 | const int input3_b_stride, const int input3_c_stride, const int input3_h_stride, const int input3_w_stride, 28 | 29 | at::Tensor& input1, at::Tensor& input2, at::Tensor& input3, 30 | 31 | at::Tensor& gradoutput, at::Tensor& gradinput1, at::Tensor& gradinput2, at::Tensor& gradinput3 32 | ); 33 | -------------------------------------------------------------------------------- /my_package/FilterInterpolation/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='filterinterpolation_cuda', 12 | ext_modules=[ 13 | CUDAExtension('filterinterpolation_cuda', [ 14 | 'filterinterpolation_cuda.cc', 15 | 'filterinterpolation_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/FlowProjection/FlowProjectionLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | import flowprojection_cuda as my_lib 5 | 6 | #Please check how the STN FUNCTION is written : 7 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/gridgen.py 8 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/stn.py 9 | 10 | class FlowProjectionLayer(Function): 11 | def __init__(self,requires_grad): 12 | super(FlowProjectionLayer,self).__init__() 13 | self.requires_grad = requires_grad 14 | 15 | @staticmethod 16 | def forward(ctx, input1, requires_grad): 17 | assert(input1.is_contiguous()) 18 | # assert(input2.is_contiguous()) 19 | # self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 20 | 21 | fillhole = 1 if requires_grad == False else 0 22 | # if input1.is_cuda: 23 | # self.device = torch.cuda.current_device() 24 | # else: 25 | # self.device = -1 26 | 27 | # count = torch.zeros(input1.size(0),1,input1.size(2),input1.size(3)) # for accumulating the homography projections 28 | # output = torch.zeros(input1.size()) 29 | 30 | if input1.is_cuda : 31 | # output = output.cuda() 32 | # count = count.cuda() 33 | count = torch.cuda.FloatTensor().resize_(input1.size(0), 1, input1.size(2), input1.size(3)).zero_() 34 | output = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 35 | err = my_lib.FlowProjectionLayer_gpu_forward(input1, count,output, fillhole) 36 | else: 37 | output = torch.cuda.FloatTensor(input1.data.size()) 38 | err = my_lib.FlowProjectionLayer_cpu_forward(input1, count, output, fillhole) 39 | if err != 0: 40 | print(err) 41 | # output = output/count # to divide the counter 42 | 43 | ctx.save_for_backward(input1, count) 44 | ctx.fillhole = fillhole 45 | # self.count = count #to keep this 46 | # print(self.input1[0, 0, :10, :10]) 47 | # print(self.count[0, 0, :10, :10]) 48 | # print(self.input1[0, 0, -10:, -10:]) 49 | # print(self.count[0, 0, -10:, -10:]) 50 | 51 | # the function returns the output to its caller 52 | return output 53 | 54 | @staticmethod 55 | def backward(ctx, gradoutput): 56 | # print("Backward of Filter Interpolation Layer") 57 | # gradinput1 = input1.new().zero_() 58 | # gradinput2 = input2.new().zero_() 59 | # gradinput1 = torch.zeros(self.input1.size()) 60 | 61 | input1, count, output = ctx.saved_tensors 62 | 63 | if input1.is_cuda: 64 | # print("CUDA backward") 65 | # gradinput1 = gradinput1.cuda(self.device) 66 | gradinput1 = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 67 | err = my_lib.FlowProjectionLayer_gpu_backward(input1, count, gradoutput, gradinput1) 68 | # print(err) 69 | if err != 0 : 70 | print(err) 71 | 72 | else: 73 | # print("CPU backward") 74 | # print(gradoutput) 75 | gradinput1 = torch.FloatTensor().resize_(input1.size()).zero_() 76 | err = my_lib.FlowProjectionLayer_cpu_backward(input1, count, gradoutput, gradinput1) 77 | # print(err) 78 | if err != 0: 79 | print(err) 80 | # print(gradinput1) 81 | # print(gradinput2) 82 | 83 | # print(gradinput1) 84 | 85 | return gradinput1, None 86 | 87 | class FlowFillholelayer(Function): 88 | def __init__(self): 89 | super(FlowFillholelayer,self).__init__() 90 | 91 | def forward(self, input1): 92 | # assert(input1.is_contiguous()) 93 | # assert(input2.is_contiguous()) 94 | self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 95 | 96 | if input1.is_cuda: 97 | self.device = torch.cuda.current_device() 98 | else: 99 | self.device = -1 100 | 101 | # count = torch.zeros(input1.size(0),1,input1.size(2),input1.size(3)) # for accumulating the homography projections 102 | output = torch.zeros(input1.size()) 103 | 104 | if input1.is_cuda : 105 | output = output.cuda() 106 | # count = count.cuda() 107 | err = my_lib.FlowFillholelayer_gpu_forward(input1, output) 108 | else: 109 | # output = torch.cuda.FloatTensor(input1.data.size()) 110 | err = my_lib.FlowFillholelayer_cpu_forward(input1, output) 111 | if err != 0: 112 | print(err) 113 | # output = output/count # to divide the counter 114 | 115 | # self.count = count #to keep this 116 | # print(self.input1[0, 0, :10, :10]) 117 | # print(self.count[0, 0, :10, :10]) 118 | # print(self.input1[0, 0, -10:, -10:]) 119 | # print(self.count[0, 0, -10:, -10:]) 120 | 121 | # the function returns the output to its caller 122 | return output 123 | 124 | #TODO: if there are multiple outputs of this function, then the order should be well considered? 125 | # def backward(self, gradoutput): 126 | # # print("Backward of Filter Interpolation Layer") 127 | # # gradinput1 = input1.new().zero_() 128 | # # gradinput2 = input2.new().zero_() 129 | # gradinput1 = torch.zeros(self.input1.size()) 130 | # if self.input1.is_cuda: 131 | # # print("CUDA backward") 132 | # gradinput1 = gradinput1.cuda(self.device) 133 | # err = my_lib.FlowProjectionLayer_gpu_backward(self.input1, self.count, gradoutput, gradinput1) 134 | # # print(err) 135 | # if err != 0 : 136 | # print(err) 137 | # 138 | # else: 139 | # # print("CPU backward") 140 | # # print(gradoutput) 141 | # err = my_lib.FlowProjectionLayer_cpu_backward(self.input1, self.count, gradoutput, gradinput1) 142 | # # print(err) 143 | # if err != 0: 144 | # print(err) 145 | # # print(gradinput1) 146 | # # print(gradinput2) 147 | # 148 | # # print(gradinput1) 149 | # 150 | # return gradinput1 -------------------------------------------------------------------------------- /my_package/FlowProjection/FlowProjectionModule.py: -------------------------------------------------------------------------------- 1 | # modules/FlowProjectionModule.py 2 | from torch.nn import Module 3 | from .FlowProjectionLayer import FlowProjectionLayer #, FlowFillholeLayer 4 | 5 | class FlowProjectionModule(Module): 6 | def __init__(self, requires_grad = True): 7 | super(FlowProjectionModule, self).__init__() 8 | 9 | self.f = FlowProjectionLayer(requires_grad) 10 | 11 | def forward(self, input1): 12 | return self.f(input1) 13 | 14 | # class FlowFillholeModule(Module): 15 | # def __init__(self,hole_value = -10000.0): 16 | # super(FlowFillholeModule, self).__init__() 17 | # self.f = FlowFillholeLayer() 18 | # 19 | # def forward(self, input1): 20 | # return self.f(input1) 21 | 22 | #we actually dont need to write the backward code for a module, since we have 23 | 24 | -------------------------------------------------------------------------------- /my_package/FlowProjection/__init__.py: -------------------------------------------------------------------------------- 1 | from .FlowProjectionModule import * -------------------------------------------------------------------------------- /my_package/FlowProjection/flowprojection_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "flowprojection_cuda_kernel.cuh" 8 | 9 | int FlowProjectionLayer_gpu_forward( 10 | at::Tensor& input1, 11 | at::Tensor& count, 12 | at::Tensor& output, 13 | int fillhole 14 | ) 15 | { 16 | 17 | int error = 1 ; 18 | 19 | int channel = input1.size( 1); 20 | if(channel!= 2) return error; 21 | int batch = input1.size(0); 22 | 23 | int h = input1.size(2); 24 | int w = input1.size(3); 25 | 26 | int input1_b_stride = input1.stride(0); 27 | int input1_c_stride = input1.stride(1); 28 | int input1_h_stride = input1.stride(2); 29 | int input1_w_stride = input1.stride(3); 30 | 31 | int count_b_stride = count.stride(0); 32 | int count_c_stride = count.stride(1); 33 | int count_h_stride = count.stride(2); 34 | int count_w_stride = count.stride(3); 35 | //TODO: do we need to assert the w_stride to be 1 36 | //if(w_stride !=1) return error; 37 | if(input1_b_stride != output.stride(0)) return error; 38 | if(input1_c_stride != output.stride(1)) return error; 39 | 40 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 41 | // printf("In gpu forward\n"); 42 | error = FlowProjection_gpu_forward_kernel( 43 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 44 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 45 | nElement,w,h,channel,batch,fillhole, 46 | 47 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 48 | count_b_stride,count_c_stride,count_h_stride,count_w_stride, 49 | 50 | input1, 51 | count, 52 | output); 53 | if (error) {AT_ERROR("CUDA call failed");} 54 | 55 | return error; 56 | 57 | } 58 | 59 | int FlowProjectionLayer_gpu_backward( 60 | at::Tensor& input1, 61 | at::Tensor& count, 62 | at::Tensor& gradoutput, 63 | at::Tensor& gradinput1 64 | ) 65 | { 66 | int error = 1 ; 67 | int channel = input1.size( 1); 68 | if(channel!=2) return error; 69 | int batch = input1.size(0); 70 | if(count.size(0) != batch) return error; 71 | if(count.size(1) != 1) return error; 72 | 73 | int h = input1.size(2); 74 | int w = input1.size(3); 75 | if(count.size(2) != h) return error;// to add some checkpoint 76 | if(count.size(3) != w) return error; 77 | 78 | int input1_b_stride = input1.stride(0); 79 | int input1_c_stride = input1.stride(1); 80 | int input1_h_stride = input1.stride(2); 81 | int input1_w_stride = input1.stride(3); 82 | 83 | int count_b_stride = count.stride(0); 84 | int count_c_stride = count.stride(1); 85 | int count_h_stride = count.stride(2); 86 | int count_w_stride = count.stride(3); 87 | //TODO: do we need to assert the w_stride to be 1 88 | //if(w_stride !=1) return error; 89 | if(input1_b_stride != gradinput1.stride(0)) return error; 90 | if(input1_c_stride != gradinput1.stride(1)) return error; 91 | 92 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 93 | // printf("GPU backward: %d,%d,%d,%d\n", count_b_stride,count_c_stride,count_h_stride,count_w_stride); 94 | 95 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 96 | 97 | error = FlowProjection_gpu_backward_kernel( 98 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 99 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 100 | nElement, //to let the nummous 101 | w,h,channel,batch, 102 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 103 | count_b_stride,count_c_stride,count_h_stride,count_w_stride, 104 | 105 | input1, 106 | count, 107 | gradoutput, 108 | gradinput1 109 | ); 110 | if (error) {AT_ERROR("CUDA call failed");} 111 | 112 | return error; 113 | 114 | } 115 | 116 | 117 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 118 | m.def("FlowProjectionLayer_gpu_forward", &FlowProjectionLayer_gpu_forward, "FlowProjection forward (CUDA)"); 119 | m.def("FlowProjectionLayer_gpu_backward", &FlowProjectionLayer_gpu_backward, "FlowProjection backward (CUDA)"); 120 | } 121 | -------------------------------------------------------------------------------- /my_package/FlowProjection/flowprojection_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int FlowProjection_gpu_forward_kernel( 8 | cudaStream_t stream, const int nElement, 9 | const int w, const int h, const int channel, const int batch, const int fillhole, 10 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 11 | const int count_b_stride, const int count_c_stride, const int count_h_stride, const int count_w_stride, 12 | 13 | at::Tensor& input1, 14 | at::Tensor& count, 15 | at::Tensor& output 16 | 17 | ); 18 | 19 | int FlowProjection_gpu_backward_kernel( 20 | cudaStream_t stream, 21 | const int nElement, 22 | const int w, 23 | const int h, 24 | const int channel, 25 | const int batch, 26 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 27 | const int count_b_stride, const int count_c_stride, const int count_h_stride, const int count_w_stride, 28 | 29 | at::Tensor& input1, 30 | at::Tensor& count, 31 | at::Tensor& gradoutput, 32 | at::Tensor& gradinput1 33 | ); 34 | 35 | 36 | -------------------------------------------------------------------------------- /my_package/FlowProjection/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='flowprojection_cuda', 12 | ext_modules=[ 13 | CUDAExtension('flowprojection_cuda', [ 14 | 'flowprojection_cuda.cc', 15 | 'flowprojection_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/Interpolation/InterpolationLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | import interpolation_cuda as my_lib 5 | 6 | #Please check how the STN FUNCTION is written : 7 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/gridgen.py 8 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/stn.py 9 | 10 | class InterpolationLayer(Function): 11 | def __init__(self): 12 | super(InterpolationLayer,self).__init__() 13 | 14 | @staticmethod 15 | def forward(ctx, input1,input2): 16 | 17 | assert(input1.is_contiguous()) 18 | assert(input2.is_contiguous()) 19 | # self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 20 | # self.input2 = input2.contiguous() # TODO: Note that this is simply a shallow copy? 21 | # if input1.is_cuda: 22 | # self.device = torch.cuda.current_device() 23 | # else: 24 | # self.device = -1 25 | 26 | # output = torch.zeros(input1.size()) 27 | 28 | 29 | if input1.is_cuda : 30 | # output = output.cuda() 31 | output = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 32 | my_lib.InterpolationLayer_gpu_forward(input1, input2, output) 33 | else: 34 | output = torch.cuda.FloatTensor(input1.data.size()) 35 | my_lib.InterpolationLayer_cpu_forward(input1, input2, output) 36 | ctx.save_for_backward(input1, input2) 37 | 38 | # the function returns the output to its caller 39 | return output 40 | 41 | @staticmethod 42 | def backward(ctx, gradoutput): 43 | # print("Backward of Interpolation Layer") 44 | # gradinput1 = input1.new().zero_() 45 | # gradinput2 = input2.new().zero_() 46 | # gradinput1 = torch.zeros(self.input1.size()) 47 | # gradinput2 = torch.zeros(self.input2.size()) 48 | input1, input2 = ctx.saved_tensors 49 | 50 | if input1.is_cuda: 51 | # print("CUDA backward") 52 | # gradinput1 = gradinput1.cuda(self.device) 53 | # gradinput2 = gradinput2.cuda(self.device) 54 | gradinput1 = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 55 | gradinput2 = torch.cuda.FloatTensor().resize_(input2.size()).zero_() 56 | 57 | # the input1 image should not require any gradients 58 | # print("Does input1 requires gradients? " + str(self.input1.requires_grad)) 59 | 60 | err = my_lib.InterpolationLayer_gpu_backward(input1,input2,gradoutput,gradinput1,gradinput2) 61 | if err != 0 : 62 | print(err) 63 | else: 64 | # print("CPU backward") 65 | # print(gradoutput) 66 | gradinput1 = torch.FloatTensor().resize_(input1.size()).zero_() 67 | gradinput2 = torch.FloatTensor().resize_(input2.size()).zero_() 68 | err = my_lib.InterpolationLayer_cpu_backward(input1, input2, gradoutput, gradinput1, gradinput2) 69 | # print(err) 70 | if err != 0 : 71 | print(err) 72 | # print(gradinput1) 73 | # print(gradinput2) 74 | 75 | # print(gradinput1) 76 | 77 | return gradinput1, gradinput2 -------------------------------------------------------------------------------- /my_package/Interpolation/InterpolationModule.py: -------------------------------------------------------------------------------- 1 | # modules/InterpolationLayer.py 2 | from torch.nn import Module 3 | from .InterpolationLayer import InterpolationLayer 4 | 5 | class InterpolationModule(Module): 6 | def __init__(self): 7 | super(InterpolationModule, self).__init__() 8 | # self.f = InterpolationLayer() 9 | 10 | def forward(self, input1, input2): 11 | return InterpolationLayer.apply(input1, input2) 12 | 13 | #we actually dont need to write the backward code for a module, since we have 14 | 15 | -------------------------------------------------------------------------------- /my_package/Interpolation/__init__.py: -------------------------------------------------------------------------------- 1 | from .InterpolationModule import * -------------------------------------------------------------------------------- /my_package/Interpolation/interpolation_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "interpolation_cuda_kernel.cuh" 8 | 9 | 10 | int InterpolationLayer_gpu_forward( 11 | at::Tensor& input1, 12 | at::Tensor& input2, 13 | at::Tensor& output 14 | ) 15 | { 16 | int error = 1 ; 17 | 18 | int channel = input1.size( 1); 19 | if(channel!=3) return error; 20 | int batch = input1.size(0); 21 | if(input2.size( 0) != batch) return error; 22 | if(input2.size(1) != 2) return error; 23 | 24 | int h = input1.size(2); 25 | int w = input1.size(3); 26 | if(input2.size(2) != h) return error;// to add some checkpoint 27 | if(input2.size(3) != w) return error; 28 | 29 | int input1_b_stride = input1.stride(0); 30 | int input1_c_stride = input1.stride(1); 31 | int input1_h_stride = input1.stride(2); 32 | int input1_w_stride = input1.stride(3); 33 | 34 | int input2_b_stride = input2.stride(0); 35 | int input2_c_stride = input2.stride(1); 36 | int input2_h_stride = input2.stride(2); 37 | int input2_w_stride = input2.stride(3); 38 | //TODO: do we need to assert the w_stride to be 1 39 | //if(w_stride !=1) return error; 40 | if(input1_b_stride != output.stride(0)) return error; 41 | if(input1_c_stride != output.stride(1)) return error; 42 | 43 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 44 | 45 | error =InterpolationLayer_gpu_forward_kernel( 46 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 47 | at::cuda::getCurrentCUDAStream(), 48 | nElement,w,h,channel,batch, 49 | 50 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 51 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 52 | 53 | input1, 54 | input2, 55 | output); 56 | if (error) {AT_ERROR("CUDA call failed");} 57 | 58 | return error; 59 | 60 | } 61 | 62 | 63 | int InterpolationLayer_gpu_backward( 64 | at::Tensor& input1, 65 | at::Tensor& input2, 66 | at::Tensor& gradoutput, 67 | at::Tensor& gradinput1, 68 | at::Tensor& gradinput2 69 | ) 70 | { 71 | int error = 1 ; 72 | int channel = input1.size( 1); 73 | if(channel!=3) return error; 74 | int batch = input1.size(0); 75 | if(input2.size( 0) != batch) return error; 76 | if(input2.size(1) != 2) return error; 77 | 78 | int h = input1.size(2); 79 | int w = input1.size(3); 80 | if(input2.size(2) != h) return error;// to add some checkpoint 81 | if(input2.size(3) != w) return error; 82 | 83 | int input1_b_stride = input1.stride(0); 84 | int input1_c_stride = input1.stride(1); 85 | int input1_h_stride = input1.stride(2); 86 | int input1_w_stride = input1.stride(3); 87 | 88 | int input2_b_stride = input2.stride(0); 89 | int input2_c_stride = input2.stride(1); 90 | int input2_h_stride = input2.stride(2); 91 | int input2_w_stride = input2.stride(3); 92 | //TODO: do we need to assert the w_stride to be 1 93 | //if(w_stride !=1) return error; 94 | if(input1_b_stride != gradinput1.stride(0)) return error; 95 | if(input2_b_stride != gradinput2.stride(0)) return error; 96 | if(input1_c_stride != gradinput1.stride(1)) return error; 97 | if(input2_c_stride != gradinput2.stride(1)) return error; 98 | 99 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 100 | 101 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 102 | 103 | error = InterpolationLayer_gpu_backward_kernel( 104 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 105 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 106 | nElement, //to let the nummous 107 | w,h,channel,batch, 108 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 109 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 110 | 111 | input1, 112 | input2, 113 | gradoutput, 114 | gradinput1, 115 | gradinput2 116 | ); 117 | if (error) {AT_ERROR("CUDA call failed");} 118 | 119 | return error; 120 | 121 | } 122 | 123 | 124 | 125 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 126 | m.def("InterpolationLayer_gpu_forward", &InterpolationLayer_gpu_forward, "Interpolation forward (CUDA)"); 127 | m.def("InterpolationLayer_gpu_backward", &InterpolationLayer_gpu_backward, "Interpolation backward (CUDA)"); 128 | } 129 | -------------------------------------------------------------------------------- /my_package/Interpolation/interpolation_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int InterpolationLayer_gpu_forward_kernel( 8 | cudaStream_t stream, 9 | const int nElement, 10 | const int w, 11 | const int h, 12 | const int channel, 13 | const int batch, 14 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 15 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 16 | 17 | at::Tensor& input1, 18 | at::Tensor& input2, 19 | at::Tensor& output 20 | 21 | ); 22 | 23 | int InterpolationLayer_gpu_backward_kernel( 24 | cudaStream_t stream, 25 | const int nElement, 26 | const int w, 27 | const int h, 28 | const int channel, 29 | const int batch, 30 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 31 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 32 | 33 | at::Tensor& input1, 34 | at::Tensor& input2, 35 | at::Tensor& gradoutput, 36 | at::Tensor& gradinput1, 37 | at::Tensor& gradinput2 38 | ); 39 | -------------------------------------------------------------------------------- /my_package/Interpolation/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='interpolation_cuda', 12 | ext_modules=[ 13 | CUDAExtension('interpolation_cuda', [ 14 | 'interpolation_cuda.cc', 15 | 'interpolation_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/InterpolationCh/InterpolationChLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | import interpolationch_cuda as my_lib 5 | 6 | #Please check how the STN FUNCTION is written : 7 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/gridgen.py 8 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/stn.py 9 | 10 | class InterpolationChLayer(Function): 11 | def __init__(self,ch): 12 | super(InterpolationChLayer,self).__init__() 13 | self.ch = ch 14 | 15 | @staticmethod 16 | def forward(ctx, input1,input2): 17 | 18 | assert(input1.is_contiguous()) 19 | assert(input2.is_contiguous()) 20 | # self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 21 | # self.input2 = input2.contiguous() # TODO: Note that this is simply a shallow copy? 22 | 23 | # if input1.is_cuda: 24 | # self.device = torch.cuda.current_device() 25 | # else: 26 | # self.device = -1 27 | 28 | # output = torch.zeros(input1.size()) 29 | 30 | if input1.is_cuda : 31 | # output = output.cuda() 32 | output = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 33 | my_lib.InterpolationChLayer_gpu_forward(input1, input2, output) 34 | else: 35 | # output = torch.cuda.FloatTensor(input1.data.size()) 36 | output = torch.FloatTensor().resize_(input1.size()).zero_() 37 | my_lib.InterpolationChLayer_cpu_forward(input1, input2, output) 38 | ctx.save_for_backward(input1, input2) 39 | # the function returns the output to its caller 40 | return output 41 | 42 | @staticmethod 43 | def backward(ctx, gradoutput): 44 | # print("Backward of Interpolation Layer") 45 | # gradinput1 = input1.new().zero_() 46 | # gradinput2 = input2.new().zero_() 47 | # gradinput1 = torch.zeros(self.input1.size()) 48 | # gradinput2 = torch.zeros(self.input2.size()) 49 | 50 | input1, input2 = ctx.saved_tensors 51 | 52 | if input1.is_cuda: 53 | # print("CUDA backward") 54 | # gradinput1 = gradinput1.cuda(self.device) 55 | # gradinput2 = gradinput2.cuda(self.device) 56 | gradinput1 = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 57 | gradinput2 = torch.cuda.FloatTensor().resize_(input2.size()).zero_() 58 | # the input1 image should not require any gradients 59 | # print("Does input1 requires gradients? " + str(self.input1.requires_grad)) 60 | 61 | err = my_lib.InterpolationChLayer_gpu_backward(input1,input2,gradoutput,gradinput1,gradinput2) 62 | if err != 0 : 63 | print(err) 64 | 65 | else: 66 | # print("CPU backward") 67 | # print(gradoutput) 68 | gradinput1 = torch.FloatTensor().resize_(input1.size()).zero_() 69 | gradinput2 = torch.FloatTensor().resize_(input2.size()).zero_() 70 | 71 | err = my_lib.InterpolationChLayer_cpu_backward(input1, input2, gradoutput, gradinput1, gradinput2) 72 | # print(err) 73 | if err != 0 : 74 | print(err) 75 | # print(gradinput1) 76 | # print(gradinput2) 77 | 78 | # print(gradinput1) 79 | 80 | return gradinput1, gradinput2 -------------------------------------------------------------------------------- /my_package/InterpolationCh/InterpolationChModule.py: -------------------------------------------------------------------------------- 1 | # modules/InterpolationLayer.py 2 | from torch.nn import Module 3 | from .InterpolationChLayer import InterpolationChLayer 4 | 5 | class InterpolationChModule(Module): 6 | def __init__(self,ch): 7 | super(InterpolationChModule, self).__init__() 8 | self.ch = ch 9 | # self.f = InterpolationChLayer(ch) 10 | 11 | def forward(self, input1, input2): 12 | return InterpolationChLayer.apply(input1, input2) 13 | 14 | #we actually dont need to write the backward code for a module, since we have 15 | 16 | -------------------------------------------------------------------------------- /my_package/InterpolationCh/__init__.py: -------------------------------------------------------------------------------- 1 | from .InterpolationChModule import * 2 | -------------------------------------------------------------------------------- /my_package/InterpolationCh/interpolationch_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "interpolationch_cuda_kernel.cuh" 8 | 9 | 10 | int InterpolationChLayer_gpu_forward( 11 | at::Tensor& input1, 12 | at::Tensor& input2, 13 | at::Tensor& output 14 | ) 15 | { 16 | int error = 1 ; 17 | 18 | int channel = input1.size( 1); 19 | // if(channel!=3) return error; 20 | int batch = input1.size(0); 21 | if(input2.size( 0) != batch) return error; 22 | if(input2.size(1) != 2) return error; 23 | 24 | int h = input1.size(2); 25 | int w = input1.size(3); 26 | if(input2.size(2) != h) return error;// to add some checkpoint 27 | if(input2.size(3) != w) return error; 28 | 29 | int input1_b_stride = input1.stride(0); 30 | int input1_c_stride = input1.stride(1); 31 | int input1_h_stride = input1.stride(2); 32 | int input1_w_stride = input1.stride(3); 33 | 34 | int input2_b_stride = input2.stride(0); 35 | int input2_c_stride = input2.stride(1); 36 | int input2_h_stride = input2.stride(2); 37 | int input2_w_stride = input2.stride(3); 38 | //TODO: do we need to assert the w_stride to be 1 39 | //if(w_stride !=1) return error; 40 | if(input1_b_stride != output.stride(0)) return error; 41 | if(input1_c_stride != output.stride(1)) return error; 42 | 43 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 44 | 45 | error =InterpolationChLayer_gpu_forward_kernel( 46 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 47 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 48 | nElement,w,h,channel,batch, 49 | 50 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 51 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 52 | 53 | input1, 54 | input2, 55 | output); 56 | if (error) {AT_ERROR("CUDA call failed");} 57 | 58 | return error; 59 | 60 | } 61 | 62 | 63 | int InterpolationChLayer_gpu_backward( 64 | at::Tensor& input1, 65 | at::Tensor& input2, 66 | at::Tensor& gradoutput, 67 | at::Tensor& gradinput1, 68 | at::Tensor& gradinput2 69 | ) 70 | { 71 | int error = 1 ; 72 | int channel = input1.size( 1); 73 | // if(channel!=3) return error; 74 | int batch = input1.size(0); 75 | if(input2.size( 0) != batch) return error; 76 | if(input2.size(1) != 2) return error; 77 | 78 | int h = input1.size(2); 79 | int w = input1.size(3); 80 | if(input2.size(2) != h) return error;// to add some checkpoint 81 | if(input2.size(3) != w) return error; 82 | 83 | 84 | int input1_b_stride = input1.stride(0); 85 | int input1_c_stride = input1.stride(1); 86 | int input1_h_stride = input1.stride(2); 87 | int input1_w_stride = input1.stride(3); 88 | 89 | int input2_b_stride = input2.stride(0); 90 | int input2_c_stride = input2.stride(1); 91 | int input2_h_stride = input2.stride(2); 92 | int input2_w_stride = input2.stride(3); 93 | //TODO: do we need to assert the w_stride to be 1 94 | //if(w_stride !=1) return error; 95 | if(input1_b_stride != gradinput1.stride(0)) return error; 96 | if(input2_b_stride != gradinput2.stride(0)) return error; 97 | if(input1_c_stride != gradinput1.stride(1)) return error; 98 | if(input2_c_stride != gradinput2.stride(1)) return error; 99 | 100 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 101 | 102 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 103 | 104 | error = InterpolationChLayer_gpu_backward_kernel( 105 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 106 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 107 | nElement, //to let the nummous 108 | w,h,channel,batch, 109 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 110 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 111 | 112 | input1, 113 | input2, 114 | gradoutput, 115 | gradinput1, 116 | gradinput2 117 | ); 118 | if (error) {AT_ERROR("CUDA call failed");} 119 | 120 | return error; 121 | 122 | } 123 | 124 | 125 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 126 | m.def("InterpolationChLayer_gpu_forward", &InterpolationChLayer_gpu_forward, "InterpolationCh forward (CUDA)"); 127 | m.def("InterpolationChLayer_gpu_backward", &InterpolationChLayer_gpu_backward, "InterpolationCh backward (CUDA)"); 128 | } 129 | -------------------------------------------------------------------------------- /my_package/InterpolationCh/interpolationch_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | 8 | int InterpolationChLayer_gpu_forward_kernel( 9 | cudaStream_t stream, 10 | const int nElement, 11 | const int w, 12 | const int h, 13 | const int channel, 14 | const int batch, 15 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 16 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 17 | 18 | at::Tensor& input1, 19 | at::Tensor& input2, 20 | at::Tensor& output 21 | 22 | ); 23 | 24 | int InterpolationChLayer_gpu_backward_kernel( 25 | cudaStream_t stream, 26 | const int nElement, 27 | const int w, 28 | const int h, 29 | const int channel, 30 | const int batch, 31 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 32 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 33 | 34 | at::Tensor& input1, 35 | at::Tensor& input2, 36 | at::Tensor& gradoutput, 37 | at::Tensor& gradinput1, 38 | at::Tensor& gradinput2 39 | ); 40 | -------------------------------------------------------------------------------- /my_package/InterpolationCh/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='interpolationch_cuda', 12 | ext_modules=[ 13 | CUDAExtension('interpolationch_cuda', [ 14 | 'interpolationch_cuda.cc', 15 | 'interpolationch_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/MinDepthFlowProjection/__init__.py: -------------------------------------------------------------------------------- 1 | from .minDepthFlowProjectionModule import * 2 | -------------------------------------------------------------------------------- /my_package/MinDepthFlowProjection/minDepthFlowProjectionLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | #import _ext.my_lib as my_lib 5 | import mindepthflowprojection_cuda as my_lib 6 | 7 | class minDepthFlowProjectionLayer(Function): 8 | def __init__(self,requires_grad): 9 | super(minDepthFlowProjectionLayer,self).__init__() 10 | # self.requires_grad = requires_grad 11 | 12 | @staticmethod 13 | def forward(ctx, input1, input2, requires_grad): 14 | # print("Depth Aware Flow Projection") 15 | assert(input1.is_contiguous()) 16 | assert(input2.is_contiguous()) 17 | # self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 18 | # self.input2 = input2.contiguous() 19 | fillhole = 1 if requires_grad == False else 0 20 | # if input1.is_cuda: 21 | # self.device = torch.cuda.current_device() 22 | # else: 23 | # self.device = -1 24 | 25 | # count = torch.zeros(input1.size(0),1,input1.size(2),input1.size(3)) # for accumulating the homography projections 26 | # output = torch.zeros(input1.size()) 27 | 28 | if input1.is_cuda: 29 | # output = output.cuda() 30 | # count = count.cuda() 31 | # print("correct") 32 | count = torch.cuda.FloatTensor().resize_(input1.size(0), 1, input1.size(2), input1.size(3)).zero_() 33 | output = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 34 | err = my_lib.minDepthFlowProjectionLayer_gpu_forward(input1,input2, count,output, fillhole) 35 | else: 36 | # output = torch.cuda.FloatTensor(input1.data.size()) 37 | count = torch.FloatTensor().resize_(input1.size(0), 1, input1.size(2), input1.size(3)).zero_() 38 | output = torch.FloatTensor().resize_(input1.size()).zero_() 39 | err = my_lib.minDepthFlowProjectionLayer_cpu_forward(input1,input2, count, output,fillhole) 40 | if err != 0: 41 | print(err) 42 | # output = output/count # to divide the counter 43 | 44 | # self.count = count #to keep this 45 | # self.output = output 46 | 47 | ctx.save_for_backward(input1, input2,count,output) 48 | ctx.fillhole = fillhole 49 | 50 | # print(self.input1[0, 0, :10, :10]) 51 | # print(self.count[0, 0, :10, :10]) 52 | # print(self.input1[0, 0, -10:, -10:]) 53 | # print(self.count[0, 0, -10:, -10:]) 54 | 55 | # the function returns the output to its caller 56 | return output 57 | 58 | @staticmethod 59 | def backward(ctx, gradoutput): 60 | # print("Backward of Filter Interpolation Layer") 61 | # gradinput1 = input1.new().zero_() 62 | # gradinput2 = input2.new().zero_() 63 | # gradinput1 = torch.zeros(self.input1.size()) 64 | 65 | input1, input2, count, output = ctx.saved_tensors 66 | # fillhole = ctx.fillhole 67 | 68 | if input1.is_cuda: 69 | # print("CUDA backward") 70 | # gradinput1 = gradinput1.cuda(self.device) 71 | gradinput1 = torch.cuda.FloatTensor().resize_(input1.size()).zero_() 72 | gradinput2 = torch.cuda.FloatTensor().resize_(input2.size()).zero_() 73 | 74 | err = my_lib.minDepthFlowProjectionLayer_gpu_backward(input1,input2, 75 | count, output, 76 | gradoutput, gradinput1,gradinput2) 77 | # print(err) 78 | if err != 0 : 79 | print(err) 80 | 81 | else: 82 | # print("CPU backward") 83 | # print(gradoutput) 84 | gradinput1 = torch.FloatTensor().resize_(input1.size()).zero_() 85 | gradinput2 = torch.FloatTensor().resize_(input2.size()).zero_() 86 | err = my_lib.minDepthFlowProjectionLayer_cpu_backward(input1, input2, 87 | count, output, 88 | gradoutput, gradinput1,gradinput2) 89 | # print(err) 90 | if err != 0: 91 | print(err) 92 | # print(gradinput1) 93 | # print(gradinput2) 94 | 95 | # print(gradinput1) 96 | 97 | return gradinput1,gradinput2,None 98 | -------------------------------------------------------------------------------- /my_package/MinDepthFlowProjection/minDepthFlowProjectionModule.py: -------------------------------------------------------------------------------- 1 | # modules/FlowProjectionModule.py 2 | from torch.nn.modules.module import Module 3 | from .minDepthFlowProjectionLayer import minDepthFlowProjectionLayer #, FlowFillholeLayer 4 | 5 | __all__ =['minDepthFlowProjectionModule'] 6 | 7 | class minDepthFlowProjectionModule(Module): 8 | def __init__(self, requires_grad = True): 9 | super(minDepthFlowProjectionModule, self).__init__() 10 | self.requires_grad = requires_grad 11 | # self.f = minDepthFlowProjectionLayer(requires_grad) 12 | 13 | def forward(self, input1, input2): 14 | return minDepthFlowProjectionLayer.apply(input1, input2,self.requires_grad) 15 | 16 | # class FlowFillholeModule(Module): 17 | # def __init__(self,hole_value = -10000.0): 18 | # super(FlowFillholeModule, self).__init__() 19 | # self.f = FlowFillholeLayer() 20 | # 21 | # def forward(self, input1): 22 | # return self.f(input1) 23 | 24 | #we actually dont need to write the backward code for a module, since we have 25 | 26 | -------------------------------------------------------------------------------- /my_package/MinDepthFlowProjection/mindepthflowprojection_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "mindepthflowprojection_cuda_kernel.cuh" 8 | 9 | 10 | int minDepthFlowProjectionLayer_gpu_forward( 11 | at::Tensor& input1, 12 | at::Tensor& input2, 13 | at::Tensor& count, 14 | at::Tensor& output, 15 | int fillhole 16 | ) 17 | { 18 | 19 | int error = 1 ; 20 | 21 | int channel = input1.size( 1); 22 | if(channel!= 2) return error; 23 | int batch = input1.size(0); 24 | 25 | int h = input1.size(2); 26 | int w = input1.size(3); 27 | 28 | if(input2.size(1) !=1 ) return error; 29 | 30 | int input1_b_stride = input1.stride(0); 31 | int input1_c_stride = input1.stride(1); 32 | int input1_h_stride = input1.stride(2); 33 | int input1_w_stride = input1.stride(3); 34 | 35 | int input2_b_stride = input2.stride(0); 36 | int input2_c_stride = input2.stride(1); 37 | int input2_h_stride = input2.stride(2); 38 | int input2_w_stride = input2.stride(3); 39 | 40 | int count_b_stride = count.stride(0); 41 | int count_c_stride = count.stride(1); 42 | int count_h_stride = count.stride(2); 43 | int count_w_stride = count.stride(3); 44 | //TODO: do we need to assert the w_stride to be 1 45 | //if(w_stride !=1) return error; 46 | if(input1_b_stride != output.stride(0)) return error; 47 | if(input1_c_stride != output.stride(1)) return error; 48 | 49 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 50 | // printf("In gpu forward\n"); 51 | error = minDepthFlowProjection_gpu_forward_kernel( 52 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 53 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 54 | nElement,w,h,channel,batch,fillhole, 55 | 56 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 57 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 58 | count_b_stride,count_c_stride,count_h_stride,count_w_stride, 59 | 60 | input1, 61 | input2, 62 | count, 63 | output); 64 | if (error) {AT_ERROR("CUDA call failed");} 65 | 66 | return error; 67 | 68 | } 69 | 70 | int minDepthFlowProjectionLayer_gpu_backward( 71 | at::Tensor& input1, 72 | at::Tensor& input2, 73 | at::Tensor& count, 74 | at::Tensor& output, 75 | at::Tensor& gradoutput, 76 | at::Tensor& gradinput1, 77 | at::Tensor& gradinput2 78 | ) 79 | { 80 | int error = 1 ; 81 | int channel = input1.size( 1); 82 | if(channel!=2) return error; 83 | int batch = input1.size(0); 84 | if(count.size( 0) != batch) return error; 85 | if(count.size(1) != 1) return error; 86 | 87 | int h = input1.size(2); 88 | int w = input1.size(3); 89 | if(input2.size(1) !=1 ) return error; 90 | if(count.size(2) != h) return error;// to add some checkpoint 91 | if(count.size(3) != w) return error; 92 | 93 | int input1_b_stride = input1.stride(0); 94 | int input1_c_stride = input1.stride(1); 95 | int input1_h_stride = input1.stride(2); 96 | int input1_w_stride = input1.stride(3); 97 | 98 | int input2_b_stride = input2.stride(0); 99 | int input2_c_stride = input2.stride(1); 100 | int input2_h_stride = input2.stride(2); 101 | int input2_w_stride = input2.stride(3); 102 | 103 | int count_b_stride = count.stride(0); 104 | int count_c_stride = count.stride(1); 105 | int count_h_stride = count.stride(2); 106 | int count_w_stride = count.stride(3); 107 | //TODO: do we need to assert the w_stride to be 1 108 | //if(w_stride !=1) return error; 109 | if(input1_b_stride != gradinput1.stride(0)) return error; 110 | if(input1_c_stride != gradinput1.stride(1)) return error; 111 | 112 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 113 | // printf("GPU backward: %d,%d,%d,%d\n", count_b_stride,count_c_stride,count_h_stride,count_w_stride); 114 | 115 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 116 | 117 | error = minDepthFlowProjection_gpu_backward_kernel( 118 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 119 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 120 | nElement, //to let the nummous 121 | w,h,channel,batch, 122 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 123 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 124 | count_b_stride,count_c_stride,count_h_stride,count_w_stride, 125 | 126 | input1, 127 | input2, 128 | count, 129 | output, 130 | gradoutput, 131 | gradinput1, 132 | gradinput2 133 | ); 134 | if (error) {AT_ERROR("CUDA call failed");} 135 | //printf("Am I good in backward function %d",error); 136 | 137 | return error; 138 | 139 | } 140 | 141 | 142 | 143 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 144 | m.def("minDepthFlowProjectionLayer_gpu_forward", &minDepthFlowProjectionLayer_gpu_forward, "minDepthFlowProjection forward (CUDA)"); 145 | m.def("minDepthFlowProjectionLayer_gpu_backward", &minDepthFlowProjectionLayer_gpu_backward, "minDepthFlowProjection backward (CUDA)"); 146 | } 147 | -------------------------------------------------------------------------------- /my_package/MinDepthFlowProjection/mindepthflowprojection_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int minDepthFlowProjection_gpu_forward_kernel( 8 | cudaStream_t stream, const int nElement, 9 | const int w, const int h, const int channel, const int batch, const int fillhole, 10 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 11 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 12 | const int count_b_stride, const int count_c_stride, const int count_h_stride, const int count_w_stride, 13 | 14 | at::Tensor& input1, at::Tensor& input2, 15 | at::Tensor& count, 16 | at::Tensor& output 17 | 18 | ); 19 | 20 | int minDepthFlowProjection_gpu_backward_kernel( 21 | cudaStream_t stream, 22 | const int nElement, 23 | const int w, 24 | const int h, 25 | const int channel, 26 | const int batch, 27 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 28 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 29 | const int count_b_stride, const int count_c_stride, const int count_h_stride, const int count_w_stride, 30 | 31 | at::Tensor& input1, 32 | at::Tensor& input2, 33 | at::Tensor& count, 34 | at::Tensor& output, 35 | at::Tensor& gradoutput, 36 | at::Tensor& gradinput1, 37 | at::Tensor& gradinput2 38 | ); 39 | -------------------------------------------------------------------------------- /my_package/MinDepthFlowProjection/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='mindepthflowprojection_cuda', 12 | ext_modules=[ 13 | CUDAExtension('mindepthflowprojection_cuda', [ 14 | 'mindepthflowprojection_cuda.cc', 15 | 'mindepthflowprojection_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/SeparableConv/SeparableConvLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | import _ext.my_lib as my_lib 5 | 6 | #Please check how the STN FUNCTION is written : 7 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/gridgen.py 8 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/stn.py 9 | 10 | class SeparableConvLayer(Function): 11 | def __init__(self,filtersize): 12 | self.filtersize = filtersize 13 | super(SeparableConvLayer,self).__init__() 14 | 15 | def forward(self, input1,input2,input3): 16 | intBatches = input1.size(0) 17 | intInputDepth = input1.size(1) 18 | intInputHeight = input1.size(2) 19 | intInputWidth = input1.size(3) 20 | intFilterSize = min(input2.size(1), input3.size(1)) 21 | intOutputHeight = min(input2.size(2), input3.size(2)) 22 | intOutputWidth = min(input2.size(3), input3.size(3)) 23 | 24 | assert(intInputHeight - self.filtersize == intOutputHeight - 1) 25 | assert(intInputWidth - self.filtersize == intOutputWidth - 1) 26 | assert(intFilterSize == self.filtersize) 27 | 28 | assert(input1.is_contiguous() == True) 29 | assert(input2.is_contiguous() == True) 30 | assert(input3.is_contiguous() == True) 31 | 32 | output = input1.new().resize_(intBatches, intInputDepth, intOutputHeight, intOutputWidth).zero_() 33 | 34 | # assert(input1.is_contiguous()) 35 | # assert(input2.is_contiguous()) 36 | self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 37 | self.input2 = input2.contiguous() # TODO: Note that this is simply a shallow copy? 38 | self.input3 = input3.contiguous() 39 | if input1.is_cuda: 40 | self.device = torch.cuda.current_device() 41 | else: 42 | self.device = -1 43 | 44 | if input1.is_cuda : 45 | output = output.cuda() 46 | err = my_lib.SeparableConvLayer_gpu_forward(input1, input2,input3, output) 47 | 48 | else: 49 | # output = torch.cuda.FloatTensor(input1.data.size()) 50 | err = my_lib.SeparableConvLayer_cpu_forward(input1, input2,input3, output) 51 | if err != 0: 52 | print(err) 53 | # the function returns the output to its caller 54 | return output 55 | 56 | #TODO: if there are multiple outputs of this function, then the order should be well considered? 57 | def backward(self, gradoutput): 58 | # print("Backward of Interpolation Layer") 59 | # gradinput1 = input1.new().zero_() 60 | # gradinput2 = input2.new().zero_() 61 | gradinput1 = torch.zeros(self.input1.size()) 62 | gradinput2 = torch.zeros(self.input2.size()) 63 | gradinput3 = torch.zeros(self.input3.size()) 64 | if self.input1.is_cuda: 65 | # print("CUDA backward") 66 | gradinput1 = gradinput1.cuda(self.device) 67 | gradinput2 = gradinput2.cuda(self.device) 68 | gradinput3 = gradinput3.cuda(self.device) 69 | 70 | # the input1 image should not require any gradients 71 | # print("Does input1 requires gradients? " + str(self.input1.requires_grad)) 72 | 73 | err = my_lib.SeparableConvLayer_gpu_backward(self.input1,self.input2,self.input3, gradoutput,gradinput1,gradinput2,gradinput3) 74 | if err != 0 : 75 | print(err) 76 | 77 | else: 78 | # print("CPU backward") 79 | # print(gradoutput) 80 | err = my_lib.SeparableConvLayer_cpu_backward(self.input1, self.input2, self.input3, gradoutput, gradinput1, gradinput2, gradinput3) 81 | # print(err) 82 | if err != 0 : 83 | print(err) 84 | # print(gradinput1) 85 | # print(gradinput2) 86 | 87 | # print(gradinput1) 88 | 89 | return gradinput1, gradinput2,gradinput3 -------------------------------------------------------------------------------- /my_package/SeparableConv/SeparableConvModule.py: -------------------------------------------------------------------------------- 1 | # modules/InterpolationLayer.py 2 | from torch.nn import Module 3 | from functions.SeparableConvLayer import SeparableConvLayer 4 | 5 | class SeparableConvModule(Module): 6 | def __init__(self,filtersize): 7 | super(SeparableConvModule, self).__init__() 8 | self.f = SeparableConvLayer(filtersize) 9 | 10 | def forward(self, input1, input2, input3): 11 | return self.f(input1, input2, input3) 12 | 13 | #we actually dont need to write the backward code for a module, since we have 14 | 15 | -------------------------------------------------------------------------------- /my_package/SeparableConv/__init__.py: -------------------------------------------------------------------------------- 1 | from .SeparableConvModule import * 2 | -------------------------------------------------------------------------------- /my_package/SeparableConv/separableconv_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "separableconv_cuda_kernel.cuh" 8 | 9 | 10 | int SeparableConvLayer_gpu_forward( 11 | at::Tensor& input1, 12 | at::Tensor& input2, 13 | at::Tensor& input3, 14 | at::Tensor& output 15 | 16 | ) 17 | { 18 | int error = 1 ; 19 | 20 | int channel = input1.size( 1); 21 | if(channel!=3) return error; 22 | int batch = input1.size(0); 23 | if(input2.size( 0) != batch) return error; 24 | if(input2.size(1) != input3.size(1)) return error; //change by zhenghe, am I right? 25 | 26 | int h = input1.size(2); 27 | int w = input1.size(3); 28 | if(input2.size(2) != h - input2.size(1) + 1) return error;// to add some checkpoint 29 | if(input2.size(3) != w - input2.size(1) + 1) return error; 30 | 31 | 32 | int input1_b_stride = input1.stride(0); 33 | int input1_c_stride = input1.stride(1); 34 | int input1_h_stride = input1.stride(2); 35 | int input1_w_stride = input1.stride(3); 36 | 37 | int input2_b_stride = input2.stride(0); 38 | int input2_c_stride = input2.stride(1); 39 | int input2_h_stride = input2.stride(2); 40 | int input2_w_stride = input2.stride(3); 41 | 42 | int input3_b_stride = input3.stride(0); 43 | int input3_c_stride = input3.stride(1); 44 | int input3_h_stride = input3.stride(2); 45 | int input3_w_stride = input3.stride(3); 46 | 47 | int output_b_stride = output.stride(0); 48 | int output_c_stride = output.stride(1); 49 | int output_h_stride = output.stride(2); 50 | int output_w_stride = output.stride(3); 51 | // printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 52 | 53 | 54 | //TODO: do we need to assert the w_stride to be 1 55 | if(input1_w_stride !=1) return error; 56 | if(input2_w_stride !=1) return error; 57 | if(input3_w_stride !=1) return error; 58 | if(output_w_stride !=1) return error; 59 | 60 | if(input2_b_stride != input3_b_stride) return error; 61 | if(input2_c_stride != input3_c_stride) return error; 62 | 63 | 64 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, output); 65 | 66 | 67 | error = SeparableConvLayer_gpu_forward_kernel( 68 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 69 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 70 | nElement,w,h,channel,batch, input2.size(1), 71 | 72 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 73 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 74 | input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride, 75 | output_b_stride,output_c_stride,output_h_stride,output_w_stride, 76 | 77 | 78 | 79 | input1, 80 | input2, 81 | input3, 82 | output); 83 | if (error) {AT_ERROR("CUDA call failed");} 84 | 85 | return error; 86 | 87 | } 88 | int SeparableConvLayer_gpu_backward( 89 | at::Tensor& input1, 90 | at::Tensor& input2, 91 | at::Tensor& input3, 92 | at::Tensor& gradoutput, 93 | at::Tensor& gradinput1, 94 | at::Tensor& gradinput2, 95 | at::Tensor& gradinput3 96 | ) 97 | { 98 | 99 | 100 | int error = 1 ; 101 | int channel = input1.size( 1); 102 | if(channel!=3) return error; 103 | int batch = input1.size(0); 104 | if(input2.size( 0) != batch) return error; 105 | if(input2.size(1) != input2.size(1)) return error; 106 | 107 | int h = input1.size(2); 108 | int w = input1.size(3); 109 | if(input2.size(2) != h - input2.size(1) + 1) return error;// to add some checkpoint 110 | if(input2.size(3) != w - input2.size(1) + 1) return error; 111 | 112 | 113 | int input1_b_stride = input1.stride(0); 114 | int input1_c_stride = input1.stride(1); 115 | int input1_h_stride = input1.stride(2); 116 | int input1_w_stride = input1.stride(3); 117 | 118 | int input2_b_stride = input2.stride(0); 119 | int input2_c_stride = input2.stride(1); 120 | int input2_h_stride = input2.stride(2); 121 | int input2_w_stride = input2.stride(3); 122 | 123 | int input3_b_stride = input3.stride(0); 124 | int input3_c_stride = input3.stride(1); 125 | int input3_h_stride = input3.stride(2); 126 | int input3_w_stride = input3.stride(3); 127 | 128 | int output_b_stride = gradoutput.stride(0); 129 | int output_c_stride = gradoutput.stride(1); 130 | int output_h_stride = gradoutput.stride(2); 131 | int output_w_stride = gradoutput.stride(3); 132 | 133 | // printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 134 | 135 | 136 | //TODO: do we need to assert the w_stride to be 1 137 | if(input1_w_stride !=1) return error; 138 | if(input2_w_stride !=1) return error; 139 | if(input3_w_stride !=1) return error; 140 | if(output_w_stride !=1) return error; 141 | 142 | if(input1_b_stride != gradinput1.stride(0)) return error; 143 | if(input2_b_stride != gradinput2.stride(0)) return error; 144 | if(input1_c_stride != gradinput1.stride(1)) return error; 145 | if(input2_c_stride != gradinput2.stride(1)) return error; 146 | if(input3_c_stride != gradinput3.stride(1)) return error; 147 | 148 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 149 | 150 | int nElement = 0;//UNUSED THCudaTensor_nElement(state, gradoutput); 151 | 152 | error = SeparableConvLayer_gpu_backward_kernel( 153 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 154 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 155 | nElement, //to let the nummous 156 | w,h,channel,batch, input2.size(1), 157 | 158 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 159 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 160 | input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride, 161 | output_b_stride,output_c_stride,output_h_stride,output_w_stride, 162 | 163 | input1, 164 | input2, 165 | input3, 166 | gradoutput, 167 | gradinput1, 168 | gradinput2, 169 | gradinput3 170 | ); 171 | if (error) {AT_ERROR("CUDA call failed");} 172 | 173 | return error; 174 | } 175 | 176 | 177 | 178 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 179 | m.def("SeparableConvLayer_gpu_forward", &SeparableConvLayer_gpu_forward, "SeparableConv forward (CUDA)"); 180 | m.def("SeparableConvLayer_gpu_backward", &SeparableConvLayer_gpu_backward, "SeparableConv backward (CUDA)"); 181 | } 182 | -------------------------------------------------------------------------------- /my_package/SeparableConv/separableconv_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int SeparableConvLayer_gpu_forward_kernel( 8 | cudaStream_t stream, 9 | const int nElement, 10 | const int w, const int h, const int channel, const int batch, const int filter_size, 11 | 12 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 13 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 14 | const int input3_b_stride, const int input3_c_stride, const int input3_h_stride, const int input3_w_stride, 15 | const int output_b_stride, const int output_c_stride, const int output_h_stride, const int output_w_stride, 16 | 17 | at::Tensor& input1, at::Tensor& input2, at::Tensor& input3, at::Tensor& output 18 | 19 | ); 20 | 21 | int SeparableConvLayer_gpu_backward_kernel( 22 | cudaStream_t stream, 23 | const int nElement, 24 | const int w, const int h, const int channel, const int batch, const int filter_size, 25 | 26 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 27 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 28 | const int input3_b_stride, const int input3_c_stride, const int input3_h_stride, const int input3_w_stride, 29 | const int output_b_stride, const int output_c_stride, const int output_h_stride, const int output_w_stride, 30 | 31 | at::Tensor& input1, at::Tensor& input2, at::Tensor& input3, 32 | 33 | at::Tensor& gradoutput, at::Tensor& gradinput1, at::Tensor& gradinput2, at::Tensor& gradinput3 34 | ); 35 | -------------------------------------------------------------------------------- /my_package/SeparableConv/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='separableconv_cuda', 12 | ext_modules=[ 13 | CUDAExtension('separableconv_cuda', [ 14 | 'separableconv_cuda.cc', 15 | 'separableconv_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/SeparableConvFlow/SeparableConvFlowLayer.py: -------------------------------------------------------------------------------- 1 | # this is for wrapping the customized layer 2 | import torch 3 | from torch.autograd import Function 4 | import separableconvflow_cuda as my_lib 5 | import warnings 6 | #Please check how the STN FUNCTION is written : 7 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/gridgen.py 8 | #https://github.com/fxia22/stn.pytorch/blob/master/script/functions/stn.py 9 | 10 | class SeparableConvFlowLayer(Function): 11 | def __init__(self,filtersize): 12 | self.filtersize = filtersize 13 | warnings.warn("\nSeparable Conv Flow Layer is not precise enough for optical flow due to a divison operation") 14 | super(SeparableConvFlowLayer,self).__init__() 15 | 16 | def forward(self, input1,input2,input3): 17 | intBatches = input1.size(0) 18 | intInputDepth = input1.size(1) 19 | intInputHeight = input1.size(2) 20 | intInputWidth = input1.size(3) 21 | intFilterSize = min(input2.size(1), input3.size(1)) 22 | intOutputHeight = min(input2.size(2), input3.size(2)) 23 | intOutputWidth = min(input2.size(3), input3.size(3)) 24 | 25 | assert(intInputHeight - self.filtersize == intOutputHeight - 1) 26 | assert(intInputWidth - self.filtersize == intOutputWidth - 1) 27 | assert(intFilterSize == self.filtersize) 28 | 29 | assert(input1.is_contiguous() == True) 30 | assert(input2.is_contiguous() == True) 31 | assert(input3.is_contiguous() == True) 32 | 33 | # output = input1.new().resize_(intBatches, intInputDepth, intOutputHeight, intOutputWidth).zero_() 34 | flow_ouput = torch.zeros(intBatches, 2,intOutputHeight, intOutputWidth) # as a byproduct of SepConv, but no 35 | 36 | # assert(input1.is_contiguous()) 37 | # assert(input2.is_contiguous()) 38 | self.input1 = input1.contiguous() # need to use in the backward process, so we need to cache it 39 | self.input2 = input2.contiguous() # TODO: Note that this is simply a shallow copy? 40 | self.input3 = input3.contiguous() 41 | if input1.is_cuda: 42 | self.device = torch.cuda.current_device() 43 | else: 44 | self.device = -1 45 | 46 | if input1.is_cuda : 47 | # output = output.cuda() 48 | flow_ouput = flow_ouput.cuda() 49 | err = my_lib.SeparableConvFlowLayer_gpu_forward(input1, input2,input3,flow_ouput) 50 | 51 | else: 52 | # output = torch.cuda.FloatTensor(input1.data.size()) 53 | err = my_lib.SeparableConvFlowLayer_cpu_forward(input1, input2,input3,flow_ouput) 54 | if err != 0: 55 | print(err) 56 | # the function returns the output to its caller 57 | return flow_ouput 58 | 59 | #TODO: if there are multiple outputs of this function, then the order should be well considered? 60 | def backward(self, gradoutput): 61 | # print("Backward of Interpolation Layer") 62 | # gradinput1 = input1.new().zero_() 63 | # gradinput2 = input2.new().zero_() 64 | gradinput1 = torch.zeros(self.input1.size()) # the input1 has zero gradient because flow backprop. nothing to gradinput1 65 | gradinput2 = torch.zeros(self.input2.size()) 66 | gradinput3 = torch.zeros(self.input3.size()) 67 | if self.input1.is_cuda: 68 | # print("CUDA backward") 69 | gradinput1 = gradinput1.cuda(self.device) 70 | gradinput2 = gradinput2.cuda(self.device) 71 | gradinput3 = gradinput3.cuda(self.device) 72 | 73 | # the input1 image should not require any gradients 74 | # print("Does input1 requires gradients? " + str(self.input1.requires_grad)) 75 | 76 | # err = my_lib.SeparableConvFlowLayer_gpu_backward(self.input1,self.input2,self.input3, gradoutput,gradinput1,gradinput2,gradinput3) 77 | err = my_lib.SeparableConvFlowLayer_gpu_backward(self.input1,self.input2,self.input3, gradoutput,gradinput1,gradinput2,gradinput3) 78 | if err != 0 : 79 | print(err) 80 | 81 | else: 82 | # print("CPU backward") 83 | # print(gradoutput) 84 | # print(err) 85 | # err = my_lib.SeparableConvFlowLayer_cpu_backward(self.input1, self.input2, self.input3, gradoutput, gradinput1, gradinput2, gradinput3) 86 | err = my_lib.SeparableConvFlowLayer_cpu_backward(self.input1, self.input2, self.input3, gradoutput, gradinput1, gradinput2, gradinput3) 87 | 88 | if err != 0 : 89 | print(err) 90 | # print(gradinput1) 91 | # print(gradinput2) 92 | 93 | # print(gradinput1) 94 | 95 | return gradinput1, gradinput2,gradinput3 -------------------------------------------------------------------------------- /my_package/SeparableConvFlow/SeparableConvFlowModule.py: -------------------------------------------------------------------------------- 1 | # modules/InterpolationLayer.py 2 | from torch.nn import Module 3 | from .SeparableConvFlowLayer import SeparableConvFlowLayer 4 | import torch 5 | class SeparableConvFlowModule(Module): 6 | def __init__(self,filtersize): 7 | super(SeparableConvFlowModule, self).__init__() 8 | self.f = SeparableConvFlowLayer(filtersize) 9 | 10 | def forward(self, input1, input2, input3): 11 | # temp2 = torch.div(input2, torch.sum(input2,dim=1,keepdim=True)) 12 | return self.f(input1, input2, input3) 13 | 14 | #we actually dont need to write the backward code for a module, since we have 15 | 16 | -------------------------------------------------------------------------------- /my_package/SeparableConvFlow/__init__.py: -------------------------------------------------------------------------------- 1 | from .SeparableConvFlowModule import * 2 | -------------------------------------------------------------------------------- /my_package/SeparableConvFlow/separableconvflow_cuda.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include //works for 1.0.0 6 | 7 | #include "separableconvflow_cuda_kernel.cuh" 8 | 9 | int SeparableConvFlowLayer_gpu_forward( 10 | at::Tensor& input1, 11 | at::Tensor& input2, 12 | at::Tensor& input3, 13 | //at::Tensor& output, 14 | at::Tensor& flow_output 15 | 16 | ) 17 | { 18 | int error = 1 ; 19 | //int point =0 ;printf("debug point %d\n", point++ ); 20 | 21 | int channel = input1.size( 1); 22 | if(channel!=3) return error; 23 | int batch = input1.size(0); 24 | if(input2.size(0) != batch) return error; 25 | if(input2.size(1) != input2.size(1)) return error; 26 | //printf("debug point %d\n", point++ ); 27 | 28 | int h = input1.size(2); 29 | int w = input1.size(3); 30 | if(input2.size(2) != h - input2.size(1) + 1) return error;// to add some checkpoint 31 | if(input2.size(3) != w - input2.size(1) + 1) return error; 32 | 33 | 34 | int input1_b_stride = input1.stride(0); 35 | int input1_c_stride = input1.stride(1); 36 | int input1_h_stride = input1.stride(2); 37 | int input1_w_stride = input1.stride(3); 38 | 39 | int input2_b_stride = input2.stride(0); 40 | int input2_c_stride = input2.stride(1); 41 | int input2_h_stride = input2.stride(2); 42 | int input2_w_stride = input2.stride(3); 43 | 44 | int input3_b_stride = input3.stride(0); 45 | int input3_c_stride = input3.stride(1); 46 | int input3_h_stride = input3.stride(2); 47 | int input3_w_stride = input3.stride(3); 48 | 49 | //int output_b_stride = output.stride(0); 50 | //int output_c_stride = output.stride(1); 51 | //int output_h_stride = output.stride(2); 52 | //int output_w_stride = output.stride(3); 53 | 54 | int flow_output_b_stride = flow_output.stride(0); 55 | int flow_output_c_stride = flow_output.stride(1); 56 | int flow_output_h_stride = flow_output.stride(2); 57 | int flow_output_w_stride = flow_output.stride(3); 58 | //printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 59 | 60 | 61 | //TODO: do we need to assert the w_stride to be 1 62 | if(input1_w_stride !=1) return error; 63 | if(input2_w_stride !=1) return error; 64 | if(input3_w_stride !=1) return error; 65 | // if(output_w_stride !=1) return error; 66 | if(flow_output_w_stride !=1) return error; 67 | 68 | 69 | if(input2_b_stride != input3_b_stride) return error; 70 | if(input2_c_stride != input3_c_stride) return error; 71 | //printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 72 | 73 | 74 | int nElement = 0;//UNUSED 0;//UNUSED THCudaTensor_nElement(state, flow_output); 75 | 76 | 77 | error = SeparableConvFlowLayer_gpu_forward_kernel( 78 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 79 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 80 | nElement,w,h,channel,batch, input2.size(1), 81 | 82 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 83 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 84 | input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride, 85 | // output_b_stride,output_c_stride,output_h_stride,output_w_stride, 86 | flow_output_b_stride,flow_output_c_stride,flow_output_h_stride,flow_output_w_stride, 87 | 88 | 89 | 90 | input1, 91 | input2, 92 | input3, 93 | //output , 94 | flow_output 95 | 96 | ); 97 | if (error) {AT_ERROR("CUDA call failed");} 98 | //printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 99 | 100 | return error; 101 | 102 | } 103 | int SeparableConvFlowLayer_gpu_backward( 104 | at::Tensor& input1, 105 | at::Tensor& input2, 106 | at::Tensor& input3, 107 | at::Tensor& gradflow_output, 108 | at::Tensor& gradinput1, 109 | at::Tensor& gradinput2, 110 | at::Tensor& gradinput3 111 | ) 112 | { 113 | 114 | 115 | int error = 1 ; 116 | int channel = input1.size( 1); 117 | if(channel!=3) return error; 118 | int batch = input1.size(0); 119 | if(input2.size( 0) != batch) return error; 120 | if(input2.size(1) != input2.size(1)) return error; 121 | 122 | int h = input1.size(2); 123 | int w = input1.size(3); 124 | if(input2.size(2) != h - input2.size(1) + 1) return error;// to add some checkpoint 125 | if(input2.size(3) != w - input2.size(1) + 1) return error; 126 | 127 | 128 | int input1_b_stride = input1.stride(0); 129 | int input1_c_stride = input1.stride(1); 130 | int input1_h_stride = input1.stride(2); 131 | int input1_w_stride = input1.stride(3); 132 | 133 | int input2_b_stride = input2.stride(0); 134 | int input2_c_stride = input2.stride(1); 135 | int input2_h_stride = input2.stride(2); 136 | int input2_w_stride = input2.stride(3); 137 | 138 | int input3_b_stride = input3.stride(0); 139 | int input3_c_stride = input3.stride(1); 140 | int input3_h_stride = input3.stride(2); 141 | int input3_w_stride = input3.stride(3); 142 | 143 | //int output_b_stride = gradoutput.stride(0); 144 | //int output_c_stride = gradoutput.stride(1); 145 | //int output_h_stride = gradoutput.stride(2); 146 | //int output_w_stride = gradoutput.stride(3); 147 | 148 | int flow_output_b_stride = gradflow_output.stride(0); 149 | int flow_output_c_stride = gradflow_output.stride(1); 150 | int flow_output_h_stride = gradflow_output.stride(2); 151 | int flow_output_w_stride = gradflow_output.stride(3); 152 | 153 | // printf("filter tensor shape: %d,%d,%d,%d\n", input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride); 154 | 155 | 156 | //TODO: do we need to assert the w_stride to be 1 157 | if(input1_w_stride !=1) return error; 158 | if(input2_w_stride !=1) return error; 159 | if(input3_w_stride !=1) return error; 160 | // if(output_w_stride !=1) return error; 161 | if(flow_output_w_stride !=1) return error; 162 | 163 | if(input1_b_stride != gradinput1.stride(0)) return error; 164 | if(input2_b_stride != gradinput2.stride(0)) return error; 165 | if(input1_c_stride != gradinput1.stride(1)) return error; 166 | if(input2_c_stride != gradinput2.stride(1)) return error; 167 | if(input3_c_stride != gradinput3.stride(1)) return error; 168 | 169 | // printf("GPU backward: %d,%d,%d,%d\n", input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride); 170 | 171 | int nElement = 0;//UNUSED 0;//UNUSED THCudaTensor_nElement(state, gradflow_output); 172 | 173 | error = SeparableConvFlowLayer_gpu_backward_kernel( 174 | // at::globalContext().getCurrentCUDAStream(), //works for 0.4.1 175 | at::cuda::getCurrentCUDAStream(), //works for 1.0.0 176 | nElement, //to let the nummous 177 | w,h,channel,batch, input2.size(1), 178 | 179 | input1_b_stride,input1_c_stride,input1_h_stride,input1_w_stride, 180 | input2_b_stride,input2_c_stride,input2_h_stride,input2_w_stride, 181 | input3_b_stride,input3_c_stride,input3_h_stride,input3_w_stride, 182 | // output_b_stride,output_c_stride,output_h_stride,output_w_stride, 183 | flow_output_b_stride,flow_output_c_stride,flow_output_h_stride,flow_output_w_stride, 184 | 185 | input1, 186 | input2, 187 | input3, 188 | gradflow_output, 189 | gradinput1, 190 | gradinput2, 191 | gradinput3 192 | ); 193 | if (error) {AT_ERROR("CUDA call failed");} 194 | 195 | return error; 196 | } 197 | 198 | 199 | 200 | 201 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 202 | m.def("SeparableConvFlowLayer_gpu_forward", &SeparableConvFlowLayer_gpu_forward, "SeparableConvFlow forward (CUDA)"); 203 | m.def("SeparableConvFlowLayer_gpu_backward", &SeparableConvFlowLayer_gpu_backward, "SeparableConvFlow backward (CUDA)"); 204 | } 205 | -------------------------------------------------------------------------------- /my_package/SeparableConvFlow/separableconvflow_cuda_kernel.cuh: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | int SeparableConvFlowLayer_gpu_forward_kernel( 8 | cudaStream_t stream, 9 | const int nElement, 10 | const int w, const int h, const int channel, const int batch, const int filter_size, 11 | 12 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 13 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 14 | const int input3_b_stride, const int input3_c_stride, const int input3_h_stride, const int input3_w_stride, 15 | // const int output_b_stride, const int output_c_stride, const int output_h_stride, const int output_w_stride, 16 | const int flow_output_b_stride, const int flow_output_c_stride, const int flow_output_h_stride, const int flow_output_w_stride, 17 | 18 | at::Tensor& input1, at::Tensor& input2, at::Tensor& input3, at::Tensor& flow_output 19 | 20 | ); 21 | 22 | int SeparableConvFlowLayer_gpu_backward_kernel( 23 | cudaStream_t stream, 24 | const int nElement, 25 | const int w, const int h, const int channel, const int batch, const int filter_size, 26 | 27 | const int input1_b_stride, const int input1_c_stride, const int input1_h_stride, const int input1_w_stride, 28 | const int input2_b_stride, const int input2_c_stride, const int input2_h_stride, const int input2_w_stride, 29 | const int input3_b_stride, const int input3_c_stride, const int input3_h_stride, const int input3_w_stride, 30 | // const int output_b_stride, const int output_c_stride, const int output_h_stride, const int output_w_stride, 31 | const int flow_output_b_stride, const int flow_output_c_stride, const int flow_output_h_stride, const int flow_output_w_stride, 32 | 33 | at::Tensor& input1, at::Tensor& input2, at::Tensor& input3, 34 | 35 | at::Tensor& gradflow_output, at::Tensor& gradinput1, at::Tensor& gradinput2, at::Tensor& gradinput3 36 | ); 37 | 38 | 39 | -------------------------------------------------------------------------------- /my_package/SeparableConvFlow/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 | from compiler_args import nvcc_args, cxx_args 9 | 10 | setup( 11 | name='separableconvflow_cuda', 12 | ext_modules=[ 13 | CUDAExtension('separableconvflow_cuda', [ 14 | 'separableconvflow_cuda.cc', 15 | 'separableconvflow_cuda_kernel.cu' 16 | ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}) 17 | ], 18 | cmdclass={ 19 | 'build_ext': BuildExtension 20 | }) 21 | -------------------------------------------------------------------------------- /my_package/build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | echo "Need pytorch>=1.0.0" 4 | source activate pytorch1.0.0 5 | 6 | export PYTHONPATH=$PYTHONPATH:$(pwd) 7 | 8 | cd MinDepthFlowProjection 9 | rm -rf build *.egg-info dist 10 | python setup.py install 11 | cd .. 12 | 13 | cd FlowProjection 14 | rm -rf build *.egg-info dist 15 | python setup.py install 16 | cd .. 17 | 18 | cd SeparableConv 19 | rm -rf build *.egg-info dist 20 | python setup.py install 21 | cd .. 22 | 23 | cd InterpolationCh 24 | rm -rf build *.egg-info dist 25 | python setup.py install 26 | cd .. 27 | 28 | cd DepthFlowProjection 29 | rm -rf build *.egg-info dist 30 | python setup.py install 31 | cd .. 32 | 33 | cd Interpolation 34 | rm -rf build *.egg-info dist 35 | python setup.py install 36 | cd .. 37 | 38 | cd SeparableConvFlow 39 | rm -rf build *.egg-info dist 40 | python setup.py install 41 | cd .. 42 | 43 | cd FilterInterpolation 44 | rm -rf build *.egg-info dist 45 | python setup.py install 46 | cd .. 47 | 48 | -------------------------------------------------------------------------------- /my_package/clean.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | echo "Need pytorch>=1.0.0" 4 | source activate pytorch1.0.0 5 | 6 | cd MinDepthFlowProjection 7 | rm -rf build *.egg-info dist 8 | #python setup.py install 9 | cd .. 10 | 11 | cd FlowProjection 12 | rm -rf build *.egg-info dist 13 | #python setup.py install 14 | cd .. 15 | 16 | cd SeparableConv 17 | rm -rf build *.egg-info dist 18 | #python setup.py install 19 | cd .. 20 | 21 | cd InterpolationCh 22 | rm -rf build *.egg-info dist 23 | #python setup.py install 24 | cd .. 25 | 26 | cd DepthFlowProjection 27 | rm -rf build *.egg-info dist 28 | #python setup.py install 29 | cd .. 30 | 31 | cd Interpolation 32 | rm -rf build *.egg-info dist 33 | #python setup.py install 34 | cd .. 35 | 36 | cd SeparableConvFlow 37 | rm -rf build *.egg-info dist 38 | #python setup.py install 39 | cd .. 40 | 41 | cd FilterInterpolation 42 | rm -rf build *.egg-info dist 43 | #python setup.py install 44 | cd .. 45 | 46 | -------------------------------------------------------------------------------- /my_package/compiler_args.py: -------------------------------------------------------------------------------- 1 | # References: https://developer.nvidia.com/cuda-gpus 2 | nvcc_args = [ 3 | # Tesla: K80, K80 4 | # Quadro: (None) 5 | # NVIDIA NVS: (None) 6 | # Jetson: (None) 7 | '-gencode', 'arch=compute_37,code=sm_37', 8 | 9 | # Tesla: (None) 10 | # Quadro: K1200, K620, M1200, M520, M5000M, M4000M, M3000M, M2000M, M1000M, K620M, M600M, M500M 11 | # NVIDIA NVS: 810 12 | # GeForce / Titan: GTX 750 Ti, GTX 750, GTX 960M, GTX 950M, 940M, 930M, GTX 860M, GTX 850M, 840M, 830M 13 | # Jetson: (None) 14 | '-gencode', 'arch=compute_50,code=sm_50', 15 | 16 | # Tesla: M60, M40 17 | # Quadro: M6000 24GB, M6000, M5000, M4000, M2000, M5500M, M2200, M620 18 | # NVIDIA NVS: (None) 19 | # GeForce / Titan: GTX TITAN X, GTX 980 Ti, GTX 980, GTX 970, GTX 960, GTX 950, GTX 980, GTX 980M, GTX 970M, GTX 965M, 910M 20 | # Jetson: (None) 21 | '-gencode', 'arch=compute_52,code=sm_52', 22 | 23 | # Tesla: P100 24 | # Quadro: GP100 25 | # NVIDIA: NVS: (None) 26 | # GeForce / Titan: (None) 27 | # Jetson: (None) 28 | '-gencode', 'arch=compute_60,code=sm_60', 29 | 30 | # Tesla: P40, P4 31 | # Quadro: P6000, P5000, P4000, P2200, P2000, P1000, P620, P600, P400, P620, P520, P5200, P4200, P3200, P5000, P4000, P3000, P2000, P1000, P600, P500 32 | # NVIDIA NVS: (None) 33 | # GeForce / Titan: TITAN Xp, TITAN X, GTX 1080 Ti, GTX 1080, GTX 1070, GTX 1060, GTX 1050, GTX 1080, GTX 1070, GTX 1060 34 | # Jetson: (None) 35 | '-gencode', 'arch=compute_61,code=sm_61', 36 | 37 | # Tesla: T4 38 | # Quadro: RTX 8000, RTX 6000, RTX 5000, RTX 4000, RTX 5000, RTX 4000, RTX 3000, T2000, T1000 39 | # NVIDIA NVS: (None) 40 | # GeForce / Titan: TITAN RTX, RTX 2080 Ti, RTX 2080, RTX 2070, RTX 2060, RTX 2080, RTX 2070, RTX 2060 41 | # Jetson: (None) 42 | '-gencode', 'arch=compute_75,code=sm_75', 43 | 44 | # '-gencode', 'arch=compute_70,code=sm_70', 45 | # '-gencode', 'arch=compute_70,code=compute_70' 46 | 47 | '-w' # Ignore compiler warnings. 48 | ] 49 | 50 | cxx_args = ['-std=c++11', '-w'] -------------------------------------------------------------------------------- /networks/__init__.py: -------------------------------------------------------------------------------- 1 | from .DAIN import DAIN 2 | from .DAIN_slowmotion import DAIN_slowmotion 3 | __all__ = ( 4 | 'DAIN', 5 | 'DAIN_slowmotion' 6 | ) 7 | 8 | --------------------------------------------------------------------------------