├── .clang-format ├── .gitignore ├── .gitmodules ├── CMakeLists.txt ├── LICENSE ├── README.md ├── cmake └── helper_macros.cmake ├── configs └── config.ini ├── create_environment.sh ├── install_pytorch_precompiled.sh ├── loss ├── traced_caffe_vgg_optim.pt └── traced_lpips.pt ├── scenes ├── .gitignore └── README.md ├── shader ├── grid_render.glsl └── point_render.glsl └── src ├── CMakeLists.txt ├── apps ├── CMakeLists.txt ├── adop_viewer.h ├── colmap2adop.cpp ├── preprocess_pointcloud.cpp ├── train.cpp ├── viewer.cpp └── viewer_base.h └── lib ├── CMakeLists.txt ├── build_config.h.in ├── config.h ├── data ├── Dataset.cpp ├── Dataset.h ├── NeuralScene.cpp ├── NeuralScene.h ├── NeuralStructure.cpp ├── NeuralStructure.h ├── SceneData.cpp ├── SceneData.h ├── Settings.cpp └── Settings.h ├── git_sha1.h.in ├── models ├── Networks.h ├── NeuralCamera.cpp ├── NeuralCamera.h ├── NeuralTexture.cpp ├── NeuralTexture.h ├── Pipeline.cpp ├── Pipeline.h ├── my_adam.cpp └── my_adam.h ├── neat-utils ├── NeAT_interop.h ├── image_utils.cpp ├── image_utils.h └── neat-base-config.ini ├── opengl ├── GridGLRenderer.cpp ├── GridGLRenderer.h ├── NeuralPointCloudOpenGL.cpp ├── NeuralPointCloudOpenGL.h ├── RealTimeRenderer.cpp ├── RealTimeRenderer.h ├── SceneViewer.cpp └── SceneViewer.h └── rendering ├── AlphaListSort.cu ├── AlphaListSort.h ├── EnvironmentMap.cu ├── EnvironmentMap.h ├── NeuralPointCloud.h ├── NeuralPointCloudCuda.cpp ├── NeuralPointCloudCuda.h ├── PointBlending.h ├── PointRenderer.cu ├── PointRenderer.h ├── PointRendererHelper.h ├── RenderBackward.cu ├── RenderConstants.h ├── RenderForward.cu ├── RenderInfo.cpp ├── RenderInfo.h ├── RenderModule.cpp └── RenderModule.h /.clang-format: -------------------------------------------------------------------------------- 1 | # Commented out parameters are those with the same value as base LLVM style 2 | # We can uncomment them if we want to change their value, or enforce the 3 | # chosen value in case the base style changes (last sync: Clang 6.0.1). 4 | --- 5 | ### General config, applies to all languages ### 6 | BasedOnStyle: Google 7 | IndentWidth: 4 8 | BreakBeforeBraces: Allman 9 | ColumnLimit: 120 10 | DerivePointerAlignment: false 11 | PointerAlignment: Left 12 | MaxEmptyLinesToKeep: 3 13 | SortIncludes: true 14 | IncludeBlocks: Regroup 15 | IncludeCategories: 16 | - Regex: '^"(saiga)/' 17 | Priority: 1 18 | - Regex: '^"(internal)/' 19 | Priority: 2 20 | - Regex: '"[[:alnum:]./]+"' 21 | Priority: 3 22 | - Regex: '<[[:alnum:]./]+>' 23 | Priority: 4 24 | IndentPPDirectives: AfterHash 25 | AlignConsecutiveAssignments: true 26 | AllowShortFunctionsOnASingleLine: Inline 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | log/ 2 | valgrind* 3 | cmake-build-*/ 4 | perf* 5 | out.perf-folded 6 | .idea/ 7 | checkpoints/ 8 | imgui.ini 9 | build*/ 10 | debug/ 11 | config.ini 12 | pretrained/ 13 | experiments/ 14 | CMakeFiles/ 15 | .vs/ -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "External/tensorboard_logger"] 2 | path = External/tensorboard_logger 3 | url = ../../RustingSword/tensorboard_logger 4 | [submodule "External/NeAT"] 5 | path = External/NeAT 6 | url = ../../lfranke/NeAT.git 7 | [submodule "External/saiga"] 8 | path = External/saiga 9 | url = ../../lfranke/saiga.git 10 | 11 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.19 FATAL_ERROR) 2 | 3 | project(VET VERSION 1.0.0 LANGUAGES C CXX CUDA) 4 | set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake" ${CMAKE_MODULE_PATH}) 5 | list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/External/saiga/cmake/") 6 | 7 | include(helper_macros) 8 | include(ExternalProject) 9 | DefaultBuildType(RelWithDebInfo) 10 | message("Build Options") 11 | 12 | OptionsHelper(HEADLESS "Skips the viewer and other apps that require a window." OFF) 13 | OptionsHelper(ADOP_ASAN "Adds sanitize=address compiler flag" OFF) 14 | OptionsHelper(ADOP_TINY_EIGEN "Use saiga's tiny eigen library." ON) 15 | 16 | if (HEADLESS) 17 | set(SAIGA_MODULE_OPENGL OFF) 18 | set(SAIGA_BUILD_GLFW OFF) 19 | endif () 20 | 21 | ############# Required LIBRARIES ############### 22 | 23 | add_subdirectory(External/tensorboard_logger) 24 | set_property(TARGET tensorboard_logger PROPERTY POSITION_INDEPENDENT_CODE ON) 25 | PackageHelperTarget(tensorboard_logger TBL_FOUND) 26 | 27 | # Saiga 28 | set(SAIGA_BUILD_SAMPLES OFF) 29 | set(SAIGA_BUILD_TESTS OFF) 30 | #set(SAIGA_BUILD_SHARED ON) 31 | set(SAIGA_MODULE_VULKAN OFF) 32 | set(SAIGA_MODULE_VISION OFF) 33 | #set(SAIGA_MODULE_CUDA OFF) 34 | if (HEADLESS) 35 | set(SAIGA_MODULE_OPENGL OFF) 36 | set(SAIGA_BUILD_GLFW OFF) 37 | message("HEADLESS MODE") 38 | endif () 39 | set(SAIGA_NO_INSTALL ON) 40 | set(SAIGA_USE_SUBMODULES ON) 41 | set(SAIGA_WITH_FFMPEG OFF) 42 | set(SAIGA_WITH_FREETYPE OFF) 43 | set(SAIGA_WITH_YAMLCPP OFF) 44 | set(SAIGA_WITH_OPENAL OFF) 45 | set(SAIGA_WITH_OPENMESH OFF) 46 | set(SAIGA_WITH_OPENVR OFF) 47 | set(SAIGA_WITH_G2O OFF) 48 | set(SAIGA_WITH_CHOLMOD OFF) 49 | set(SAIGA_WITH_OPENNI OFF) 50 | set(SAIGA_WITH_K4A OFF) 51 | set(SAIGA_WITH_TINY_EIGEN ON) 52 | 53 | add_subdirectory(External/saiga) 54 | 55 | PackageHelperTarget(saiga_core SAIGA_FOUND) 56 | if (NOT HEADLESS) 57 | PackageHelperTarget(saiga_opengl SAIGA_FOUND) 58 | endif () 59 | #PackageHelperTarget(saiga_opengl SAIGA_FOUND) 60 | PackageHelperTarget(saiga_cuda SAIGA_FOUND) 61 | 62 | 63 | # Torch 64 | find_package(Torch REQUIRED) 65 | PackageHelperTarget(torch TORCH_FOUND) 66 | 67 | #replace flags 68 | set(CMAKE_EXE_LINKER_FLAGS "-Wl,--as-needed") 69 | 70 | 71 | include_directories(.) 72 | include_directories(${PACKAGE_INCLUDES}) 73 | 74 | ############# COMPILER FLAGS ############### 75 | SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=native -pthread") 76 | 77 | if (MSVC) 78 | #multiprocessor compilation for visual studio 79 | SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /MP") 80 | else () 81 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Werror=return-type") 82 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-aliasing") 83 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-sign-compare") 84 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -ftemplate-backtrace-limit=0") 85 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-psabi") 86 | endif () 87 | 88 | 89 | message(STATUS CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS}) 90 | message(STATUS CMAKE_EXE_LINKER_FLAGS ${CMAKE_EXE_LINKER_FLAGS}) 91 | 92 | set(LIBS ${LIBS} ${LIB_TARGETS} "${TORCH_LIBRARIES}") 93 | 94 | ############# C++ Standard and Filesystem stuff ############### 95 | 96 | set(CMAKE_CXX_STANDARD 17) 97 | set(CMAKE_CXX_STANDARD_REQUIRED ON) 98 | 99 | if (MSVC) 100 | #dll has all symbols, to avoid __declspecs everywhere 101 | set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON) 102 | include(GenerateExportHeader) 103 | endif () 104 | ############# SOURCE ############### 105 | 106 | add_subdirectory(src) 107 | 108 | 109 | if (CMAKE_SYSTEM_NAME MATCHES "Linux|Darwin") 110 | set(NEAT_CMAKE_PREFIX_PATH ${CMAKE_PREFIX_PATH}) 111 | string(REPLACE ";" "|" CMAKE_PREFIX_PATH_ALT_SEP "${NEAT_CMAKE_PREFIX_PATH}") 112 | if (CMAKE_GENERATOR MATCHES "Ninja") 113 | ExternalProject_Add(NeAT 114 | SOURCE_DIR ../External/NeAT 115 | PREFIX ${CMAKE_CURRENT_BINARY_DIR}/NeAT 116 | LIST_SEPARATOR | 117 | CMAKE_ARGS -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH_ALT_SEP} 118 | INSTALL_COMMAND echo "Skipping install step.") 119 | else () 120 | ExternalProject_Add(NeAT 121 | SOURCE_DIR ../External/NeAT 122 | PREFIX ${CMAKE_CURRENT_BINARY_DIR}/NeAT 123 | LIST_SEPARATOR | 124 | CMAKE_ARGS -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH_ALT_SEP} 125 | BUILD_COMMAND make -j20 126 | INSTALL_COMMAND echo "Skipping install step.") 127 | endif () 128 | endif () -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Darius Rückert and Linus Franke 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # VET: Visual Error Tomography for Point Cloud Completion and High-Quality Neural Rendering 2 | 3 |
Linus Franke, Darius Rückert, Laura Fink, Matthias Innmann, Marc Stamminger
4 | 5 | 6 | 7 | **Abstract:** In the last few years, deep neural networks opened the doors for big advances in novel view synthesis. 8 | Many of these approaches are based on a (coarse) proxy geometry obtained by structure from motion algorithms. 9 | Small deficiencies in this proxy can be fixed by neural rendering, but larger holes or missing parts, as they commonly 10 | appear for thin structures or for glossy regions, still lead to distracting artifacts and temporal instability. 11 | In this paper, we present a novel neural-rendering-based approach to detect and fix such deficiencies. 12 | As a proxy, we use a point cloud, which allows us to easily remove outlier geometry and to fill in missing geometry 13 | without complicated topological operations. 14 | Keys to our approach are (i) a differentiable, blending point-based renderer that can blend out redundant points, as 15 | well as (ii) the concept of Visual Error Tomography (VET), which allows us to lift 2D error maps to identify 3D-regions 16 | lacking geometry and to spawn novel points accordingly. 17 | Furthermore, (iii) by adding points as nested environment maps, our approach allows us to generate high-quality 18 | renderings of the surroundings in the same pipeline. 19 | In our results, we show that our approach can improve the quality of a point cloud obtained by structure from motion and 20 | thus increase novel view synthesis quality significantly. 21 | In contrast to point growing techniques, the approach can also fix large-scale holes and missing thin structures 22 | effectively. 23 | Rendering quality outperforms state-of-the-art methods and temporal stability is significantly improved, while rendering 24 | is possible at real-time frame rates. 25 | 26 | [[Project Page]](https://lfranke.github.io/vet/) [[Paper]](https://arxiv.org/abs/2311.04634) [[Youtube]](https://youtu.be/adH6GyqC4Jk) 27 | 28 | ## Citation 29 | 30 | ``` 31 | @article{franke2023vet, 32 | title={VET: Visual Error Tomography for Point Cloud Completion and High-Quality Neural Rendering}, 33 | author={Linus Franke and Darius R{\"u}ckert and Laura Fink and Matthias Innmann and Marc Stamminger}, 34 | booktitle = {ACM SIGGRAPH Asia 2023 Conference Proceedings}, 35 | year = {2023} 36 | } 37 | 38 | ``` 39 | 40 | ## Install Requirements 41 | 42 | Supported Operating Systems: Ubuntu 22.04 43 | 44 | Supported Compiler: g++-9 45 | 46 | Software Requirement: Conda (Anaconda/Miniconda) 47 | 48 | 49 | 50 | ## Install Instructions 51 | 52 | * Install Ubuntu Dependancies 53 | ``` 54 | sudo apt install git build-essential gcc-9 g++-9 55 | ``` 56 | For the viewer, also install: 57 | ``` 58 | sudo apt install xorg-dev 59 | ``` 60 | (There exists a headless mode without window management meant for training on a cluster, see below) 61 | 62 | * Clone Repo 63 | ``` 64 | git clone git@github.com:lfranke/VET.git 65 | cd VET/ 66 | git submodule update --init --recursive --jobs 0 67 | ``` 68 | 69 | * Create Conda Environment 70 | 71 | ```shell 72 | cd VET 73 | ./create_environment.sh 74 | ``` 75 | 76 | * Install Pytorch 77 | 78 | ```shell 79 | cd VET 80 | ./install_pytorch_precompiled.sh 81 | ``` 82 | 83 | * Compile VET 84 | 85 | ```shell 86 | cd VET 87 | 88 | conda activate vet 89 | 90 | export CONDA=${CONDA_PREFIX:-"$(dirname $(which conda))/../"} 91 | export CC=gcc-9 92 | export CXX=g++-9 93 | export CUDAHOSTCXX=g++-9 94 | 95 | mkdir build 96 | cd build 97 | 98 | cmake -DCMAKE_PREFIX_PATH="${CONDA}/lib/python3.9/site-packages/torch/;${CONDA}" .. 99 | 100 | make -j10 101 | 102 | ``` 103 | 104 | 105 | 106 | ## Running on pretrained models 107 | 108 | Supplemental materials link: [https://zenodo.org/records/10477744](https://zenodo.org/records/10477744) 109 | 110 | After a successful compilation, the best way to get started is to run `viewer` on the *tanks and temples* scenes 111 | using our pretrained models. 112 | First, download the scenes and extract them 113 | into `scenes/`. 114 | Now, download the model checkpoints and extract 115 | them into `experiments/`. 116 | Your folder structure should look like this: 117 | 118 | ```shell 119 | VET/ 120 | build/ 121 | ... 122 | scenes/ 123 | tt_train/ 124 | tt_playground/ 125 | ... 126 | experiments/ 127 | checkpoint_train_vet 128 | checkpoint_playground_vet 129 | ... 130 | ``` 131 | 132 | ## Viewer 133 | 134 | Start the viewer with 135 | 136 | ```shell 137 | conda activate vet 138 | export CONDA=${CONDA_PREFIX:-"$(dirname $(which conda))/../"} 139 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA/lib 140 | ./build/bin/viewer --scene_dir scenes/tt_train 141 | 142 | ``` 143 | Your working directory should be the vet root directory. 144 | 145 | 146 | ## Scene Description 147 | 148 | * VET uses ADOP's scene format. 149 | * ADOP uses a simple, text-based scene description format. 150 | * To run on your scenes you have to convert them into this format. 151 | * If you have created your scene with COLMAP (like us) you can use the colmap2adop converter. 152 | * More infos on this topic can be found here: [scenes/README.md](scenes/README.md) 153 | 154 | ## Training 155 | 156 | The pipeline is fitted to your scenes by the `train` executable. 157 | All training parameters are stored in a separate config file. 158 | The basic syntax is: 159 | 160 | ```shell 161 | ./build/bin/adop_train --config configs/config.ini 162 | ``` 163 | 164 | Make again sure that the working directory is the root. 165 | Otherwise, the loss models will not be found. 166 | 167 | ## Headless Mode 168 | 169 | If you do not want the viewer application, consider calling cmake with an additional `-DHEADLESS`. 170 | 171 | 172 | ## Troubleshooting 173 | 174 | * VET is build upon [ADOP](https://github.com/darglein/ADOP), take a look at the instructions there as well. 175 | 176 | ## License 177 | 178 | The code here is licensed under MIT, however note that some submodules are not and are compiled against. 179 | -------------------------------------------------------------------------------- /cmake/helper_macros.cmake: -------------------------------------------------------------------------------- 1 | macro(GroupSources startDir curdir) 2 | file(GLOB children RELATIVE ${startDir}/${curdir} ${startDir}/${curdir}/*) 3 | foreach(child ${children}) 4 | if(IS_DIRECTORY ${startDir}/${curdir}/${child}) 5 | GroupSources(${startDir} ${curdir}/${child}) 6 | else() 7 | string(REPLACE "/" "\\" groupname ${curdir}) 8 | source_group(${groupname} FILES ${startDir}/${curdir}/${child}) 9 | endif() 10 | endforeach() 11 | endmacro() 12 | 13 | macro(GroupSources2 startDir) 14 | file(GLOB children RELATIVE ${startDir} ${startDir}/*) 15 | foreach(child ${children}) 16 | if(IS_DIRECTORY ${startDir}/${child}) 17 | GroupSources(${startDir} ${child}) 18 | else() 19 | source_group("" FILES ${startDir}/${child}) 20 | endif() 21 | endforeach() 22 | endmacro() 23 | 24 | macro(OptionsHelper _variableName _description _defaultValue) 25 | option (${_variableName} "${_description}" "${_defaultValue}") 26 | # Create a padding string to align the console output 27 | string(LENGTH ${_variableName} SIZE) 28 | math(EXPR SIZE 20-${SIZE}) 29 | string(RANDOM LENGTH ${SIZE} ALPHABET " " padding) 30 | message(STATUS "Option ${_variableName} ${padding} ${${_variableName}}") 31 | endmacro() 32 | 33 | macro(PackageHelper _name _found _include_dir _libraries) 34 | if (${_found}) 35 | SET(LIB_MSG "Yes") 36 | SET(LIBS ${LIBS} ${_libraries}) 37 | 38 | if(NOT "${_include_dir}" STREQUAL "" ) 39 | #include_directories(${_include_dir}) 40 | SET(PACKAGE_INCLUDES ${PACKAGE_INCLUDES} ${_include_dir}) 41 | SET(LIB_MSG "Yes, at ${_include_dir}") 42 | endif() 43 | else () 44 | SET(LIB_MSG "No") 45 | endif () 46 | # Create a padding string to align the console output 47 | string(LENGTH ${_name} SIZE) 48 | math(EXPR SIZE 25-${SIZE}) 49 | string(RANDOM LENGTH ${SIZE} ALPHABET " " padding) 50 | message(STATUS "Package ${_name} ${padding} ${LIB_MSG}") 51 | endmacro() 52 | 53 | macro(PackageHelperTarget _target _found) 54 | if (TARGET ${_target}) 55 | SET(LIB_MSG "Yes") 56 | SET(LIB_TARGETS ${LIB_TARGETS} ${_target}) 57 | set(${_found} 1) 58 | get_target_property(tmp_interface_includes ${_target} INTERFACE_INCLUDE_DIRECTORIES) 59 | if(NOT "${tmp_interface_includes}" STREQUAL "" ) 60 | SET(LIB_MSG "Yes, at ${tmp_interface_includes}") 61 | endif() 62 | else () 63 | SET(LIB_MSG "No") 64 | set(${_found} 0) 65 | endif () 66 | # Create a padding string to align the console output 67 | string(LENGTH ${_target} SIZE) 68 | math(EXPR SIZE 25-${SIZE}) 69 | string(RANDOM LENGTH ${SIZE} ALPHABET " " padding) 70 | message(STATUS "Package ${_target} ${padding} ${LIB_MSG}") 71 | endmacro() 72 | 73 | #################################################################################### 74 | 75 | macro(DefaultBuildType _default_value) 76 | # Set a default build type if none was specified 77 | set(default_build_type ${_default_value}) 78 | 79 | # Create build type drop down 80 | if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) 81 | message(STATUS "Setting build type to '${default_build_type}' as none was specified.") 82 | set(CMAKE_BUILD_TYPE "${default_build_type}" CACHE 83 | STRING "Choose the type of build." FORCE) 84 | # Set the possible values of build type for cmake-gui 85 | set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS 86 | "Debug" "Release" "MinSizeRel" "RelWithDebInfo") 87 | endif() 88 | 89 | message(STATUS "\nBuild Type: ${CMAKE_BUILD_TYPE}") 90 | endmacro() 91 | -------------------------------------------------------------------------------- /configs/config.ini: -------------------------------------------------------------------------------- 1 | [TrainParams] 2 | random_seed = 3746934646 3 | do_train = true 4 | do_eval = false 5 | batch_size = 4 6 | inner_batch_size = 1 7 | inner_sample_size = 3 8 | num_epochs = 600 9 | save_checkpoints_its = 50 10 | eval_only_on_checkpoint = true 11 | name = train_vet 12 | debug = false 13 | output_file_type = .jpg 14 | checkpoint_directory = 15 | split_method = 16 | max_images = 0 17 | duplicate_train_factor = 1 18 | shuffle_initial_indices = false 19 | shuffle_train_indices = true 20 | split_remove_neighbors = 0 21 | split_index_file_train = scenes/tt_train/train.txt 22 | split_index_file_test = scenes/tt_train/test.txt 23 | train_on_eval = true 24 | train_factor = 0 25 | num_workers_train = 8 26 | num_workers_eval = 4 27 | train_crop_size = 512 28 | train_mask_border = 16 29 | reduced_check_point = false 30 | write_images_at_checkpoint = true 31 | keep_all_scenes_in_memory = true 32 | use_image_masks = false 33 | write_test_images = true 34 | texture_random_init = true 35 | texture_color_init = false 36 | train_use_crop = true 37 | experiment_dir = experiments/ 38 | scene_base_dir = scenes/ 39 | scene_names = tt_train 40 | loss_vgg = 1 41 | loss_l1 = 1 42 | loss_mse = 0 43 | min_zoom = 0.75 44 | max_zoom = 1.5 45 | crop_prefere_border = true 46 | crop_gaussian_sample = false 47 | crop_rotation = true 48 | optimize_eval_camera = false 49 | interpolate_eval_settings = false 50 | noise_pose_r = 0 51 | noise_pose_t = 0 52 | noise_intr_k = 0 53 | noise_intr_d = 0 54 | noise_point = 0 55 | lr_decay_factor = 0.8500000238 56 | lr_decay_patience = 30 57 | lock_camera_params_epochs = 100 58 | lock_structure_params_epochs = 25 59 | lock_dynamic_refinement_epochs = 50 60 | temp_image_dir = 61 | 62 | 63 | [RenderParams] 64 | render_outliers = false 65 | check_normal = false 66 | ghost_gradients = false 67 | drop_out_points_by_radius = true 68 | outlier_count = 1000000 69 | drop_out_radius_threshold = 0.6000000238 70 | dropout = 0.25 71 | depth_accept = 0.009999999776 72 | depth_accept_blend = 0.009999999776 73 | test_backward_mode = 0 74 | distortion_gradient_factor = 0.004999999888 75 | K_gradient_factor = 0.5 76 | add_depth_to_network = false 77 | no_envmap_at_points = true 78 | 79 | 80 | [PipelineParams] 81 | train = true 82 | verbose_eval = false 83 | log_render = false 84 | log_texture = false 85 | non_subzero_texture = false 86 | skip_neural_render_network = false 87 | enable_environment_map = false 88 | use_points_for_env_map = true 89 | env_num_points = 500000 90 | env_map_w = 512 91 | env_map_h = 512 92 | env_map_channels = 4 93 | env_inner_radius = 30 94 | env_radius_factor = 10 95 | env_spheres = 4 96 | env_axis = 0 97 | num_texture_channels = 4 98 | cat_env_to_color = false 99 | cat_masks_to_color = false 100 | render_modes_start_epochs = -2 -2 -2 0 101 | use_env_points_at_start = false 102 | 103 | 104 | [PointAddingParams] 105 | start_adding_points_epoch = 202 106 | point_adding_epoch_interval = 200 107 | add_points_amount_max_per_cell = 10 108 | start_removing_points_epoch = 250 109 | point_removal_epoch_interval = 100 110 | removal_confidence_cutoff = 0.3000000119 111 | scene_add_initially_random_points = 0 112 | dont_use_initial_pointcloud = false 113 | cells_worldspace_size = 0.1000000015 114 | fixed_ct_reco_path = 115 | sigmoid_narrowing_factor = 0 116 | only_use_point_growing = false 117 | neat_use_as_subprocess_ct_reco = true 118 | full_path_to_neat_executable = build/NeAT/src/NeAT-build/bin/reconstruct 119 | neat_loss_folder_name = l1_loss_grey 120 | neat_scene_scale = 1 121 | neat_zmin = 0 122 | neat_tv_loss = 9.999999747e-05 123 | debug_max_list_length = -1 124 | push_point_confidences_down = 9.999999975e-07 125 | 126 | 127 | [OptimizerParams] 128 | network_checkpoint_directory = 129 | texture_optimizer = adam 130 | fix_render_network = false 131 | fix_texture = false 132 | fix_environment_map = false 133 | fix_points = false 134 | fix_poses = false 135 | fix_intrinsics = true 136 | fix_dynamic_refinement = true 137 | fix_vignette = false 138 | fix_response = false 139 | fix_wb = false 140 | fix_exposure = false 141 | fix_motion_blur = true 142 | fix_rolling_shutter = true 143 | lr_render_network = 0.0002 144 | lr_texture = 0.01 145 | lr_background_color = 0.004 146 | lr_dynamic_refinement = 0.005 147 | lr_environment_map = 0.02 148 | lr_confidence = 0.001 149 | lr_environment_map_density = 0.02 150 | lr_points = 0.0001 151 | lr_poses = 0.005 152 | lr_intrinsics = 0.01 153 | response_smoothness = 1 154 | lr_vignette = 1e-05 155 | lr_response = 0.001 156 | lr_wb = 0.0005 157 | lr_exposure = 0.0005 158 | lr_motion_blur = 0.005 159 | lr_rolling_shutter = 2e-06 160 | 161 | 162 | [NeuralCameraParams] 163 | enable_vignette = true 164 | enable_exposure = true 165 | enable_response = true 166 | enable_white_balance = true 167 | enable_motion_blur = false 168 | enable_rolling_shutter = false 169 | response_params = 25 170 | response_gamma = 0.4545454681 171 | response_leak_factor = 0.009999999776 172 | 173 | 174 | [MultiScaleUnet2dParams] 175 | num_input_layers = 4 176 | num_input_channels = 4 177 | num_output_channels = 3 178 | feature_factor = 4 179 | num_layers = 4 180 | add_input_to_filters = false 181 | channels_last = false 182 | half_float = false 183 | upsample_mode = bilinear 184 | norm_layer_down = id 185 | norm_layer_up = id 186 | last_act = id 187 | conv_block = gated 188 | conv_block_up = gated 189 | pooling = average 190 | activation = elu 191 | network_version = MultiScaleUnet2d 192 | filters_network = 4 8 16 32 64 193 | sh_bands = 4 194 | -------------------------------------------------------------------------------- /create_environment.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #git submodule update --init --recursive --jobs 0 4 | 5 | 6 | source $(conda info --base)/etc/profile.d/conda.sh 7 | 8 | conda update -n base -c defaults conda 9 | 10 | conda create -y -n vet python=3.9.7 11 | conda activate vet 12 | 13 | conda install -y ncurses=6.3 -c conda-forge 14 | conda install -y cuda -c nvidia/label/cuda-11.8.0 15 | conda install -y configargparse=1.4 astunparse=1.6.3 numpy=1.21.2 ninja=1.10.2 pyyaml mkl=2022.0.1 mkl-include=2022.0.1 setuptools=58.0.4 cffi=1.15.0 typing_extensions=4.1.1 future=0.18.2 six=1.16.0 requests=2.27.1 dataclasses=0.8 tqdm 16 | conda install -y -c conda-forge jpeg=9e 17 | conda install -y -c conda-forge cmake=3.26.1 coin-or-cbc=2.10.5 glog=0.5.0 gflags=2.2.2 protobuf=3.13.0.1 freeimage tensorboard=2.8.0 18 | 19 | conda install -y -c conda-forge cudnn=8.9.2 20 | 21 | #./install_pytorch_precompiled.sh 22 | -------------------------------------------------------------------------------- /install_pytorch_precompiled.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | source $(conda info --base)/etc/profile.d/conda.sh 3 | 4 | conda activate vet 5 | 6 | CONDA=${CONDA_PREFIX:-"$(dirname $(which conda))/../"} 7 | 8 | mkdir External/ 9 | cd External/ 10 | 11 | 12 | wget https://download.pytorch.org/libtorch/cu116/libtorch-cxx11-abi-shared-with-deps-1.13.1%2Bcu116.zip -O libtorch.zip 13 | unzip libtorch.zip -d . 14 | 15 | 16 | cp -rv libtorch/ $CONDA/lib/python3.9/site-packages/torch/ 17 | 18 | -------------------------------------------------------------------------------- /loss/traced_caffe_vgg_optim.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lfranke/VET/644343f1c73b1d6ad08a1473c0e874d001145381/loss/traced_caffe_vgg_optim.pt -------------------------------------------------------------------------------- /loss/traced_lpips.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lfranke/VET/644343f1c73b1d6ad08a1473c0e874d001145381/loss/traced_lpips.pt -------------------------------------------------------------------------------- /scenes/.gitignore: -------------------------------------------------------------------------------- 1 | */ 2 | -------------------------------------------------------------------------------- /scenes/README.md: -------------------------------------------------------------------------------- 1 | ## Custom Datasets 2 | 3 | 4 | 5 | ### COLMAP Scenes 6 | 7 | All of our sample scenes have been generated by COLMAP. 8 | To run on your COLMAP reconstruction use the `colmap2adop` executable. 9 | 10 | Usage Example: 11 | 12 | ```shell 13 | export SCENE_BASE=/local.home/Projects/ADOP/additional_material/colmap/Playground/ 14 | 15 | build/bin/colmap2adop --sparse_dir SCENE_BASE/sparse/0/ \ 16 | --image_dir scenes/tt_playground/images/ \ 17 | --point_cloud_file SCENE_BASE/dense_point_cloud.ply \ 18 | --output_path scenes/playground_test \ 19 | --scale_intrinsics 1 --render_scale 1 20 | ``` 21 | 22 | ### render_scale 23 | 24 | The `render_scale` (see command above) dictates how large the rendered image is relative to the ground truth images. 25 | For example, if the input images are 4000x3000 and the render_scale is 0.5 then the `viewer` will render the images in 2000x1500. 26 | This of course impacts efficiency and memory consumption. 27 | 28 | The `render_scale` can also be updated manually by modifying the `dataset.ini` in the corresponding scene folder. 29 | 30 | We recommend to set `render_scale` **before training**, because training and inference should be done at the same scale. 31 | 32 | ### Point Cloud Augmentation 33 | 34 | ADOP relies on the point cloud to store texture and geometry information. 35 | If the point cloud is too sparse compared to the image resolution, holes might appear or texture detail is lost. 36 | A simple trick to improve the render quality is therefore to just duplicate the points and add a small offset to the new points. 37 | This can be done by running the `preprocess_pointcloud` executable. 38 | Usage: 39 | 40 | ```shell 41 | # Double the number of points of the tt_train scene 42 | ./build/bin/preprocess_pointcloud --scene_path scenes/tt_train --point_factor 2 43 | ``` 44 | 45 | The tanks and temples scenes that we have included in the supplementary material have been processed with the following settings: 46 | 47 | | Scene | point_factor | 48 | | ------------- | ------------- | 49 | | tt_playground | 2 | 50 | | tt_lighthouse | 4 | 51 | | tt_m60 | 2 | 52 | | tt_train | 3 | 53 | 54 | ### Other Scenes 55 | 56 | To train on your scene you must convert it to the following format. If you have reconstructed the scene with COLMAP you can use the COLMAP2ADOP converter below. 57 | 58 | ``` 59 | + scenes/my_dataset/ 60 | - dataset.ini [required] 61 | Copy this from a sample scene and update the pathes. 62 | - Camera files [required] 63 | For each camera one ini file. Checkout sample scenes. 64 | If all images are captured by the same camera, only one camera file is required! 65 | camera0.ini 66 | camera1.ini 67 | ... 68 | - point_cloud.ply [required] 69 | positions, normals 70 | - images.txt [required] 71 | 0000.jpg 72 | 0001.jpg 73 | ... 74 | - camera_indices.txt [required] 75 | For each image the corresponding camera. Example: 76 | 0 77 | 0 78 | 1 79 | 0 80 | ... 81 | - masks.txt [optional] 82 | 0000.png 83 | 0001.png 84 | ... 85 | - poses.txt [required] 86 | Camera -> World Transformation 87 | qx qy qz qw tx ty tz 88 | - exposure.txt [optional] 89 | Initialization of the Exposure Value (EV). For example from the jpg EXIF data. 90 | 13.1 91 | 14 92 | 10.5 93 | ... 94 | ``` 95 | -------------------------------------------------------------------------------- /shader/grid_render.glsl: -------------------------------------------------------------------------------- 1 | ##GL_VERTEX_SHADER 2 | #version 430 3 | 4 | layout(location=0) in vec4 in_position; 5 | layout(location=1) in vec4 in_color; 6 | layout(location=2) in int in_index; 7 | layout(location=3) in vec3 in_normal; 8 | layout(location=4) in vec4 in_data; 9 | 10 | 11 | layout(location=0) uniform mat4 model; 12 | layout(location=1) uniform mat4 view; 13 | layout(location=2) uniform mat4 proj; 14 | 15 | 16 | layout(location=3) uniform mat4 cv_view; 17 | layout(location=4) uniform mat3 cv_proj; 18 | layout(location=5) uniform vec2 cv_viewport; 19 | layout(location=6) uniform float cv_scale; 20 | layout(location=7) uniform vec4 cv_dis1; 21 | layout(location=8) uniform vec4 cv_dis2; 22 | layout(location = 12) uniform int render_cam_model = 0; 23 | 24 | #include "vision/distortion.glsl" 25 | 26 | 27 | 28 | 29 | out vec4 v_color; 30 | out vec4 v_data; 31 | flat out int v_index; 32 | out float discard_float; 33 | 34 | #define M_PI 3.14159265358979323846 35 | vec4 toPano(vec3 p) 36 | { 37 | float phi = atan(normalize(p).z, normalize(p).x) / (M_PI); 38 | float chi = -(acos(normalize(p).y) - M_PI / 2) / 2; 39 | return vec4(phi, chi, length(p)/100, 1); 40 | } 41 | 42 | void main() { 43 | v_color = in_color; 44 | v_data = in_data; 45 | v_index = in_index; 46 | vec4 world_p = vec4(in_position.xyz, 1); 47 | vec4 world_n = vec4(in_normal.xyz, 0); 48 | discard_float=0; 49 | // world_p.xyz = world_p.xzy; 50 | // // world_p.x*=-1; 51 | // world_p.y*=-1; 52 | // // world_p.z*=-1; 53 | 54 | 55 | vec3 view_p = vec3(cv_view * model * world_p); 56 | vec3 view_n = normalize(vec3(cv_view * model * world_n)); 57 | 58 | float z = view_p.z; 59 | 60 | if (z <= 0 && render_cam_model == 0 ) 61 | { 62 | gl_Position = vec4(-100, -100, -100, 1); 63 | discard_float = 100; 64 | return; 65 | } 66 | 67 | if (dot(normalize(view_p), view_n) > 0 && render_cam_model == 0) 68 | { 69 | // gl_Position = vec4(0, 0, -100, 0); 70 | // return; 71 | } 72 | 73 | vec2 image_p = view_p.xy / z; 74 | image_p = distortNormalizedPoint(image_p, cv_dis1, cv_dis2); 75 | image_p = vec2(cv_proj * vec3(image_p, 1)); 76 | 77 | float znear = 0.1; 78 | float zfar = 500; 79 | float f = zfar; 80 | float n = znear; 81 | float d = z * (f + n) / (f - n) - (2 * f * n) / (f - n); 82 | float d2 = z; 83 | 84 | // remove viewport transform 85 | image_p = 2 * image_p / cv_viewport - vec2(1); 86 | image_p.y *= -1; 87 | 88 | gl_Position = vec4(image_p * d2, d, d2); 89 | vec2 norm_p = gl_Position.xy/gl_Position.w; 90 | if(norm_p.x < -1 || norm_p.y < -1 || norm_p.x > 1 || norm_p.y > 1) 91 | discard_float = 100; 92 | 93 | if (render_cam_model == 2) 94 | gl_Position = toPano(vec3(cv_view * model * world_p)); 95 | // if (render_cam_model == 2) 96 | // gl_Position = toPano(vec3(cv_view * model * world_p)); 97 | 98 | } 99 | 100 | 101 | ##GL_FRAGMENT_SHADER 102 | #version 430 103 | 104 | #include "include/saiga/colorize.h" 105 | in vec4 v_color; 106 | in vec4 v_data; 107 | flat in int v_index; 108 | in float discard_float; 109 | layout(location=9) uniform float color_scale = 1; 110 | layout(location=10) uniform int display_mode = 0; 111 | layout(location=11) uniform float cutoff_val = 0; 112 | 113 | 114 | layout(location=0) out vec4 out_color; 115 | layout(location=1) out int out_index; 116 | 117 | void main() { 118 | if(discard_float>0.1) 119 | { 120 | discard; 121 | return; 122 | } 123 | vec4 color = vec4(1,0,0,0); 124 | if(display_mode == 0) {color = v_color;} 125 | 126 | if(display_mode >= 1) { 127 | // float v = v_data[render_mode-1] / max_value; 128 | // color = vec4(colorizeFusion(v), 1); 129 | color = v_data; 130 | if(v_data.x 0 && render_cam_model == 0) 71 | { 72 | // gl_Position = vec4(0, 0, -100, 0); 73 | // return; 74 | } 75 | 76 | vec2 image_p = view_p.xy / z; 77 | image_p = distortNormalizedPoint(image_p, cv_dis1, cv_dis2); 78 | image_p = vec2(cv_proj * vec3(image_p, 1)); 79 | 80 | float znear = 0.1; 81 | float zfar = 500; 82 | float f = zfar; 83 | float n = znear; 84 | float d = z * (f + n) / (f - n) - (2 * f * n) / (f - n); 85 | float d2 = z; 86 | 87 | // remove viewport transform 88 | image_p = 2 * image_p / cv_viewport - vec2(1); 89 | image_p.y *= -1; 90 | 91 | gl_Position = vec4(image_p * d2, d, d2); 92 | if (render_cam_model == 2) 93 | gl_Position = toPano(vec3(cv_view * model * world_p)); 94 | 95 | 96 | gl_PointSize = 3; 97 | } 98 | 99 | 100 | ##GL_FRAGMENT_SHADER 101 | #version 430 102 | 103 | #include "include/saiga/colorize.h" 104 | in vec4 v_color; 105 | in vec4 v_data; 106 | flat in int v_index; 107 | layout(location=9) uniform float color_scale = 1; 108 | layout(location=10) uniform int render_mode = 0; 109 | layout(location=11) uniform float max_value = 1; 110 | 111 | 112 | layout(location=0) out vec4 out_color; 113 | layout(location=1) out int out_index; 114 | 115 | void main() { 116 | vec4 color = vec4(1,0,0,0); 117 | if(render_mode == 0) {color = v_color;} 118 | 119 | if(render_mode >= 1) { 120 | float v = v_data[render_mode-1] / max_value; 121 | color = vec4(colorizeFusion(v), 1); 122 | } 123 | 124 | 125 | out_color = color_scale * color; 126 | out_index = v_index; 127 | // out_color = vec4(0,0,1,1); 128 | } 129 | 130 | 131 | -------------------------------------------------------------------------------- /src/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(lib) 2 | add_subdirectory(apps) 3 | -------------------------------------------------------------------------------- /src/apps/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | function(add_app TARGET_NAME ADD_DLL_POSTBUILD) 2 | 3 | set_target_properties(${TARGET_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/bin") 4 | target_include_directories(${TARGET_NAME} PUBLIC ".") 5 | target_link_libraries(${TARGET_NAME} NeuralPoints) 6 | 7 | if (PR_NO_WINDOW) 8 | target_compile_definitions(${TARGET_NAME} PUBLIC PR_NO_WINDOW) 9 | endif () 10 | add_definitions(-DEXECUTABLE_DIR="${PROJECT_SOURCE_DIR}") 11 | 12 | # ---------------------------------------------------------- 13 | # dll copy 14 | if (WIN32) 15 | if (ADD_DLL_POSTBUILD) 16 | add_custom_command(TARGET ${TARGET_NAME} POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy_if_different "${PROJECT_BINARY_DIR}/src/lib/$/NeuralPoints.dll" $) 17 | 18 | set(copyDest "${PROJECT_BINARY_DIR}/bin/$/") 19 | 20 | function(SANITIZE_DLL_PATHS_AND_COPY PATHS_DLLS PATHS_OUTPUT) 21 | #Sanitizes paths to remove backslashes 22 | # STRING(REGEX REPLACE "/" "\\\\" copyInput \"${PATHS_DLLS}\") 23 | # STRING(REGEX REPLACE "/" "\\\\" copyDestination \"${PATHS_OUTPUT}\") 24 | message("dll post build copy: copying ${copyInput} to ${copyDestination}.") 25 | 26 | add_custom_command(TARGET ${TARGET_NAME} POST_BUILD COMMAND COMMAND copy ${copyInput} ${copyDestination}) 27 | endfunction() 28 | 29 | SANITIZE_DLL_PATHS_AND_COPY("${CONDA_P_PATH}/Lib/site-packages/torch/lib/*.dll" ${copyDest}) 30 | SANITIZE_DLL_PATHS_AND_COPY("${CONDA_P_PATH}/Library/bin/win64/Release/*.dll" ${copyDest}) 31 | SANITIZE_DLL_PATHS_AND_COPY("${CONDA_P_PATH}/Library/bin/uv.dll" ${copyDest}) 32 | SANITIZE_DLL_PATHS_AND_COPY("${CONDA_P_PATH}/Library/bin/nvToolsExt64_1.dll" ${copyDest}) 33 | SANITIZE_DLL_PATHS_AND_COPY("${CONDA_P_PATH}/Library/bin/libiomp5md.dll" ${copyDest}) 34 | SANITIZE_DLL_PATHS_AND_COPY("${PROJECT_BINARY_DIR}/External/saiga/bin/$/*.dll" ${copyDest}) 35 | SANITIZE_DLL_PATHS_AND_COPY("${PROJECT_BINARY_DIR}/External/torchvision/$/*.dll" ${copyDest}) 36 | endif () 37 | endif () 38 | 39 | message(STATUS "App enabled: ${TARGET_NAME}") 40 | endfunction() 41 | 42 | add_executable(colmap2adop colmap2adop.cpp) 43 | add_app(colmap2adop FALSE) 44 | 45 | 46 | add_executable(preprocess_pointcloud preprocess_pointcloud.cpp) 47 | add_app(preprocess_pointcloud FALSE) 48 | 49 | 50 | add_executable(train train.cpp "../../External/NeAT/src/utils/cimg_wrapper.cpp") 51 | add_dependencies(train NeAT) 52 | add_app(train TRUE) 53 | 54 | if (TARGET saiga_opengl) 55 | add_executable(viewer viewer.cpp adop_viewer.h) 56 | add_app(viewer TRUE) 57 | 58 | endif () 59 | -------------------------------------------------------------------------------- /src/apps/adop_viewer.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #include "saiga/core/geometry/cameraAnimation.h" 10 | #include "saiga/core/glfw/all.h" 11 | #include "saiga/opengl/assets/AssetRenderSystem.h" 12 | #include "saiga/opengl/assets/all.h" 13 | #include "saiga/opengl/rendering/deferredRendering/deferredRendering.h" 14 | #include "saiga/opengl/rendering/forwardRendering/forwardRendering.h" 15 | #include "saiga/opengl/rendering/renderer.h" 16 | #include "saiga/opengl/window/WindowTemplate.h" 17 | #include "saiga/opengl/window/glfw_window.h" 18 | #include "saiga/opengl/world/LineSoup.h" 19 | #include "saiga/opengl/world/TextureDisplay.h" 20 | #include "saiga/opengl/world/pointCloud.h" 21 | 22 | #include "config.h" 23 | #include "opengl/GridGLRenderer.h" 24 | #include "opengl/RealTimeRenderer.h" 25 | #include "opengl/SceneViewer.h" 26 | 27 | #include "viewer_base.h" 28 | using namespace Saiga; 29 | 30 | 31 | using WindowType = glfw_Window; 32 | constexpr WindowManagement wm = WindowManagement::GLFW; 33 | 34 | 35 | class ADOPViewer : public StandaloneWindow, 36 | public glfw_KeyListener, 37 | glfw_MouseListener, 38 | ViewerBase 39 | { 40 | public: 41 | ADOPViewer(std::string scene_dir, std::unique_ptr renderer_, std::unique_ptr window_); 42 | ~ADOPViewer() {} 43 | 44 | 45 | void LoadSceneImpl(); 46 | 47 | 48 | void update(float dt) override 49 | { 50 | if (next_scene != current_scene) 51 | { 52 | LoadSceneImpl(); 53 | } 54 | 55 | renderer->tone_mapper.params_dirty = true; 56 | 57 | 58 | scene->scene_camera.update(dt); 59 | } 60 | 61 | void interpolate(float dt, float interpolation) override 62 | { 63 | if (renderer->use_mouse_input_in_3dview || renderer->use_mouse_input_in_3dview || mouse_in_gt) 64 | { 65 | scene->scene_camera.interpolate(dt, interpolation); 66 | } 67 | mouse_in_gt = false; 68 | } 69 | 70 | void render(RenderInfo render_info) override; 71 | bool mouse_in_gt = false; 72 | 73 | 74 | void Recording(ImageInfo& fd); 75 | 76 | virtual void keyPressed(int key, int scancode, int mods) override 77 | { 78 | switch (key) 79 | { 80 | case GLFW_KEY_Q: 81 | { 82 | auto& f = scene->scene->frames[scene->selected_capture]; 83 | renderer->tone_mapper.params.exposure_value = f.exposure_value; 84 | camera->setModelMatrix(f.OpenglModel()); 85 | camera->updateFromModel(); 86 | break; 87 | } 88 | case GLFW_KEY_ESCAPE: 89 | { 90 | window->close(); 91 | break; 92 | } 93 | case GLFW_KEY_T: 94 | { 95 | camera->rotateGlobal(camera->getDirection().head<3>(), 5.f); 96 | break; 97 | } 98 | case GLFW_KEY_G: 99 | { 100 | camera->rotateGlobal(camera->getDirection().head<3>(), -5.f); 101 | break; 102 | } 103 | default: 104 | break; 105 | }; 106 | } 107 | 108 | 109 | void mousePressed(int key, int x, int y) override 110 | { 111 | ivec2 global_pixel = ivec2(x, y); 112 | ivec2 local_pixel = renderer->WindowCoordinatesToViewport(global_pixel); 113 | 114 | switch (key) 115 | { 116 | case GLFW_MOUSE_BUTTON_RIGHT: 117 | { 118 | auto ray = ::camera->PixelRay(local_pixel.cast(), renderer->viewport_size.x(), 119 | renderer->viewport_size.y(), true); 120 | scene->Select(ray); 121 | 122 | break; 123 | } 124 | } 125 | } 126 | float render_scale = 1.f; 127 | 128 | bool debug_directionality = false; 129 | vec3 debug_refl_dir = vec3(0, 0, 1); 130 | // used to rotate all results left -1, not 0, right +1 131 | int rotate_result_90deg = 0; 132 | bool render_grid = false; 133 | float grid_cutoff_val = 0.05f; 134 | bool grid_cutoff_as_percent = true; 135 | int grid_mode = 0; 136 | 137 | 138 | private: 139 | std::shared_ptr spline_mesh; 140 | SplinePath camera_spline; 141 | 142 | ViewMode view_mode; 143 | std::shared_ptr sun; 144 | TextureDisplay display; 145 | 146 | std::shared_ptr grid_renderer; 147 | 148 | std::unique_ptr target_framebuffer; 149 | 150 | // video recording variables 151 | std::string recording_dir; 152 | }; 153 | -------------------------------------------------------------------------------- /src/apps/colmap2adop.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #include "saiga/core/util/commandLineArguments.h" 8 | #include "saiga/core/util/exif/TinyEXIF.h" 9 | #include "saiga/core/util/file.h" 10 | #include "saiga/vision/cameraModel/OCam.h" 11 | #include "saiga/vision/util/ColmapReader.h" 12 | 13 | #include "data/SceneData.h" 14 | 15 | 16 | std::vector ExposureValuesFromImages(std::vector files, std::string image_dir) 17 | { 18 | std::vector exposures; 19 | 20 | for (auto f : files) 21 | { 22 | auto path = image_dir + "/" + f; 23 | if (!std::filesystem::exists(path)) 24 | { 25 | SAIGA_EXIT_ERROR("Could not find image file " + path); 26 | } 27 | 28 | auto data = File::loadFileBinary(path); 29 | TinyEXIF::EXIFInfo info; 30 | info.parseFrom((unsigned char*)data.data(), data.size()); 31 | 32 | if (info.FNumber == 0 || info.ExposureTime == 0 || info.ISOSpeedRatings == 0) 33 | { 34 | std::cout << "No EXIF exposure value found for image " << f << std::endl; 35 | exposures.push_back(0); 36 | } 37 | else 38 | { 39 | double EV_log2 = log2((info.FNumber * info.FNumber) / info.ExposureTime) + 40 | log2(info.ISOSpeedRatings / 100.0) - info.ExposureBiasValue; 41 | exposures.push_back(EV_log2); 42 | } 43 | } 44 | 45 | std::cout << "EV Statistic:" << std::endl; 46 | Statistics stat(exposures); 47 | std::cout << stat << std::endl; 48 | std::cout << "dynamic range: " << exp2(stat.max - stat.min) << std::endl; 49 | 50 | return exposures; 51 | } 52 | 53 | static std::shared_ptr ColmapScene(std::string sparse_dir, std::string image_dir, 54 | std::string point_cloud_file, std::string output_scene_path, 55 | double scale_intrinsics = 1., double render_scale = 1.) 56 | { 57 | std::cout << "Preprocessing Colmap scene " << sparse_dir << " -> " << output_scene_path << std::endl; 58 | std::filesystem::create_directories(output_scene_path); 59 | 60 | 61 | 62 | ColmapReader reader(sparse_dir); 63 | 64 | std::map col_cam_to_id; 65 | 66 | std::vector camera_files; 67 | std::vector image_files; 68 | 69 | for (auto s : reader.images) 70 | { 71 | image_files.push_back(s.name); 72 | } 73 | 74 | std::vector exposures = ExposureValuesFromImages(image_files, image_dir); 75 | 76 | { 77 | SceneCameraParams params; 78 | for (int i = 0; i < reader.cameras.size(); ++i) 79 | { 80 | auto c = reader.cameras[i]; 81 | 82 | params.w = c.w * scale_intrinsics; 83 | params.h = c.h * scale_intrinsics; 84 | 85 | params.K = c.K.scale(scale_intrinsics).cast(); 86 | params.distortion = c.dis.cast(); 87 | 88 | col_cam_to_id[c.camera_id] = i; 89 | 90 | auto f = "camera" + std::to_string(i) + ".ini"; 91 | 92 | std::filesystem::remove(output_scene_path + "/" + f); 93 | params.Save(output_scene_path + "/" + f); 94 | 95 | camera_files.push_back(f); 96 | } 97 | } 98 | 99 | { 100 | std::filesystem::remove(output_scene_path + "/dataset.ini"); 101 | SceneDatasetParams params; 102 | params.file_model = ""; 103 | params.image_dir = image_dir; 104 | params.camera_files = camera_files; 105 | params.scene_exposure_value = Statistics(exposures).mean; 106 | params.render_scale = render_scale; 107 | 108 | params.scene_up_vector = vec3(0, -1, 0); 109 | params.Save(output_scene_path + "/dataset.ini"); 110 | } 111 | 112 | 113 | { 114 | std::ofstream ostream1(output_scene_path + "/images.txt"); 115 | std::ofstream ostream2(output_scene_path + "/camera_indices.txt"); 116 | 117 | 118 | for (auto s : reader.images) 119 | { 120 | ostream1 << s.name << std::endl; 121 | ostream2 << col_cam_to_id[s.camera_id] << std::endl; 122 | } 123 | } 124 | 125 | { 126 | auto pc_in = point_cloud_file; 127 | auto pc_out = output_scene_path + "/point_cloud.ply"; 128 | 129 | std::filesystem::remove(pc_out); 130 | std::filesystem::remove(output_scene_path + "/point_cloud.bin"); 131 | 132 | SAIGA_ASSERT(std::filesystem::exists(pc_in)); 133 | SAIGA_ASSERT(std::filesystem::is_regular_file(pc_in)); 134 | 135 | std::string command = "cp -v -n " + pc_in + " " + pc_out; 136 | auto res = system(command.c_str()); 137 | if (res != 0) 138 | { 139 | SAIGA_EXIT_ERROR("Copy failed!"); 140 | } 141 | } 142 | 143 | std::vector poses; 144 | for (int i = 0; i < reader.images.size(); ++i) 145 | { 146 | SE3 view(reader.images[i].q, reader.images[i].t); 147 | poses.push_back(view.inverse()); 148 | } 149 | SceneData::SavePoses(poses, output_scene_path + "/poses.txt"); 150 | 151 | 152 | std::shared_ptr sd = std::make_shared(output_scene_path); 153 | { 154 | SAIGA_ASSERT(sd->frames.size() == reader.images.size()); 155 | 156 | for (int i = 0; i < reader.images.size(); ++i) 157 | { 158 | sd->frames[i].exposure_value = exposures[i]; 159 | } 160 | } 161 | sd->Save(); 162 | 163 | return sd; 164 | } 165 | 166 | 167 | 168 | int main(int argc, char* argv[]) 169 | { 170 | std::string sparse_dir; 171 | std::string image_dir; 172 | std::string point_cloud_file; 173 | std::string output_path; 174 | double scale_intrinsics = 1; 175 | double render_scale = 1; 176 | 177 | CLI::App app{"COLMAP to ADOP Scene Converter", "colmap2adop"}; 178 | app.add_option("--sparse_dir", sparse_dir)->required(); 179 | app.add_option("--image_dir", image_dir)->required(); 180 | app.add_option("--point_cloud_file", point_cloud_file)->required(); 181 | app.add_option("--output_path", output_path)->required(); 182 | app.add_option("--scale_intrinsics", scale_intrinsics)->required(); 183 | app.add_option("--render_scale", render_scale)->required(); 184 | 185 | 186 | CLI11_PARSE(app, argc, argv); 187 | 188 | 189 | ColmapScene(sparse_dir, image_dir, point_cloud_file, output_path, scale_intrinsics, render_scale); 190 | 191 | 192 | return 0; 193 | } 194 | -------------------------------------------------------------------------------- /src/apps/preprocess_pointcloud.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #include "saiga/core/util/commandLineArguments.h" 8 | 9 | #include "data/SceneData.h" 10 | 11 | void PreprocessPointcloud(std::string scene_dir, int dup_factor) 12 | { 13 | auto scene = std::make_shared(scene_dir); 14 | 15 | scene->DuplicatePoints(dup_factor, 1); 16 | scene->ComputeRadius(); 17 | scene->RemoveClosePoints(0.00005); 18 | scene->RemoveLonelyPoints(5, 0.02); 19 | scene->point_cloud.RemoveDoubles(0.0002); 20 | std::cout << "Remaining Points: " << scene->point_cloud.NumVertices() << std::endl; 21 | 22 | scene->point_cloud.ReorderMorton64(); 23 | scene->point_cloud.RandomBlockShuffle(default_point_block_size); 24 | scene->ComputeRadius(); 25 | scene->Save(); 26 | 27 | 28 | UnifiedModel(scene->point_cloud).Save(scene_dir + "/point_cloud_dup.ply"); 29 | } 30 | 31 | int main(int argc, char* argv[]) 32 | { 33 | CLI::App app{"Point Cloud Preprocessor", "preprocess_pointcloud"}; 34 | 35 | std::string scene_path; 36 | int point_factor = 1; 37 | app.add_option("--scene_path", scene_path)->required(); 38 | app.add_option("--point_factor", point_factor, 39 | "Increase number of point by this factor. Set to 1 to keep the number of points.") 40 | ->required(); 41 | 42 | 43 | CLI11_PARSE(app, argc, argv); 44 | 45 | PreprocessPointcloud(scene_path, point_factor); 46 | 47 | return 0; 48 | } 49 | -------------------------------------------------------------------------------- /src/apps/viewer_base.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | 8 | #include "config.h" 9 | #include "opengl/RealTimeRenderer.h" 10 | #include "opengl/SceneViewer.h" 11 | 12 | #include 13 | using namespace Saiga; 14 | enum class ViewMode 15 | { 16 | // Only the model 17 | MODEL, 18 | // Only neural render 19 | NEURAL, 20 | // 4x4 Split with neural, model, and point cloud 21 | SPLIT_NEURAL, 22 | // Only Debug 23 | DEBUG_ONLY, 24 | // GT only 25 | GT_ONLY, 26 | }; 27 | 28 | 29 | class ViewerBase 30 | { 31 | public: 32 | ViewerBase() 33 | { 34 | { 35 | // ignore performance warning when downloading auto exposure data 36 | std::vector ids; 37 | ids.push_back(131186); 38 | Error::ignoreGLError(ids); 39 | } 40 | 41 | // Define GUI layout 42 | auto editor_layout = std::make_unique(); 43 | editor_layout->RegisterImguiWindow("Video Recording", EditorLayoutLSplit2x2::WINDOW_POSITION_LEFT); 44 | editor_layout->RegisterImguiWindow("Extra", EditorLayoutLSplit2x2::WINDOW_POSITION_LEFT); 45 | editor_layout->RegisterImguiWindow("Model Viewer", EditorLayoutLSplit2x2::WINDOW_POSITION_LEFT); 46 | editor_layout->RegisterImguiWindow("Capture View", EditorLayoutLSplit2x2::WINDOW_POSITION_LEFT_BOTTOM); 47 | editor_layout->RegisterImguiWindow("Neural View", EditorLayoutLSplit2x2::WINDOW_POSITION_MAIN_12); 48 | editor_layout->RegisterImguiWindow("Debug View", EditorLayoutLSplit2x2::WINDOW_POSITION_MAIN_21); 49 | editor_layout->RegisterImguiWindow("Closest Ground Truth", EditorLayoutLSplit2x2::WINDOW_POSITION_MAIN_22); 50 | editor_layout->RegisterImguiWindow("Neural Renderer", EditorLayoutLSplit2x2::WINDOW_POSITION_LEFT_BOTTOM); 51 | editor_gui.SetLayout(std::move(editor_layout)); 52 | } 53 | 54 | void LoadScene(std::string scene_dir) 55 | { 56 | neural_renderer = nullptr; 57 | scene = std::make_shared(std::make_shared(scene_dir)); 58 | current_scene = next_scene; 59 | object_tex = {}; 60 | object_col = {}; 61 | } 62 | 63 | 64 | void imgui() 65 | { 66 | ImGui::Begin("Model Viewer"); 67 | 68 | 69 | ImGui::Checkbox("renderPoints", &render_points); 70 | ImGui::Checkbox("renderWireframe", &renderWireframe); 71 | ImGui::Checkbox("renderObject", &renderObject); 72 | ImGui::Checkbox("renderTexture", &renderTexture); 73 | 74 | ImGui::Checkbox("renderAABB", &render_scene_aabb); 75 | ImGui::Checkbox("renderCoordinateSystem", &render_coordinate_system); 76 | 77 | 78 | if (ImGui::Button("reset objects")) 79 | { 80 | object_tex = {}; 81 | object_col = {}; 82 | } 83 | 84 | if (ImGui::Button("RemoveInvalidPoints")) 85 | { 86 | scene->RemoveInvalidPoints(); 87 | } 88 | 89 | 90 | 91 | scene->imgui(); 92 | 93 | ImGui::End(); 94 | 95 | if (neural_renderer) 96 | { 97 | neural_renderer->imgui(); 98 | } 99 | } 100 | 101 | protected: 102 | int current_scene = -1; 103 | int next_scene = 0; 104 | std::shared_ptr scene; 105 | 106 | std::shared_ptr object_tex; 107 | std::shared_ptr object_col; 108 | 109 | std::unique_ptr neural_renderer; 110 | 111 | bool renderObject = true; 112 | bool renderTexture = true; 113 | bool renderWireframe = false; 114 | bool render_points = true; 115 | bool render_debug = true; 116 | 117 | bool render_scene_aabb = true; 118 | bool render_coordinate_system = true; 119 | }; 120 | -------------------------------------------------------------------------------- /src/lib/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(TARGET_NAME "NeuralPoints") 2 | 3 | 4 | ## Generate source file with git commit 5 | include(GetGitRevisionDescription) 6 | get_git_head_revision(GIT_REFSPEC MY_GIT_SHA1 ALLOW_LOOKING_ABOVE_CMAKE_SOURCE_DIR) 7 | set(MY_BUILDCONFIG_DIR "${CMAKE_CURRENT_BINARY_DIR}/include") 8 | 9 | configure_file( 10 | "${CMAKE_CURRENT_SOURCE_DIR}/git_sha1.h.in" 11 | "${MY_BUILDCONFIG_DIR}/git_sha1.h" 12 | @ONLY 13 | ) 14 | 15 | configure_file( 16 | "${CMAKE_CURRENT_SOURCE_DIR}/build_config.h.in" 17 | "${MY_BUILDCONFIG_DIR}/build_config.h" 18 | @ONLY 19 | ) 20 | 21 | #cmake_policy(SET CMP0104 OLD) 22 | add_library(${TARGET_NAME} SHARED "") 23 | if (MSVC) 24 | generate_export_header(${TARGET_NAME}) 25 | endif () 26 | target_include_directories(${TARGET_NAME} PUBLIC "." ${MY_BUILDCONFIG_DIR}) 27 | target_sources(${TARGET_NAME} PRIVATE "${MY_BUILDCONFIG_DIR}/build_config.h" "${MY_BUILDCONFIG_DIR}/git_sha1.h") 28 | target_link_libraries(${TARGET_NAME} ${LIBS}) 29 | 30 | if (ADOP_ASAN) 31 | target_compile_options( 32 | ${TARGET_NAME} PUBLIC 33 | $<$:-fsanitize=address> 34 | ) 35 | target_link_options(${TARGET_NAME} PUBLIC -fsanitize=address) 36 | endif () 37 | 38 | list(APPEND MY_CUDA_FLAGS "--expt-relaxed-constexpr") 39 | list(APPEND MY_CUDA_FLAGS "-lineinfo") 40 | #list(APPEND MY_CUDA_FLAGS "-Xcudafe --diag_suppress=20236") 41 | 42 | #list(APPEND MY_CUDA_FLAGS "-rdynamic") 43 | #list(APPEND MY_CUDA_FLAGS "-G") 44 | 45 | 46 | target_compile_options(${TARGET_NAME} PUBLIC $<$:${MY_CUDA_FLAGS}>) 47 | #set_property(TARGET ${TARGET_NAME} PROPERTY CUDA_ARCHITECTURES 70-virtual) 48 | #set_property(TARGET ${TARGET_NAME} PROPERTY CUDA_ARCHITECTURES native) 49 | 50 | 51 | ############# source ############# 52 | 53 | 54 | # Source files in "src/" to level directory 55 | FILE(GLOB_RECURSE DATA_SRC data/*.cpp) 56 | FILE(GLOB_RECURSE MODELS_SRC models/*.cpp) 57 | if (TARGET saiga_opengl) 58 | FILE(GLOB_RECURSE OPENGL_SRC opengl/*.cpp) 59 | endif () 60 | FILE(GLOB_RECURSE RENDER_SRC rendering/*.cpp) 61 | FILE(GLOB_RECURSE NEATUTILS_SRC neat-utils/*.cpp) 62 | FILE(GLOB_RECURSE RENDER_CUDA_SRC rendering/*.cu) 63 | FILE(GLOB_RECURSE PROG_HEADER config.h *.h ${SAIGA_BUILDCONFIG_FILE}) 64 | target_sources(${TARGET_NAME} PRIVATE ${DATA_SRC} ${MODELS_SRC} ${OPENGL_SRC} ${RENDER_SRC} ${RENDER_CUDA_SRC} ${NEATUTILS_SRC} PUBLIC ${PROG_HEADER}) 65 | 66 | 67 | set_target_properties(${TARGET_NAME} PROPERTIES CXX_STANDARD 17) 68 | set_target_properties(${TARGET_NAME} PROPERTIES ARCHIVE_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/bin") 69 | set_target_properties(${TARGET_NAME} PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/bin") 70 | 71 | if (PR_NO_WINDOW) 72 | target_compile_definitions(${TARGET_NAME} PUBLIC PR_NO_WINDOW) 73 | endif () 74 | 75 | 76 | if (WIN32) 77 | target_compile_definitions(${TARGET_NAME} PRIVATE -DBUILD_NEURAL_POINTS_DLL) 78 | target_compile_definitions(${TARGET_NAME} PRIVATE -DBUILD_SHARED_LIBS) 79 | 80 | 81 | #torch_cpu and torch_cuda interface compiler flags which need to be "escaped" with -Xcompiler 82 | function(CUDA_CONVERT_FLAGS EXISTING_TARGET) 83 | get_property(old_flags TARGET ${EXISTING_TARGET} PROPERTY INTERFACE_COMPILE_OPTIONS) 84 | message("replacing ${EXISTING_TARGET} cuda flags:") 85 | message(${old_flags}) 86 | if (NOT "${old_flags}" STREQUAL "") 87 | string(REPLACE ";" "," CUDA_flags "${old_flags}") 88 | set_property(TARGET ${EXISTING_TARGET} PROPERTY INTERFACE_COMPILE_OPTIONS 89 | "$<$>:${old_flags}>$<$>:-Xcompiler=${CUDA_flags}>" 90 | ) 91 | endif () 92 | endfunction() 93 | 94 | CUDA_CONVERT_FLAGS(torch_cpu) 95 | CUDA_CONVERT_FLAGS(torch_cuda) 96 | 97 | endif () -------------------------------------------------------------------------------- /src/lib/build_config.h.in: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #include "saiga/core/util/FileSystem.h" 10 | 11 | #define PROJECT_DIR std::filesystem::path("@PROJECT_SOURCE_DIR@") 12 | -------------------------------------------------------------------------------- /src/lib/config.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #include "saiga/core/geometry/kdtree.h" 10 | #include "saiga/core/imgui/imgui.h" 11 | #include "saiga/core/math/CoordinateSystems.h" 12 | #include "saiga/core/math/random.h" 13 | #include "saiga/core/model/model_from_shape.h" 14 | #include "saiga/core/util/FileSystem.h" 15 | #include "saiga/core/util/directory.h" 16 | #include "saiga/core/util/ini/ini.h" 17 | #include "saiga/vision/VisionTypes.h" 18 | 19 | #include "build_config.h" 20 | 21 | #ifdef _WIN32 22 | # ifdef BUILD_NEURAL_POINTS_DLL 23 | # define _NP_API __declspec(dllexport) 24 | # define _NP_INLINE_API inline 25 | # define _NP_EXT_API 26 | # else 27 | # define _NP_API __declspec(dllimport) 28 | # define _NP_INLINE_API 29 | # define _NP_EXT_API extern 30 | # endif 31 | 32 | #else 33 | // UNIX, ignore _API 34 | # define _NP_API 35 | # define _NP_EXT_API 36 | # define _NP_INLINE_API inline 37 | #endif 38 | 39 | 40 | namespace Saiga 41 | { 42 | class Camera; 43 | } 44 | _NP_INLINE_API _NP_API Saiga::Camera* camera; 45 | 46 | using Saiga::mat3; 47 | using Saiga::mat4; 48 | using Saiga::ucvec3; 49 | using Saiga::vec3; 50 | using Saiga::vec4; 51 | -------------------------------------------------------------------------------- /src/lib/data/Dataset.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | #include "saiga/core/Core.h" 8 | #include "saiga/core/util/FileSystem.h" 9 | #include "saiga/core/util/directory.h" 10 | #include "saiga/core/util/tostring.h" 11 | #include "saiga/vision/torch/ImageTensor.h" 12 | 13 | #include "SceneData.h" 14 | #include "config.h" 15 | #include "rendering/RenderInfo.h" 16 | 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include 22 | #include 23 | 24 | TemplatedImage InitialUVImage(int h, int w); 25 | 26 | 27 | using NeuralTrainData = std::shared_ptr; 28 | 29 | class SceneDataTrainSampler 30 | { 31 | public: 32 | SceneDataTrainSampler() {} 33 | SceneDataTrainSampler(std::shared_ptr dataset, std::vector indices, bool down_scale, 34 | ivec2 crop_size, int inner_batch_size, bool use_image_mask, bool crop_rotate); 35 | 36 | std::vector Get(int index); 37 | 38 | int Size() const { return indices.size(); } 39 | 40 | public: 41 | int inner_batch_size = 1; 42 | 43 | // different for each camera 44 | std::vector image_size_crop; 45 | std::vector image_size_input; 46 | std::vector> uv_target; 47 | 48 | std::shared_ptr scene; 49 | std::vector indices; 50 | 51 | int num_classes = -1; 52 | 53 | bool down_scale = false; 54 | bool crop_rotate = false; 55 | bool random_translation = true; 56 | bool sample_gaussian = false; 57 | bool random_zoom = true; 58 | bool prefere_border = true; 59 | bool use_image_mask = false; 60 | int inner_sample_size = 1; 61 | 62 | vec2 min_max_zoom = vec2(0.75, 1.5); 63 | }; 64 | 65 | namespace torch 66 | { 67 | class MultiDatasetSampler : public torch::data::samplers::Sampler<> 68 | { 69 | public: 70 | MultiDatasetSampler(std::vector sizes, int batch_size, bool shuffle); 71 | 72 | void reset(torch::optional new_size = torch::nullopt) override {} 73 | 74 | /// Returns the next batch of indices. 75 | optional> next(size_t batch_size) override; 76 | 77 | /// Serializes the `RandomSampler` to the `archive`. 78 | void save(serialize::OutputArchive& archive) const override {} 79 | 80 | /// Deserializes the `RandomSampler` from the `archive`. 81 | void load(serialize::InputArchive& archive) override {} 82 | 83 | /// Returns the current index of the `RandomSampler`. 84 | size_t index() const noexcept { return current_index; } 85 | 86 | int Size() { return batch_offset_size.size(); } 87 | int NumImages() { return combined_indices.size(); } 88 | 89 | size_t current_index = 0; 90 | std::vector> combined_indices; 91 | 92 | // a pointer into the array above 93 | std::vector> batch_offset_size; 94 | 95 | 96 | std::vector sizes; 97 | int batch_size; 98 | }; 99 | } // namespace torch 100 | 101 | 102 | class TorchSingleSceneDataset : public torch::data::Dataset 103 | { 104 | public: 105 | TorchSingleSceneDataset(std::vector sampler); 106 | virtual torch::optional size() const override { return sampler.front().Size(); } 107 | virtual std::vector get2(size_t index); 108 | 109 | 110 | virtual NeuralTrainData get(size_t index) override { return {}; } 111 | 112 | std::vector get_batch(torch::ArrayRef indices) override; 113 | 114 | private: 115 | std::vector sampler; 116 | }; 117 | -------------------------------------------------------------------------------- /src/lib/data/NeuralScene.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | 8 | #include "saiga/core/Core.h" 9 | 10 | #include "SceneData.h" 11 | #include "Settings.h" 12 | #include "config.h" 13 | #include "data/NeuralStructure.h" 14 | #include "models/NeuralCamera.h" 15 | #include "models/NeuralTexture.h" 16 | #include "rendering/EnvironmentMap.h" 17 | #include "rendering/NeuralPointCloudCuda.h" 18 | 19 | 20 | using namespace Saiga; 21 | 22 | class NeuralScene 23 | { 24 | public: 25 | NeuralScene(std::shared_ptr scene, std::shared_ptr params, bool eval_only = false); 26 | 27 | void BuildOutlierCloud(int n); 28 | 29 | void Train(int epoch_id, bool train); 30 | 31 | void to(torch::Device device) 32 | { 33 | if (environment_map) 34 | { 35 | environment_map->to(device); 36 | } 37 | texture->to(device); 38 | 39 | camera->to(device); 40 | intrinsics->to(device); 41 | poses->to(device); 42 | point_cloud_cuda->to(device); 43 | if (outlier_point_cloud_cuda) 44 | { 45 | outlier_point_cloud_cuda->to(device); 46 | } 47 | } 48 | 49 | void SaveCheckpoint(const std::string& dir, bool reduced); 50 | 51 | void LoadCheckpoint(const std::string& dir); 52 | 53 | void CreateTextureOptimizer(); 54 | 55 | void CreateStructureOptimizer(); 56 | 57 | void ShrinkTextureOptimizer(torch::Tensor indices_to_keep); 58 | 59 | void AppendToTextureOptimizer(int new_size); 60 | 61 | void Log(const std::string& log_dir); 62 | 63 | void OptimizerStep(int epoch_id, bool structure_only); 64 | 65 | void OptimizerClear(int epoch_id, bool structure_only); 66 | 67 | void UpdateLearningRate(int epoch_id, double factor); 68 | 69 | // Download + Save in 'scene' 70 | void DownloadIntrinsics(); 71 | 72 | void DownloadPoses(); 73 | 74 | void AddPointsViaPointGrowing(int factor = 2, float distance = 1.f, bool update_optimizer = true); 75 | 76 | void AddNewPoints(std::vector positions, std::vector normal); 77 | 78 | void AddNewPoints(torch::Tensor random_values, bool update_optimizer = true); 79 | void AddNewPoints(std::vector& positions, bool update_optimizer = true); 80 | 81 | void AddNewRandomForEnvSphere(int num_spheres, float inner_radius, float env_radius_factor, int num_points, 82 | bool update_optimizer = true); 83 | 84 | void AddNewRandomPoints(float factor); 85 | 86 | void AddNewRandomPointsInValuefilledBB(int num_points_to_add, float percent_of_boxes = 0.05); 87 | 88 | void AddNewRandomPointsFromCTStack(int num_points_to_add, std::string path, float ct_volume_scale = 5.f, 89 | vec3 ct_volume_translation = vec3(0, 0, 1)); 90 | 91 | void AddNewRandomPointsFromCTHdr(torch::Tensor hdr_img_stack, int max_num_points_to_add, float ct_volume_scale, 92 | vec3 ct_volume_translation, AABB aabb); 93 | 94 | void RemovePoints(torch::Tensor indices, bool update_optimizer = true); 95 | 96 | public: 97 | friend class NeuralPipeline; 98 | 99 | std::shared_ptr scene; 100 | 101 | NeuralPointCloudCuda point_cloud_cuda = nullptr; 102 | NeuralPointCloudCuda outlier_point_cloud_cuda = nullptr; 103 | 104 | NeuralPointTexture texture = nullptr; 105 | 106 | EnvironmentMap environment_map = nullptr; 107 | NeuralCamera camera = nullptr; 108 | PoseModule poses = nullptr; 109 | IntrinsicsModule intrinsics = nullptr; 110 | 111 | 112 | //[batch, n_points, 3] 113 | torch::Tensor dynamic_refinement_t; 114 | 115 | std::shared_ptr camera_adam_optimizer, camera_sgd_optimizer; 116 | std::shared_ptr texture_optimizer; 117 | std::shared_ptr structure_optimizer; 118 | 119 | torch::DeviceType device = torch::kCUDA; 120 | std::shared_ptr params; 121 | }; 122 | -------------------------------------------------------------------------------- /src/lib/data/NeuralStructure.cpp: -------------------------------------------------------------------------------- 1 | #include "NeuralStructure.h" 2 | 3 | #include "rendering/PointRenderer.h" 4 | 5 | 6 | PoseModuleImpl::PoseModuleImpl(std::shared_ptr scene) 7 | { 8 | SAIGA_ASSERT(!scene->frames.empty()); 9 | 10 | tangent_poses = torch::zeros({(long)scene->frames.size(), 6L}, 11 | torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCUDA)) 12 | .clone(); 13 | tangent_poses = tangent_poses.set_requires_grad(true); 14 | 15 | std::vector pose_data; 16 | for (auto& f : scene->frames) 17 | { 18 | pose_data.push_back(f.pose.inverse()); 19 | } 20 | static_assert(sizeof(Sophus::SE3d) == sizeof(double) * 8); 21 | 22 | poses_se3 = 23 | torch::from_blob(pose_data.data(), {(long)pose_data.size(), 8L}, torch::TensorOptions().dtype(torch::kFloat64)) 24 | .clone() 25 | .cuda(); 26 | 27 | // The last element is just padding 28 | poses_se3.slice(1, 7, 8).zero_(); 29 | register_parameter("tangent_poses", tangent_poses); 30 | register_buffer("poses_se3", poses_se3); 31 | } 32 | 33 | PoseModuleImpl::PoseModuleImpl(Sophus::SE3d pose) 34 | { 35 | pose = pose.inverse(); 36 | poses_se3 = torch::from_blob(pose.data(), {1L, 8L}, torch::TensorOptions().dtype(torch::kFloat64)).clone().cuda(); 37 | tangent_poses = torch::zeros({1L, 6L}, torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCUDA)).clone(); 38 | register_buffer("poses_se3", poses_se3); 39 | register_parameter("tangent_poses", tangent_poses); 40 | } 41 | 42 | void PoseModuleImpl::ApplyTangent() 43 | { 44 | ApplyTangentToPose(tangent_poses, poses_se3); 45 | } 46 | std::vector PoseModuleImpl::Download() 47 | { 48 | torch::Tensor cp = poses_se3.contiguous().cpu(); 49 | 50 | std::vector result(cp.size(0)); 51 | 52 | memcpy(result[0].data(), cp.data_ptr(), sizeof(Sophus::SE3d) * result.size()); 53 | 54 | return result; 55 | } 56 | void PoseModuleImpl::SetPose(int id, Sophus::SE3d pose) 57 | { 58 | pose = pose.inverse(); 59 | torch::NoGradGuard ngg; 60 | auto new_pose = 61 | torch::from_blob(pose.data(), {1L, 8L}, torch::TensorOptions().dtype(torch::kFloat64)).clone().cuda(); 62 | poses_se3.slice(0, id, id + 1) = new_pose; 63 | } 64 | 65 | IntrinsicsModuleImpl::IntrinsicsModuleImpl(std::shared_ptr scene) 66 | { 67 | SAIGA_ASSERT(!scene->scene_cameras.empty()); 68 | 69 | if (scene->dataset_params.camera_model == CameraModel::PINHOLE_DISTORTION) 70 | { 71 | std::vector> intrinsics_data; 72 | 73 | 74 | for (auto& cam : scene->scene_cameras) 75 | { 76 | vec5 a = cam.K.coeffs(); 77 | vec8 b = cam.distortion.Coeffs(); 78 | 79 | 80 | Eigen::Matrix data; 81 | 82 | data.head<5>() = a; 83 | data.tail<8>() = b; 84 | 85 | intrinsics_data.push_back(data); 86 | } 87 | 88 | 89 | 90 | intrinsics = torch::from_blob(intrinsics_data.data(), {(long)intrinsics_data.size(), 13L}, 91 | torch::TensorOptions().dtype(torch::kFloat32)) 92 | .clone() 93 | .cuda(); 94 | 95 | 96 | register_parameter("intrinsics", intrinsics); 97 | std::cout << "Pinhole Intrinsics:" << std::endl; 98 | PrintTensorInfo(intrinsics); 99 | } 100 | else if (scene->dataset_params.camera_model == CameraModel::OCAM) 101 | { 102 | std::cout << "ocam Intrinsics:" << std::endl; 103 | 104 | long count = scene->scene_cameras.front().ocam.NumProjectParams(); 105 | 106 | std::vector intrinsics_data(count * scene->scene_cameras.size()); 107 | 108 | for (int i = 0; i < scene->scene_cameras.size(); ++i) 109 | { 110 | auto c = scene->scene_cameras[i].ocam; 111 | SAIGA_ASSERT(c.NumProjectParams() == count); 112 | 113 | float* ptr = intrinsics_data.data() + (count * i); 114 | 115 | auto params = c.ProjectParams(); 116 | 117 | for (int i = 0; i < params.size(); ++i) 118 | { 119 | ptr[i] = params[i]; 120 | } 121 | } 122 | 123 | intrinsics = torch::from_blob(intrinsics_data.data(), {(long)scene->scene_cameras.size(), count}, 124 | torch::TensorOptions().dtype(torch::kFloat32)) 125 | .clone() 126 | .cuda(); 127 | PrintTensorInfo(intrinsics); 128 | register_parameter("intrinsics", intrinsics); 129 | } 130 | else 131 | { 132 | SAIGA_EXIT_ERROR("unknown camera model"); 133 | } 134 | } 135 | 136 | IntrinsicsModuleImpl::IntrinsicsModuleImpl(IntrinsicsPinholef K) 137 | { 138 | Eigen::Matrix data; 139 | data.setZero(); 140 | data.head<5>() = K.coeffs(); 141 | intrinsics = torch::from_blob(data.data(), {1L, 13L}, torch::TensorOptions().dtype(torch::kFloat32)).clone().cuda(); 142 | 143 | register_parameter("intrinsics", intrinsics); 144 | } 145 | 146 | std::vector IntrinsicsModuleImpl::DownloadK() 147 | { 148 | torch::Tensor cp = intrinsics.clone().slice(1, 0, 5).contiguous().cpu(); 149 | std::vector> intrinsics_data(cp.size(0)); 150 | memcpy(intrinsics_data[0].data(), cp.data_ptr(), sizeof(Eigen::Matrix) * intrinsics_data.size()); 151 | 152 | 153 | std::vector result; 154 | for (auto v : intrinsics_data) 155 | { 156 | IntrinsicsPinholef f = v.head<5>().eval(); 157 | result.push_back(f); 158 | } 159 | return result; 160 | } 161 | 162 | std::vector IntrinsicsModuleImpl::DownloadDistortion() 163 | { 164 | torch::Tensor cp = intrinsics.contiguous().cpu(); 165 | std::vector> intrinsics_data(cp.size(0)); 166 | memcpy(intrinsics_data[0].data(), cp.data_ptr(), sizeof(Eigen::Matrix) * intrinsics_data.size()); 167 | 168 | 169 | std::vector result; 170 | for (auto v : intrinsics_data) 171 | { 172 | Distortionf f = v.tail<8>().eval(); 173 | result.push_back(f); 174 | } 175 | 176 | return result; 177 | } 178 | 179 | /* 180 | std::vector IntrinsicsModuleImpl::DownloadOcam(int width, int height) 181 | { 182 | torch::Tensor cp = intrinsics.contiguous().cpu(); 183 | std::vector> intrinsics_data(cp.size(0)); 184 | for(auto& v: intrinsics_data){ 185 | v.resize(cp.size(1)); 186 | } 187 | memcpy(intrinsics_data[0].data(), cp.data_ptr(), sizeof(float) * cp.size(1) * intrinsics_data.size()); 188 | 189 | 190 | std::vector> result; 191 | for (auto v : intrinsics_data) 192 | { 193 | OCam f = OCam(width,height, vec5(v[0],v[1],v[2],v[3],v[4]), std::vector(v.begin()+5, 194 | v.end())); result.push_back(f); 195 | } 196 | 197 | return result; 198 | }*/ 199 | 200 | void IntrinsicsModuleImpl::SetPinholeIntrinsics(int id, IntrinsicsPinholef K, Distortionf dis) 201 | { 202 | Eigen::Matrix intrinsics_data; 203 | 204 | intrinsics_data.head<5>() = K.coeffs(); 205 | intrinsics_data.tail<8>() = dis.Coeffs(); 206 | 207 | 208 | auto new_intr = torch::from_blob(intrinsics_data.data(), {1L, 13L}, torch::TensorOptions().dtype(torch::kFloat32)) 209 | .clone() 210 | .cuda(); 211 | 212 | torch::NoGradGuard ngg; 213 | // intrinsics.slice(0, id, id + 1) = new_intr; 214 | intrinsics = new_intr; 215 | } 216 | 217 | void IntrinsicsModuleImpl::SetOcamIntrinsics(int id, OCam ocam) 218 | { 219 | // Eigen::Matrix intrinsics_data; 220 | long count = ocam.NumProjectParams(); 221 | 222 | std::vector intrinsics_data(count); 223 | { 224 | auto params = ocam.ProjectParams(); 225 | for (int i = 0; i < params.size(); ++i) 226 | { 227 | intrinsics_data[i] = params[i]; 228 | } 229 | } 230 | 231 | auto new_intr = 232 | torch::from_blob(intrinsics_data.data(), {(long)1, count}, torch::TensorOptions().dtype(torch::kFloat32)) 233 | .clone() 234 | .cuda(); 235 | 236 | torch::NoGradGuard ngg; 237 | // intrinsics.slice(0, id, id + 1) = new_intr; 238 | intrinsics = new_intr; 239 | // PrintTensorInfo(intrinsics.slice(0, id, id + 1) ); 240 | } 241 | -------------------------------------------------------------------------------- /src/lib/data/NeuralStructure.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | #include "saiga/core/Core.h" 8 | 9 | #include "SceneData.h" 10 | #include "Settings.h" 11 | #include "config.h" 12 | 13 | #include 14 | 15 | 16 | using namespace Saiga; 17 | 18 | 19 | class PoseModuleImpl : public torch::nn::Module 20 | { 21 | public: 22 | PoseModuleImpl(std::shared_ptr scene); 23 | 24 | // Adds only a single pose 25 | PoseModuleImpl(Sophus::SE3d pose); 26 | 27 | std::vector Download(); 28 | void SetPose(int id, Sophus::SE3d pose); 29 | 30 | // double: [num_cameras, 8] 31 | torch::Tensor poses_se3; 32 | 33 | // double: [num_cameras, 6] 34 | torch::Tensor tangent_poses; 35 | 36 | void ApplyTangent(); 37 | }; 38 | TORCH_MODULE(PoseModule); 39 | 40 | 41 | class IntrinsicsModuleImpl : public torch::nn::Module 42 | { 43 | public: 44 | IntrinsicsModuleImpl(std::shared_ptr scene); 45 | 46 | // Adds only a single intrinsic 47 | IntrinsicsModuleImpl(IntrinsicsPinholef K); 48 | 49 | std::vector DownloadDistortion(); 50 | std::vector DownloadK(); 51 | 52 | 53 | void SetPinholeIntrinsics(int id, IntrinsicsPinholef K, Distortionf dis); 54 | void SetOcamIntrinsics(int id, OCam ocam); 55 | 56 | // [num_cameras, num_model_params] 57 | // Pinhole + Distortion: [num_cameras, 5 + 8] 58 | // OCam: [num_cameras, 5 + world2cam_coefficients] 59 | torch::Tensor intrinsics; 60 | }; 61 | TORCH_MODULE(IntrinsicsModule); 62 | -------------------------------------------------------------------------------- /src/lib/data/Settings.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #include "Settings.h" 7 | 8 | #include "saiga/core/imgui/imgui.h" 9 | 10 | void CombinedParams::Check() 11 | { 12 | if (net_params.conv_block == "partial" || net_params.conv_block == "partial_multi" || 13 | pipeline_params.enable_environment_map || pipeline_params.cat_masks_to_color) 14 | { 15 | render_params.output_background_mask = true; 16 | } 17 | 18 | if (pipeline_params.skip_neural_render_network) 19 | { 20 | pipeline_params.num_texture_channels = 3; 21 | net_params.num_input_layers = 1; 22 | } 23 | 24 | net_params.num_input_channels = pipeline_params.num_texture_channels; 25 | render_params.num_texture_channels = pipeline_params.num_texture_channels; 26 | render_params.use_point_adding_and_removing_module = pipeline_params.use_point_adding_and_removing_module; 27 | render_params.use_environment_map = pipeline_params.enable_environment_map; 28 | render_params.debug_max_list_length = points_adding_params.debug_max_list_length; 29 | 30 | if (!optimizer_params.fix_points || !optimizer_params.fix_poses || !optimizer_params.fix_intrinsics || 31 | !optimizer_params.fix_dynamic_refinement) 32 | { 33 | render_params.need_point_gradient = true; 34 | std::cout << "POINT GRADIENTS ARE COMPUTED." << std::endl; 35 | } 36 | else 37 | { 38 | render_params.need_point_gradient = false; 39 | 40 | std::cout << "Point gradients omitted." << std::endl; 41 | } 42 | 43 | 44 | 45 | SAIGA_ASSERT(!train_params.texture_color_init || pipeline_params.num_texture_channels == 3); 46 | 47 | if (pipeline_params.cat_env_to_color) 48 | { 49 | net_params.num_input_channels += pipeline_params.env_map_channels; 50 | } 51 | else 52 | { 53 | pipeline_params.env_map_channels = pipeline_params.num_texture_channels; 54 | } 55 | 56 | if (pipeline_params.cat_masks_to_color) 57 | { 58 | net_params.num_input_channels += 1; 59 | } 60 | if (render_params.add_depth_to_network) 61 | { 62 | net_params.num_input_channels += 1; 63 | } 64 | } 65 | void CombinedParams::imgui() 66 | { 67 | ImGui::Checkbox("render_points", &render_params.render_points); 68 | ImGui::Checkbox("render_outliers", &render_params.render_outliers); 69 | ImGui::Checkbox("drop_out_points_by_radius", &render_params.drop_out_points_by_radius); 70 | ImGui::SliderFloat("drop_out_radius_threshold", &render_params.drop_out_radius_threshold, 0, 5); 71 | ImGui::Checkbox("super_sampling", &render_params.super_sampling); 72 | 73 | ImGui::Checkbox("check_normal", &render_params.check_normal); 74 | 75 | 76 | ImGui::Checkbox("debug_weight_color", &render_params.debug_weight_color); 77 | ImGui::Checkbox("debug_depth_color", &render_params.debug_depth_color); 78 | ImGui::SliderFloat("debug_max_weight", &render_params.debug_max_weight, 0, 100); 79 | ImGui::Checkbox("debug_print_num_rendered_points", &render_params.debug_print_num_rendered_points); 80 | 81 | ImGui::SliderFloat("dropout", &render_params.dropout, 0, 1); 82 | ImGui::SliderFloat("depth_accept", &render_params.depth_accept, 0, 0.1); 83 | ImGui::SliderFloat("depth_accept_blend", &render_params.depth_accept_blend, 0.001, 1); 84 | 85 | ImGui::SliderFloat("dist_cutoff", &render_params.dist_cutoff, 0, 1); 86 | 87 | ImGui::Separator(); 88 | } 89 | -------------------------------------------------------------------------------- /src/lib/git_sha1.h.in: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #define GIT_SHA1 "@MY_GIT_SHA1@" 10 | -------------------------------------------------------------------------------- /src/lib/models/NeuralCamera.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/core/camera/HDR.h" 9 | #include "saiga/core/imgui/imgui.h" 10 | #include "saiga/core/util/FileSystem.h" 11 | #include "saiga/core/util/file.h" 12 | #include "saiga/vision/VisionIncludes.h" 13 | #include "saiga/vision/torch/TorchHelper.h" 14 | 15 | #include "data/Settings.h" 16 | 17 | using namespace Saiga; 18 | 19 | class VignetteNetImpl : public torch::nn::Module 20 | { 21 | public: 22 | VignetteNetImpl(ivec2 image_size); 23 | 24 | void Set(vec3 params, vec2 center) 25 | { 26 | torch::NoGradGuard ngg; 27 | vignette_params[0] = params(0); 28 | vignette_params[1] = params(1); 29 | vignette_params[2] = params(1); 30 | 31 | vignette_center[0][0][0][0] = center(0); 32 | vignette_center[0][1][0][0] = center(1); 33 | } 34 | 35 | void PrintParams(const std::string& log_dir, std::string& name); 36 | 37 | void ApplyConstraints() {} 38 | 39 | torch::Tensor forward(torch::Tensor uv); 40 | 41 | ivec2 image_size; 42 | bool use_calibrated_center = false; 43 | vec2 calibrated_center; 44 | 45 | torch::Tensor vignette_params; 46 | torch::Tensor vignette_center; 47 | }; 48 | 49 | TORCH_MODULE(VignetteNet); 50 | 51 | // This is just test. 52 | // It does not work very well. 53 | class RollingShutterNetImpl : public torch::nn::Module 54 | { 55 | public: 56 | RollingShutterNetImpl(int frames, int size_x = 32, int size_y = 32); 57 | 58 | void ApplyConstraints() { torch::NoGradGuard ngg; } 59 | 60 | torch::Tensor forward(torch::Tensor x, torch::Tensor frame_index, torch::Tensor uv); 61 | 62 | // [n_frames, size, 1 ,1] 63 | torch::Tensor transform_grid; 64 | 65 | 66 | torch::Tensor uv_local; 67 | 68 | torch::nn::functional::GridSampleFuncOptions options; 69 | }; 70 | TORCH_MODULE(RollingShutterNet); 71 | 72 | // Also a test 73 | // Not really used 74 | class MotionblurNetImpl : public torch::nn::Module 75 | { 76 | public: 77 | MotionblurNetImpl(int frames, int radius = 3, float initial_blur = 0); 78 | 79 | void ApplyConstraints() 80 | { 81 | torch::NoGradGuard ngg; 82 | // Not exact 0,1 so we don't get zero gradients 83 | blur_values.clamp_(0.0001, 0.9999); 84 | } 85 | 86 | torch::Tensor forward(torch::Tensor x, torch::Tensor frame_index, torch::Tensor scale); 87 | 88 | // [n, 1, 1 ,1] 89 | torch::Tensor blur_values; 90 | 91 | torch::Tensor kernel_raw_x, kernel_raw_y; 92 | 93 | int padding; 94 | }; 95 | TORCH_MODULE(MotionblurNet); 96 | 97 | 98 | class CameraResponseNetImpl : public torch::nn::Module 99 | { 100 | public: 101 | CameraResponseNetImpl(int params, int num_channels, float initial_gamma, float leaky_clamp_value = 0); 102 | 103 | std::vector> GetCRF(); 104 | 105 | void ApplyConstraints() {} 106 | 107 | torch::Tensor forward(torch::Tensor image); 108 | 109 | torch::Tensor ParamLoss(); 110 | 111 | int NumParameters() { return response.size(3); } 112 | 113 | torch::nn::functional::GridSampleFuncOptions options; 114 | torch::Tensor response; 115 | torch::Tensor leaky_value; 116 | }; 117 | 118 | TORCH_MODULE(CameraResponseNet); 119 | 120 | class NeuralCameraImpl : public torch::nn::Module 121 | { 122 | public: 123 | NeuralCameraImpl(ivec2 image_size, NeuralCameraParams params, int frames, std::vector initial_exposure, 124 | std::vector initial_wb); 125 | 126 | torch::Tensor ParamLoss(torch::Tensor frame_index); 127 | 128 | torch::Tensor forward(torch::Tensor x, torch::Tensor frame_index, torch::Tensor uv, torch::Tensor scale, 129 | float fixed_exposure, vec3 fixed_white_balance); 130 | 131 | float param_loss_exposure = 20.f; 132 | float param_loss_wb = 20.f; 133 | 134 | // [n, 1, 1 ,1] 135 | torch::Tensor exposures_values; 136 | torch::Tensor exposures_values_reference; 137 | 138 | // [n, 3, 1, 1] 139 | torch::Tensor white_balance_values; 140 | torch::Tensor white_balance_values_reference; 141 | 142 | RollingShutterNet rolling_shutter = nullptr; 143 | VignetteNet vignette_net = nullptr; 144 | CameraResponseNet camera_response = nullptr; 145 | NeuralCameraParams params; 146 | 147 | std::vector DownloadExposure(); 148 | std::vector DownloadWhitebalance(); 149 | 150 | 151 | // Interpolates exposure and wb for all passed indices from the neighbors. 152 | // Usually used for train/test split. 153 | void InterpolateFromNeighbors(std::vector indices); 154 | 155 | void SaveCheckpoint(const std::string& checkpoint_dir); 156 | void LoadCheckpoint(const std::string& checkpoint_dir); 157 | 158 | void ApplyConstraints(); 159 | }; 160 | TORCH_MODULE(NeuralCamera); 161 | -------------------------------------------------------------------------------- /src/lib/models/NeuralTexture.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #include "NeuralTexture.h" 8 | 9 | using namespace Saiga; 10 | 11 | NeuralPointTextureImpl::NeuralPointTextureImpl(int num_channels, int num_points, bool random_init, bool log_texture, 12 | bool use_stability) 13 | : log_texture(log_texture), random_init(random_init) 14 | { 15 | if (random_init) 16 | { 17 | // Random init in the range [-factor/2, factor/2] 18 | float factor = 2; 19 | texture_raw = (torch::rand({num_channels, num_points}) * factor); 20 | texture_raw = texture_raw - torch::ones_like(texture_raw) * (factor / 2); 21 | 22 | texture_raw = torch::empty({num_channels, num_points}); 23 | texture_raw.uniform_(0, 1); 24 | } 25 | else 26 | { 27 | texture_raw = torch::ones({num_channels, num_points}) * 0.5f; 28 | } 29 | texture = texture_raw.clone(); 30 | 31 | background_color_raw = torch::ones({num_channels}) * -1; 32 | 33 | confidence_value_of_point = torch::ones({1, num_points}) * 0.5; 34 | confidence_raw = torch::ones({1, num_points}) * 0.5; 35 | 36 | if (log_texture) 37 | { 38 | std::cout << "initialized first texture channel (before checkpoint loading) with 1" << std::endl; 39 | texture_raw = torch::log(torch::ones({num_channels, num_points}) * 0.5); 40 | background_color = torch::ones({num_channels}) * -1; 41 | } 42 | 43 | 44 | // register_parameter("texture", texture); 45 | register_parameter("texture", texture_raw); 46 | register_parameter("background_color", background_color_raw); 47 | register_parameter("confidence_value_of_point", confidence_raw); 48 | 49 | 50 | std::cout << "GPU memory - Texture " << texture_raw.sizes() << " : " 51 | << (texture_raw.nbytes() + confidence_value_of_point.nbytes() + confidence_raw.nbytes() + 52 | background_color_raw.nbytes()) / 53 | 1000000.0 54 | << "MB" << std::endl; 55 | } 56 | 57 | NeuralPointTextureImpl::NeuralPointTextureImpl(const UnifiedMesh& model, int channels, int num_points) 58 | : NeuralPointTextureImpl(channels, (num_points > 0) ? num_points : model.NumVertices(), false, false) 59 | { 60 | int num_p = (num_points > 0) ? num_points : model.NumVertices(); 61 | if (channels == 3) 62 | { 63 | std::vector colors; 64 | for (auto c : model.color) 65 | { 66 | if (colors.size() < num_p) colors.push_back(c.head<3>()); 67 | } 68 | for (int i = colors.size(); i < num_p; ++i) 69 | { 70 | colors.push_back(vec3(1, 1, 1)); 71 | } 72 | 73 | auto t = 74 | torch::from_blob(colors.data(), {(long)colors.size(), 3}, torch::TensorOptions().dtype(torch::kFloat32)) 75 | .to(texture_raw.device()) 76 | .permute({1, 0}) 77 | .contiguous(); 78 | SAIGA_ASSERT(t.sizes() == texture_raw.sizes()); 79 | 80 | { 81 | torch::NoGradGuard ngg; 82 | texture_raw.set_(t); 83 | } 84 | 85 | SetBackgroundColor({0, 0, 0}); 86 | } 87 | else if (channels == 4) 88 | { 89 | std::vector colors; 90 | for (auto c : model.color) 91 | { 92 | c(3) = 1; 93 | if (colors.size() < num_p) colors.push_back(c.head<4>()); 94 | } 95 | for (int i = model.color.size(); i < num_p; ++i) 96 | { 97 | colors.push_back(vec4(1, 1, 1, 0)); 98 | } 99 | 100 | 101 | auto t = 102 | torch::from_blob(colors.data(), {(long)colors.size(), 4}, torch::TensorOptions().dtype(torch::kFloat32)) 103 | .to(texture_raw.device()) 104 | .permute({1, 0}) 105 | .contiguous(); 106 | SAIGA_ASSERT(t.sizes() == texture_raw.sizes()); 107 | 108 | { 109 | torch::NoGradGuard ngg; 110 | texture_raw.set_(t); 111 | } 112 | 113 | SetBackgroundColor({0, 0, 0, 1}); 114 | } 115 | } 116 | 117 | 118 | torch::Tensor NeuralPointTextureImpl::ResizeAndFill(torch::Tensor& t, int new_point_size, bool fill_with_zeros) 119 | { 120 | using namespace torch::indexing; 121 | SAIGA_ASSERT(t.sizes().size() == 2 || t.sizes().size() == 3); 122 | 123 | torch::Tensor new_vals; 124 | if (t.sizes().size() == 2) 125 | { 126 | new_vals = torch::ones({t.sizes()[0], long(new_point_size) - t.sizes()[1]}, t.options()) * 0.5; 127 | } 128 | else if (t.sizes().size() == 3) 129 | { 130 | new_vals = torch::ones({t.sizes()[0], t.sizes()[1], new_point_size - t.sizes()[2]}, t.options()) * 0.5; 131 | } 132 | if (fill_with_zeros) 133 | { 134 | new_vals *= 0; 135 | } 136 | else if (random_init) 137 | { 138 | float factor = 2; 139 | new_vals = (torch::rand_like(new_vals) * factor) - (new_vals * (2 * (factor / 2))); 140 | } 141 | 142 | torch::Tensor t_n; 143 | if (t.sizes().size() == 2) 144 | { 145 | t_n = torch::cat({t.clone(), new_vals}, 1); 146 | } 147 | else 148 | { 149 | t_n = torch::cat({t.clone(), new_vals}, 2); 150 | } 151 | 152 | { 153 | torch::NoGradGuard ngg; 154 | t.set_(t_n); 155 | } 156 | 157 | return new_vals; 158 | } 159 | 160 | void NeuralPointTextureImpl::RemoveAndFlatten(torch::Tensor& t, torch::Tensor& indices_to_keep) 161 | { 162 | using namespace torch::indexing; 163 | { 164 | torch::NoGradGuard ngg; 165 | // std::cout << TensorInfo(indices_to_keep.squeeze().unsqueeze(0)) << std::endl; 166 | // std::cout << TensorInfo(t) << std::endl; 167 | // auto values_keep = t.index({indices_to_keep.squeeze().unsqueeze(0) }); 168 | auto values_keep = t.index_select(t.sizes().size() - 1, indices_to_keep.squeeze().to(t.device())); 169 | // std::cout << TensorInfo(values_keep) << std::endl; 170 | 171 | t.set_(values_keep); 172 | } 173 | } 174 | -------------------------------------------------------------------------------- /src/lib/models/NeuralTexture.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/core/model/UnifiedMesh.h" 9 | #include "saiga/core/util/FileSystem.h" 10 | #include "saiga/core/util/file.h" 11 | #include "saiga/cuda/cudaTimer.h" 12 | #include "saiga/vision/torch/TorchHelper.h" 13 | 14 | 15 | using namespace Saiga; 16 | 17 | // [num_channels, num_points] 18 | 19 | class NeuralPointTextureImpl : public torch::nn::Module 20 | { 21 | public: 22 | NeuralPointTextureImpl(int num_channels, int num_points, bool random_init, bool log_texture, 23 | bool use_stability = false); 24 | 25 | NeuralPointTextureImpl(const Saiga::UnifiedMesh& model, int channels = 4, int num_points = -1); 26 | 27 | std::vector GetDescriptorSlow(int i) 28 | { 29 | std::vector desc; 30 | for (int j = 0; j < texture_raw.size(0); ++j) 31 | { 32 | float f = texture_raw[j][i].item().toFloat(); 33 | desc.push_back(f); 34 | } 35 | return desc; 36 | } 37 | 38 | void PrepareConfidence(float narrowing_param_times_epoch) 39 | { 40 | confidence_value_of_point = torch::sigmoid((10.f + narrowing_param_times_epoch) * confidence_raw); 41 | } 42 | 43 | void PrepareTexture(bool abs) 44 | { 45 | if (abs) 46 | { 47 | texture = torch::abs(texture_raw); 48 | background_color = torch::abs(background_color_raw); 49 | } 50 | else 51 | { 52 | texture = texture_raw.clone(); 53 | background_color = background_color_raw.clone(); 54 | } 55 | } 56 | 57 | std::vector GetBackgroundColor() 58 | { 59 | std::vector desc; 60 | for (int j = 0; j < background_color.size(0); ++j) 61 | { 62 | float f = background_color[j].item().toFloat(); 63 | desc.push_back(f); 64 | } 65 | return desc; 66 | } 67 | 68 | void SetBackgroundColor(std::vector col) 69 | { 70 | torch::NoGradGuard ngg; 71 | background_color_raw.set_( 72 | torch::from_blob(col.data(), {(long)col.size()}).to(background_color_raw.device()).clone()); 73 | } 74 | 75 | int NumPoints() { return texture_raw.size(1); } 76 | int TextureChannels() 77 | { 78 | SAIGA_ASSERT(texture_raw.dim() == 2); 79 | return texture_raw.size(0); 80 | } 81 | torch::Tensor ResizeAndFill(torch::Tensor& t, int new_point_size, bool fill_with_zeros = false); 82 | 83 | void RemoveAndFlatten(torch::Tensor& t, torch::Tensor& indices_to_keep); 84 | 85 | void RemoveAndFlattenTexture(torch::Tensor& indices_to_remove) 86 | { 87 | RemoveAndFlatten(texture_raw, indices_to_remove); 88 | RemoveAndFlatten(confidence_value_of_point, indices_to_remove); 89 | RemoveAndFlatten(confidence_raw, indices_to_remove); 90 | std::cout << "Removed and flattened Texture to " << texture.sizes() << std::endl; 91 | } 92 | 93 | // increases size of all buffers and initializes data 94 | void EnlargeTexture(int new_num_points) 95 | { 96 | ResizeAndFill(texture_raw, new_num_points); 97 | 98 | confidence_value_of_point.reset(); 99 | texture.reset(); 100 | ResizeAndFill(confidence_raw, new_num_points, true); 101 | 102 | std::cout << "Resized Texture to " << texture_raw.sizes() << std::endl; 103 | } 104 | 105 | void to_except_texture(torch::Device dev) 106 | { 107 | confidence_value_of_point = confidence_value_of_point.to(dev); 108 | background_color = background_color.to(dev); 109 | } 110 | 111 | bool log_texture; 112 | bool random_init; 113 | 114 | // [channels, points] 115 | torch::Tensor texture; 116 | torch::Tensor texture_raw; 117 | 118 | 119 | 120 | // [1, points] 121 | torch::Tensor confidence_value_of_point; 122 | torch::Tensor confidence_raw; 123 | 124 | // [channels] 125 | torch::Tensor background_color; 126 | torch::Tensor background_color_raw; 127 | }; 128 | 129 | TORCH_MODULE(NeuralPointTexture); 130 | -------------------------------------------------------------------------------- /src/lib/models/Pipeline.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/core/util/file.h" 9 | #include "saiga/cuda/imgui_cuda.h" 10 | #include "saiga/vision/torch/ImageSimilarity.h" 11 | #include "saiga/vision/torch/TorchHelper.h" 12 | #include "saiga/vision/torch/VGGLoss.h" 13 | 14 | #include "data/NeuralScene.h" 15 | #include "models/Networks.h" 16 | #include "models/NeuralTexture.h" 17 | #include "models/Pipeline.h" 18 | #include "rendering/PointRenderer.h" 19 | #include "rendering/RenderModule.h" 20 | using namespace Saiga; 21 | 22 | struct LossResult 23 | { 24 | float loss_vgg = 0; 25 | float loss_l1 = 0; 26 | float loss_mse = 0; 27 | float loss_psnr = 0; 28 | float loss_ssim = 0; 29 | float loss_lpips = 0; 30 | 31 | float loss_float = 0; 32 | float loss_float_param = 0; 33 | 34 | int count = 0; 35 | 36 | LossResult& operator+=(const LossResult& other) 37 | { 38 | loss_vgg += other.loss_vgg; 39 | loss_l1 += other.loss_l1; 40 | loss_mse += other.loss_mse; 41 | loss_psnr += other.loss_psnr; 42 | loss_ssim += other.loss_ssim; 43 | loss_lpips += other.loss_lpips; 44 | loss_float += other.loss_float; 45 | loss_float_param += other.loss_float_param; 46 | count += other.count; 47 | return *this; 48 | } 49 | 50 | LossResult& operator/=(float value) 51 | { 52 | loss_vgg /= value; 53 | loss_l1 /= value; 54 | loss_mse /= value; 55 | loss_psnr /= value; 56 | loss_ssim /= value; 57 | loss_lpips /= value; 58 | loss_float /= value; 59 | loss_float_param /= value; 60 | return *this; 61 | } 62 | 63 | LossResult Average() 64 | { 65 | LossResult cpy = *this; 66 | cpy /= count; 67 | return cpy; 68 | } 69 | 70 | void AppendToFile(const std::string& file, int epoch) 71 | { 72 | std::ofstream strm(file, std::ios_base::app); 73 | 74 | Table tab({10, 15, 15, 15, 15, 15, 15, 15}, strm, ','); 75 | if (epoch == 0) 76 | { 77 | tab << "ep" 78 | << "vgg" 79 | << "lpips" 80 | << "l1" 81 | << "psrn" 82 | << "ssim" 83 | << "param" 84 | << "count"; 85 | } 86 | tab << epoch << loss_vgg << loss_lpips << loss_l1 << loss_psnr << loss_ssim << loss_float_param << count; 87 | } 88 | 89 | void Print() 90 | { 91 | console << "Param " << loss_float_param << " VGG " << loss_vgg << " L1 " << loss_l1 << " MSE " << loss_mse 92 | << " PSNR " << loss_psnr << " SSIM " << loss_ssim << " LPIPS " << loss_lpips << " count " << count 93 | << std::endl; 94 | } 95 | }; 96 | 97 | struct ForwardResult 98 | { 99 | std::vector> outputs; 100 | std::vector> targets; 101 | std::vector image_ids; 102 | 103 | torch::Tensor x; 104 | torch::Tensor loss; 105 | torch::Tensor target; 106 | 107 | LossResult float_loss; 108 | }; 109 | 110 | 111 | class RenderNet 112 | { 113 | public: 114 | RenderNet() {} 115 | // RenderNet(MultiScaleUnet2dParams params): multiScaleUnet2d(std::make_shared(params)){} 116 | virtual void to(torch::DeviceType d) = 0; 117 | virtual void to(torch::ScalarType d) = 0; 118 | 119 | virtual void eval() = 0; 120 | virtual void train() = 0; 121 | virtual void train(bool t) = 0; 122 | virtual bool valid() = 0; 123 | 124 | virtual void save(std::string path) = 0; 125 | virtual void load(std::string path) = 0; 126 | 127 | virtual std::vector parameters() = 0; 128 | 129 | virtual at::Tensor forward(ArrayView inputs) = 0; 130 | virtual at::Tensor forward(ArrayView inputs, ArrayView masks) = 0; 131 | }; 132 | 133 | template 134 | class DerivedRenderNet : public RenderNet 135 | { 136 | public: 137 | std::shared_ptr network = nullptr; 138 | DerivedRenderNet(MultiScaleUnet2dParams params) : network(std::make_shared(params)) {} 139 | 140 | virtual void to(torch::DeviceType d) { (*network)->to(d); } 141 | virtual void to(torch::ScalarType d) { (*network)->to(d); } 142 | virtual void eval() { (*network)->eval(); } 143 | virtual void train() { (*network)->train(); } 144 | virtual void train(bool t) { (*network)->train(t); } 145 | virtual bool valid() { return network != nullptr; }; 146 | 147 | virtual std::vector parameters() { return (*network)->parameters(); }; 148 | 149 | 150 | virtual at::Tensor forward(ArrayView inputs) { return (*network)->forward(inputs); } 151 | virtual at::Tensor forward(ArrayView inputs, ArrayView masks) 152 | { 153 | return (*network)->forward(inputs, masks); 154 | } 155 | 156 | virtual void save(std::string path) { torch::save(*network, path); } 157 | 158 | virtual void load(std::string path) 159 | { 160 | if (*network && std::filesystem::exists(path)) 161 | { 162 | std::cout << "Load Checkpoint render" << std::endl; 163 | torch::load(*network, path); 164 | } 165 | } 166 | }; 167 | 168 | using RenderNetwork = RenderNet; 169 | using RenderNetworkParams = MultiScaleUnet2dParams; 170 | 171 | class NeuralPipeline 172 | { 173 | public: 174 | NeuralPipeline(std::shared_ptr params); 175 | 176 | void Train(bool train); 177 | 178 | void SaveCheckpoint(const std::string& dir) 179 | { 180 | render_network->save(dir + "/render_net.pth"); 181 | 182 | } 183 | void LoadCheckpoint(const std::string& dir) 184 | { 185 | render_network->load(dir + "/render_net.pth"); 186 | 187 | } 188 | 189 | // void SaveCheckpoint(const std::string& dir) { torch::save(render_network, dir + "/render_net.pth"); } 190 | // void LoadCheckpoint(const std::string& dir) 191 | // { 192 | // if (render_network.valid() && std::filesystem::exists(dir + "/render_net.pth")) 193 | // { 194 | // std::cout << "Load Checkpoint render" << std::endl; 195 | // torch::load(render_network, dir + "/render_net.pth"); 196 | // } 197 | // } 198 | 199 | void Log(const std::string& dir); 200 | 201 | void OptimizerStep(int epoch_id); 202 | void OptimizerClear(int epoch_id); 203 | void UpdateLearningRate(double factor); 204 | 205 | ForwardResult Forward(NeuralScene& scene, std::vector& batch, torch::Tensor global_mask, 206 | bool loss_statistics, int current_epoch = -1, bool keep_image = false, 207 | float fixed_exposure = std::numeric_limits::infinity(), 208 | vec3 fixed_white_balance = vec3(std::numeric_limits::infinity(), 0, 0)); 209 | 210 | std::shared_ptr render_network = nullptr; 211 | torch::DeviceType device = torch::kCUDA; 212 | 213 | PointRenderModule render_module = nullptr; 214 | std::shared_ptr params; 215 | CUDA::CudaTimerSystem* timer_system = nullptr; 216 | 217 | 218 | std::shared_ptr render_optimizer; 219 | std::shared_ptr refinement_optimizer; 220 | 221 | // Loss stuff 222 | std::shared_ptr loss_vgg = nullptr; 223 | PSNR loss_psnr = PSNR(0, 1); 224 | LPIPS loss_lpips = LPIPS("loss/traced_lpips.pt"); 225 | SSIM loss_ssim = SSIM(); 226 | }; -------------------------------------------------------------------------------- /src/lib/models/my_adam.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #include "my_adam.h" 7 | 8 | #include "saiga/core/util/assert.h" 9 | #include "saiga/vision/torch/TorchHelper.h" 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | 20 | namespace torch 21 | { 22 | namespace optim 23 | { 24 | 25 | Tensor MyAdam::step(LossClosure closure) 26 | { 27 | NoGradGuard no_grad; 28 | Tensor loss = {}; 29 | if (closure != nullptr) 30 | { 31 | at::AutoGradMode enable_grad(true); 32 | loss = closure(); 33 | } 34 | for (auto& group : param_groups_) 35 | { 36 | for (auto& p : group.params()) 37 | { 38 | if (!p.grad().defined()) 39 | { 40 | continue; 41 | } 42 | auto grad = p.grad(); 43 | TORCH_CHECK(!grad.is_sparse(), 44 | "Adam does not support sparse gradients" /*, please consider SparseAdam instead*/); 45 | auto param_state = state_.find(c10::guts::to_string(p.unsafeGetTensorImpl())); 46 | auto& options = static_cast(group.options()); 47 | 48 | // State initialization 49 | if (param_state == state_.end()) 50 | { 51 | auto state = std::make_unique(); 52 | state->step(0); 53 | // Exponential moving average of gradient values 54 | state->exp_avg(torch::zeros_like(p, MemoryFormat::Preserve)); 55 | // Exponential moving average of squared gradient values 56 | state->exp_avg_sq(torch::zeros_like(p, MemoryFormat::Preserve)); 57 | if (options.amsgrad()) 58 | { 59 | // Maintains max of all exp. moving avg. of sq. grad. values 60 | state->max_exp_avg_sq(torch::zeros_like(p, MemoryFormat::Preserve)); 61 | } 62 | state_[c10::guts::to_string(p.unsafeGetTensorImpl())] = std::move(state); 63 | } 64 | 65 | auto& state = static_cast(*state_[c10::guts::to_string(p.unsafeGetTensorImpl())]); 66 | auto& exp_avg = state.exp_avg(); 67 | auto& exp_avg_sq = state.exp_avg_sq(); 68 | auto& max_exp_avg_sq = state.max_exp_avg_sq(); 69 | 70 | state.step(state.step() + 1); 71 | auto beta1 = std::get<0>(options.betas()); 72 | auto beta2 = std::get<1>(options.betas()); 73 | 74 | auto bias_correction1 = 1 - std::pow(beta1, state.step()); 75 | auto bias_correction2 = 1 - std::pow(beta2, state.step()); 76 | 77 | if (options.weight_decay() != 0) 78 | { 79 | grad = grad.add(p, options.weight_decay()); 80 | } 81 | 82 | // std::cout << Saiga::TensorInfo(grad) << " " << Saiga::TensorInfo(exp_avg) << " " << beta1 << " " 83 | // << 1 - beta1 << std::endl; 84 | // Decay the first and second moment running average coefficient 85 | exp_avg.mul_(beta1).add_(grad, 1 - beta1); 86 | exp_avg_sq.mul_(beta2).addcmul_(grad, grad, 1 - beta2); 87 | 88 | Tensor denom; 89 | if (options.amsgrad()) 90 | { 91 | // Maintains the maximum of all 2nd moment running avg. till now 92 | torch::max_out(max_exp_avg_sq, exp_avg_sq, max_exp_avg_sq); 93 | // Use the max. for normalizing running avg. of gradient 94 | denom = (max_exp_avg_sq.sqrt() / sqrt(bias_correction2)).add_(options.eps()); 95 | } 96 | else 97 | { 98 | denom = (exp_avg_sq.sqrt() / sqrt(bias_correction2)).add_(options.eps()); 99 | } 100 | 101 | auto step_size = options.lr() / bias_correction1; 102 | p.addcdiv_(exp_avg, denom, -step_size); 103 | } 104 | } 105 | return loss; 106 | } 107 | void MyAdam::shrinkenInternalState(int param_group_index, torch::Tensor indices_to_keep) 108 | { 109 | SAIGA_ASSERT(param_group_index < param_groups_.size()); 110 | auto& group = param_groups_[param_group_index]; 111 | 112 | { 113 | for (auto& p : group.params()) 114 | { 115 | auto param_state = state_.find(c10::guts::to_string(p.unsafeGetTensorImpl())); 116 | // not created yet -> will we initialized later 117 | if (param_state == state_.end()) continue; 118 | auto& state = static_cast(*state_[c10::guts::to_string(p.unsafeGetTensorImpl())]); 119 | auto& exp_avg = state.exp_avg(); 120 | auto& exp_avg_sq = state.exp_avg_sq(); 121 | auto& max_exp_avg_sq = state.max_exp_avg_sq(); 122 | 123 | 124 | auto remove_selected = [&](torch::Tensor& t) 125 | { 126 | auto values_keep = t.index_select(t.sizes().size() - 1, indices_to_keep.squeeze().to(t.device())); 127 | t.set_(values_keep); 128 | }; 129 | 130 | if (exp_avg.defined()) remove_selected(exp_avg); 131 | if (exp_avg_sq.defined()) remove_selected(exp_avg_sq); 132 | if (max_exp_avg_sq.defined()) remove_selected(max_exp_avg_sq); 133 | } 134 | } 135 | } 136 | 137 | 138 | void MyAdam::appendToInternalState(int param_group_index, int new_size) 139 | { 140 | SAIGA_ASSERT(param_group_index < param_groups_.size()); 141 | auto& group = param_groups_[param_group_index]; 142 | { 143 | for (auto& p : group.params()) 144 | { 145 | auto param_state = state_.find(c10::guts::to_string(p.unsafeGetTensorImpl())); 146 | // not created yet -> will we initialized later 147 | if (param_state == state_.end()) continue; 148 | auto& state = static_cast(*state_[c10::guts::to_string(p.unsafeGetTensorImpl())]); 149 | auto& exp_avg = state.exp_avg(); 150 | auto& exp_avg_sq = state.exp_avg_sq(); 151 | auto& max_exp_avg_sq = state.max_exp_avg_sq(); 152 | 153 | auto add_selected = [&](torch::Tensor& t) 154 | { 155 | torch::Tensor new_vals; 156 | int new_point_size = new_size - t.size(-1); 157 | std::vector sizes_tensor(t.sizes().size()); 158 | for (int i = 0; i < t.sizes().size(); ++i) sizes_tensor[i] = t.sizes()[i]; 159 | sizes_tensor[sizes_tensor.size() - 1] = new_point_size; 160 | 161 | new_vals = torch::zeros(sizes_tensor, t.options()); 162 | auto t_n = torch::cat({t.clone(), new_vals}, -1); 163 | t.set_(t_n); 164 | }; 165 | if (exp_avg.defined()) add_selected(exp_avg); 166 | if (exp_avg_sq.defined()) add_selected(exp_avg_sq); 167 | if (max_exp_avg_sq.defined()) add_selected(max_exp_avg_sq); 168 | } 169 | } 170 | } 171 | 172 | void MyAdam::save(serialize::OutputArchive& archive) const 173 | { 174 | SAIGA_ASSERT(false); 175 | } 176 | 177 | void MyAdam::load(serialize::InputArchive& archive) 178 | { 179 | SAIGA_ASSERT(false); 180 | } 181 | } // namespace optim 182 | } // namespace torch -------------------------------------------------------------------------------- /src/lib/models/my_adam.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | namespace torch 16 | { 17 | namespace serialize 18 | { 19 | class OutputArchive; 20 | class InputArchive; 21 | } // namespace serialize 22 | } // namespace torch 23 | 24 | namespace torch 25 | { 26 | namespace optim 27 | { 28 | 29 | class TORCH_API MyAdam : public Optimizer 30 | { 31 | public: 32 | explicit MyAdam(std::vector param_groups, AdamOptions defaults = {}) 33 | : Optimizer(std::move(param_groups), std::make_unique(defaults)) 34 | { 35 | TORCH_CHECK(defaults.lr() >= 0, "Invalid learning rate: ", defaults.lr()); 36 | TORCH_CHECK(defaults.eps() >= 0, "Invalid epsilon value: ", defaults.eps()); 37 | auto betas = defaults.betas(); 38 | TORCH_CHECK(0 <= std::get<0>(betas) && std::get<0>(betas) < 1.0, 39 | "Invalid beta parameter at index 0: ", std::get<0>(betas)); 40 | TORCH_CHECK(0 <= std::get<1>(betas) && std::get<1>(betas) < 1.0, 41 | "Invalid beta parameter at index 1: ", std::get<1>(betas)); 42 | TORCH_CHECK(defaults.weight_decay() >= 0, "Invalid weight_decay value: ", defaults.weight_decay()); 43 | } 44 | explicit MyAdam(std::vector params, AdamOptions defaults = {}) 45 | : MyAdam({OptimizerParamGroup(std::move(params))}, defaults) 46 | { 47 | } 48 | 49 | torch::Tensor step(LossClosure closure = nullptr) override; 50 | void save(serialize::OutputArchive& archive) const override; 51 | void load(serialize::InputArchive& archive) override; 52 | void shrinkenInternalState(int param_group_index, torch::Tensor indices_to_keep); 53 | void appendToInternalState(int param_group_index, int new_size); 54 | 55 | private: 56 | // template 57 | // static void serialize(Self& self, Archive& archive) 58 | //{ 59 | // _TORCH_OPTIM_SERIALIZE_WITH_TEMPLATE_ARG(Adam); 60 | // } 61 | }; 62 | } // namespace optim 63 | } // namespace torch -------------------------------------------------------------------------------- /src/lib/neat-utils/NeAT_interop.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/core/util/ini/ini.h" 9 | 10 | 11 | struct CameraBase : public ParamsBase 12 | { 13 | SAIGA_PARAM_STRUCT(CameraBase); 14 | SAIGA_PARAM_STRUCT_FUNCTIONS; 15 | template 16 | void Params(ParamIterator* it) 17 | 18 | { 19 | SAIGA_PARAM(w); 20 | SAIGA_PARAM(h); 21 | 22 | auto vector2string = [](auto vector) 23 | { 24 | std::stringstream sstrm; 25 | sstrm << std::setprecision(15) << std::scientific; 26 | for (unsigned int i = 0; i < vector.size(); ++i) 27 | { 28 | sstrm << vector[i]; 29 | if (i < vector.size() - 1) sstrm << " "; 30 | } 31 | return sstrm.str(); 32 | }; 33 | 34 | 35 | 36 | { 37 | std::vector K = split(vector2string(this->K.cast().coeffs()), ' '); 38 | SAIGA_PARAM_LIST_COMMENT(K, ' ', "# fx fy cx cy s"); 39 | SAIGA_ASSERT(K.size() == 5); 40 | 41 | Vector K_coeffs; 42 | for (int i = 0; i < 5; ++i) 43 | { 44 | double d = to_double(K[i]); 45 | K_coeffs(i) = d; 46 | } 47 | this->K = IntrinsicsPinholed(K_coeffs); 48 | } 49 | } 50 | 51 | 52 | int w = 0; 53 | int h = 0; 54 | IntrinsicsPinholed K; 55 | }; 56 | 57 | struct DatasetParams : public ParamsBase 58 | { 59 | SAIGA_PARAM_STRUCT(DatasetParams); 60 | SAIGA_PARAM_STRUCT_FUNCTIONS; 61 | // virtual void Params(Saiga::SimpleIni* ini, CLI::App* app) override 62 | template 63 | void Params(ParamIterator* it) 64 | { 65 | SAIGA_PARAM(image_dir); 66 | SAIGA_PARAM(mask_dir); 67 | SAIGA_PARAM(projection_factor); 68 | SAIGA_PARAM(vis_volume_intensity_factor); 69 | SAIGA_PARAM(scene_scale); 70 | SAIGA_PARAM(xray_min); 71 | SAIGA_PARAM(xray_max); 72 | SAIGA_PARAM(volume_file); 73 | SAIGA_PARAM(log_space_input); 74 | SAIGA_PARAM(use_log10_conversion); 75 | 76 | SAIGA_PARAM(z_min); 77 | } 78 | 79 | // only set if a ground truth volume exists 80 | std::string volume_file = ""; 81 | 82 | 83 | // linear multiplier to the projection 84 | // otherwise it is just transformed by xray/min/max parameters 85 | double projection_factor = 1; 86 | 87 | 88 | // Only for visualization! 89 | // multiplied to the intensity of the projection (after normalization) 90 | double vis_volume_intensity_factor = 1; 91 | 92 | // the camera position is multiplied by this factor to "scale" the scene 93 | double scene_scale = 1; 94 | 95 | std::string image_dir = ""; 96 | std::string mask_dir = ""; 97 | 98 | // "real" raw xray is usually NOT in log space (background is white) 99 | // if the data is already preprocessed and converted to log space (background is black) 100 | // set this flag in the dataset 101 | bool log_space_input = false; 102 | 103 | // pepper:13046, 65535 104 | // ropeball: 26000, 63600 105 | double xray_min = 0; 106 | double xray_max = 65535; 107 | 108 | double z_min = 0; 109 | 110 | // true: log10 111 | // false: loge 112 | bool use_log10_conversion = true; 113 | }; 114 | 115 | 116 | struct MyNeATTrainParams : public TrainParams 117 | { 118 | using ParamStructType = MyNeATTrainParams; 119 | // MyTrainParams() {} 120 | // MyTrainParams(const std::string file) { Load(file); } 121 | 122 | // using ParamStructType = MyTrainParams; 123 | // SAIGA_PARAM_STRUCT(MyTrainParams); 124 | 125 | MyNeATTrainParams() {} 126 | MyNeATTrainParams(const std::string file) : TrainParams(file) {} 127 | 128 | 129 | SAIGA_PARAM_STRUCT_FUNCTIONS; 130 | 131 | // SAIGA_PARAM_STRUCT_FUNCTIONS; 132 | 133 | // virtual void Params(Saiga::SimpleIni* ini, CLI::App* app) override 134 | template 135 | void Params(ParamIterator* it) 136 | { 137 | TrainParams::Params(it); 138 | 139 | SAIGA_PARAM(scene_dir); 140 | SAIGA_PARAM_LIST2(scene_name, ' '); 141 | SAIGA_PARAM(split_name); 142 | 143 | SAIGA_PARAM(optimize_structure_every_epochs); 144 | SAIGA_PARAM(optimize_structure_convergence); 145 | SAIGA_PARAM(per_node_batch_size); 146 | SAIGA_PARAM(rays_per_image); 147 | SAIGA_PARAM(output_volume_size); 148 | 149 | SAIGA_PARAM(lr_exex_grid_adam); 150 | SAIGA_PARAM(optimize_pose); 151 | SAIGA_PARAM(optimize_intrinsics); 152 | SAIGA_PARAM(lr_pose); 153 | SAIGA_PARAM(lr_intrinsics); 154 | SAIGA_PARAM(lr_decay_factor); 155 | SAIGA_PARAM(optimize_structure_after_epochs); 156 | SAIGA_PARAM(optimize_tree_structure_after_epochs); 157 | SAIGA_PARAM(optimize_tone_mapper_after_epochs); 158 | SAIGA_PARAM(init_bias_with_bg); 159 | SAIGA_PARAM(grid_init); 160 | SAIGA_PARAM(loss_tv); 161 | SAIGA_PARAM(loss_edge); 162 | SAIGA_PARAM(eval_scale); 163 | 164 | SAIGA_PARAM(experiment_name_str); 165 | SAIGA_PARAM(experiment_dir_override); 166 | 167 | SAIGA_PARAM(temp_image_dir); 168 | } 169 | 170 | std::string scene_dir = ""; 171 | std::vector scene_name = {"not_set"}; 172 | std::string split_name = "exp_uniform_50"; 173 | std::string experiment_name_str = ""; 174 | std::string experiment_dir_override = ""; 175 | std::string temp_image_dir = ""; 176 | 177 | int optimize_structure_every_epochs = 1; 178 | float optimize_structure_convergence = 0.95; 179 | 180 | std::string grid_init = "uniform"; 181 | 182 | int rays_per_image = 50000; 183 | int per_node_batch_size = 256; 184 | int output_volume_size = 256; 185 | 186 | double lr_decay_factor = 1.0; 187 | 188 | double eval_scale = 0.25; 189 | 190 | double loss_tv = 1e-4; 191 | double loss_edge = 1e-3; 192 | 193 | 194 | float lr_exex_grid_adam = 0.04; 195 | 196 | // On each image we compute the median value of a top right corner crop 197 | // This is used to initialize the tone-mapper's bias value 198 | bool init_bias_with_bg = false; 199 | 200 | // In the first few epochs we keep the camera pose/model fixed! 201 | int optimize_tree_structure_after_epochs = 1; 202 | int optimize_structure_after_epochs = 5; 203 | int optimize_tone_mapper_after_epochs = 1; 204 | bool optimize_pose = true; 205 | bool optimize_intrinsics = true; 206 | float lr_pose = 0.001; 207 | float lr_intrinsics = 100; 208 | }; 209 | 210 | struct Netparams : public ParamsBase 211 | { 212 | SAIGA_PARAM_STRUCT(Netparams); 213 | SAIGA_PARAM_STRUCT_FUNCTIONS; 214 | // virtual void Params(Saiga::SimpleIni* ini, CLI::App* app) override 215 | template 216 | void Params(ParamIterator* it) 217 | { 218 | SAIGA_PARAM(grid_size); 219 | SAIGA_PARAM(grid_features); 220 | SAIGA_PARAM(last_activation_function); 221 | SAIGA_PARAM(softplus_beta); 222 | 223 | SAIGA_PARAM(decoder_skip); 224 | SAIGA_PARAM(decoder_lr); 225 | SAIGA_PARAM(decoder_activation); 226 | SAIGA_PARAM(decoder_hidden_layers); 227 | SAIGA_PARAM(decoder_hidden_features); 228 | } 229 | 230 | int grid_size = 17; 231 | int grid_features = 4; 232 | 233 | // relu, id, abs 234 | std::string last_activation_function = "softplus"; 235 | float softplus_beta = 2; 236 | 237 | bool decoder_skip = false; 238 | float decoder_lr = 0.00005; 239 | std::string decoder_activation = "silu"; 240 | int decoder_hidden_layers = 1; 241 | int decoder_hidden_features = 64; 242 | }; 243 | 244 | 245 | struct TreeOptimizerParams 246 | { 247 | int num_threads = 16; 248 | bool use_saved_errors = true; 249 | int max_active_nodes = 512; 250 | bool verbose = false; 251 | 252 | double error_merge_factor = 1.1; 253 | double error_split_factor = 0.75; 254 | }; 255 | // Params for the HyperTree 256 | struct OctreeParams : public ParamsBase 257 | { 258 | SAIGA_PARAM_STRUCT(OctreeParams); 259 | SAIGA_PARAM_STRUCT_FUNCTIONS; 260 | // virtual void Params(Saiga::SimpleIni* ini, CLI::App* app) override 261 | template 262 | void Params(ParamIterator* it) 263 | { 264 | SAIGA_PARAM(start_layer); 265 | SAIGA_PARAM(tree_depth); 266 | SAIGA_PARAM(optimize_structure); 267 | SAIGA_PARAM(max_samples_per_node); 268 | SAIGA_PARAM(culling_start_epoch); 269 | SAIGA_PARAM(node_culling); 270 | SAIGA_PARAM(culling_threshold); 271 | 272 | 273 | SAIGA_PARAM(tree_optimizer_params.use_saved_errors); 274 | SAIGA_PARAM(tree_optimizer_params.max_active_nodes); 275 | SAIGA_PARAM(tree_optimizer_params.error_merge_factor); 276 | SAIGA_PARAM(tree_optimizer_params.error_split_factor); 277 | SAIGA_PARAM(tree_optimizer_params.verbose); 278 | } 279 | 280 | int start_layer = 3; 281 | int tree_depth = 4; 282 | bool optimize_structure = true; 283 | 284 | int max_samples_per_node = 32; 285 | 286 | int culling_start_epoch = 4; 287 | bool node_culling = true; 288 | 289 | // 0.01 for mean, 0.4 for max 290 | float culling_threshold = 0.1; 291 | 292 | 293 | TreeOptimizerParams tree_optimizer_params; 294 | }; 295 | 296 | struct PhotometricCalibrationParams : public ParamsBase 297 | { 298 | SAIGA_PARAM_STRUCT(PhotometricCalibrationParams); 299 | SAIGA_PARAM_STRUCT_FUNCTIONS; 300 | 301 | // virtual void Params(Saiga::SimpleIni* ini, CLI::App* app) override 302 | template 303 | void Params(ParamIterator* it) 304 | { 305 | SAIGA_PARAM(response_enable); 306 | SAIGA_PARAM(response_range); 307 | SAIGA_PARAM(response_lr); 308 | 309 | SAIGA_PARAM(exposure_enable); 310 | SAIGA_PARAM(exposure_mult); 311 | SAIGA_PARAM(exposure_lr); 312 | 313 | 314 | SAIGA_PARAM(sensor_bias_enable); 315 | SAIGA_PARAM(sensor_bias_size_w); 316 | SAIGA_PARAM(sensor_bias_lr); 317 | } 318 | 319 | bool response_enable = true; 320 | float response_range = 2; 321 | float response_lr = 0.1; 322 | 323 | bool exposure_enable = false; 324 | bool exposure_mult = false; 325 | float exposure_lr = 0.01; 326 | 327 | // the size in h will be computed from the aspect ratio 328 | bool sensor_bias_enable = false; 329 | int sensor_bias_size_w = 32; 330 | float sensor_bias_lr = 0.05; 331 | }; 332 | 333 | 334 | struct NeATCombinedParams 335 | { 336 | MyNeATTrainParams train_params; 337 | OctreeParams octree_params; 338 | Netparams net_params; 339 | PhotometricCalibrationParams photo_calib_params; 340 | 341 | NeATCombinedParams() {} 342 | NeATCombinedParams(const std::string& combined_file) 343 | : train_params(combined_file), 344 | octree_params(combined_file), 345 | net_params(combined_file), 346 | photo_calib_params(combined_file) 347 | { 348 | } 349 | 350 | void Save(const std::string file) 351 | { 352 | train_params.Save(file); 353 | octree_params.Save(file); 354 | net_params.Save(file); 355 | photo_calib_params.Save(file); 356 | } 357 | 358 | void Load(std::string file) 359 | { 360 | train_params.Load(file); 361 | octree_params.Load(file); 362 | net_params.Load(file); 363 | photo_calib_params.Load(file); 364 | } 365 | 366 | void Load(CLI::App& app) 367 | { 368 | train_params.Load(app); 369 | octree_params.Load(app); 370 | net_params.Load(app); 371 | photo_calib_params.Load(app); 372 | } 373 | }; -------------------------------------------------------------------------------- /src/lib/neat-utils/image_utils.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #include "image_utils.h" 8 | 9 | TemplatedImage UndistortImage(ImageView img, IntrinsicsPinholef K, Distortionf D, 10 | IntrinsicsPinholef target_K, ivec2 sizes_image, float scene_scale, 11 | unsigned int border_pixels) 12 | { 13 | TemplatedImage cp(sizes_image.y(), sizes_image.x()); 14 | cp.create(sizes_image.y(), sizes_image.x()); 15 | cp.makeZero(); 16 | #pragma omp parallel for 17 | for (int i_y = 0; i_y < cp.h; ++i_y) 18 | { 19 | for (int j_x = 0; j_x < cp.w; ++j_x) 20 | { 21 | vec2 p(j_x, i_y); 22 | 23 | p = target_K.unproject2(p); 24 | p = distortNormalizedPoint(p, D); 25 | p = K.normalizedToImage(p); 26 | 27 | p *= scene_scale; 28 | // if (p.y() < cp.h && p.y() >= 0 && p.x() < cp.w && p.x() >= 0) 29 | // if(img.inImage(ivec2(p.x(),p.y()))) 30 | // leave edges black 31 | if (p.y() >= border_pixels && p.y() < (img.height - border_pixels) && p.x() >= border_pixels && 32 | p.x() <= (img.w - border_pixels)) 33 | { 34 | cp(i_y, j_x) = img.inter(p.y(), p.x()); 35 | } 36 | } 37 | } 38 | 39 | return cp; 40 | } 41 | 42 | TemplatedImage UndistortOCAMImage(ImageView img, IntrinsicsPinholef K, IntrinsicsPinholef targetK, 43 | OCam ocam, vec2 size_of_ocam_target_image, float scene_scale, 44 | unsigned int border_pixels) 45 | { 46 | vec2 sizes_img = size_of_ocam_target_image; 47 | TemplatedImage cp(sizes_img.y(), sizes_img.x()); 48 | cp.create(sizes_img.y(), sizes_img.x()); 49 | cp.makeZero(); 50 | auto oc = ocam.cast(); 51 | 52 | #pragma omp parallel for 53 | for (int i_y = 0; i_y < cp.h; ++i_y) 54 | { 55 | for (int j_x = 0; j_x < cp.w; ++j_x) 56 | { 57 | vec3 px(j_x, i_y, 1); 58 | px = targetK.unproject(vec2(px.x(), px.y()), px.z()); 59 | vec3 p_s = px; 60 | vec2 p = oc.Project(p_s); 61 | p *= scene_scale; 62 | 63 | // leave edges black 64 | if (p.y() >= border_pixels && p.y() < (img.height - border_pixels) && p.x() >= border_pixels && 65 | p.x() <= (img.w - border_pixels)) 66 | { 67 | cp(i_y, j_x) = img.inter(p.y(), p.x()); 68 | } 69 | } 70 | } 71 | 72 | return cp; 73 | } 74 | 75 | 76 | torch::Tensor process_l1_image(torch::Tensor l1_img) 77 | { 78 | l1_img *= 1.3; 79 | l1_img -= 0.3; 80 | l1_img.clamp_(0, 1); 81 | return l1_img; 82 | } 83 | torch::Tensor process_ssim_image(torch::Tensor ssim_map) 84 | { 85 | ssim_map *= 1.3; 86 | ssim_map -= 0.3; 87 | ssim_map.clamp_(0, 1); 88 | return ssim_map; 89 | } 90 | 91 | torch::Tensor process_l2_image(torch::Tensor l2_map) 92 | { 93 | const float fac = 0.001; 94 | l2_map *= (1.f + fac); 95 | l2_map -= fac; 96 | l2_map.clamp_(0, 1); 97 | return l2_map; 98 | } 99 | 100 | void write16bitImg(TemplatedImage img, std::string path) 101 | { 102 | TemplatedImage img_16b(img.h, img.w); 103 | img_16b.create(img.h, img.w); 104 | img_16b.makeZero(); 105 | for (int y = 0; y < img.height; ++y) 106 | { 107 | for (int x = 0; x < img.width; ++x) 108 | { 109 | img_16b(y, x) = ushort(clamp(img(y, x) * 65535.f, 0, 65535.f)); 110 | } 111 | } 112 | img_16b.save(path); 113 | } -------------------------------------------------------------------------------- /src/lib/neat-utils/image_utils.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | #include "saiga/core/Core.h" 8 | #include "saiga/core/camera/HDR.h" 9 | #include "saiga/core/imgui/imgui.h" 10 | #include "saiga/core/math/CoordinateSystems.h" 11 | #include "saiga/core/sophus/Sophus.h" 12 | #include "saiga/core/util/FileSystem.h" 13 | #include "saiga/core/util/directory.h" 14 | #include "saiga/core/util/file.h" 15 | #include "saiga/core/util/tostring.h" 16 | #include "saiga/vision/VisionIncludes.h" 17 | #include "saiga/vision/cameraModel/CameraModel.h" 18 | #include "saiga/vision/cameraModel/Distortion.h" 19 | #include "saiga/vision/cameraModel/OCam.h" 20 | #include "saiga/vision/torch/TorchHelper.h" 21 | 22 | using namespace Saiga; 23 | 24 | TemplatedImage UndistortImage(ImageView img, IntrinsicsPinholef K, Distortionf D, 25 | IntrinsicsPinholef target_K, ivec2 sizes_image, float scene_scale = 1.f, 26 | unsigned int border_pixels = 16); 27 | 28 | TemplatedImage UndistortOCAMImage(ImageView img, IntrinsicsPinholef K, IntrinsicsPinholef targetK, 29 | OCam ocam, vec2 size_of_ocam_target_image, float scene_scale = 0.5f, 30 | unsigned int border_pixels = 16); 31 | 32 | torch::Tensor process_l1_image(torch::Tensor l1_img); 33 | torch::Tensor process_ssim_image(torch::Tensor ssim_map); 34 | torch::Tensor process_l2_image(torch::Tensor l2_map); 35 | void write16bitImg(TemplatedImage img, std::string path); 36 | -------------------------------------------------------------------------------- /src/lib/neat-utils/neat-base-config.ini: -------------------------------------------------------------------------------- 1 | [TrainParams] 2 | random_seed = 3746934646 3 | do_train = true 4 | do_eval = true 5 | batch_size = 16000 6 | inner_batch_size = 1 7 | inner_sample_size = 1 8 | num_epochs = 20 9 | save_checkpoints_its = 20 10 | eval_only_on_checkpoint = false 11 | name = default 12 | debug = false 13 | output_file_type = .jpg 14 | checkpoint_directory = 15 | split_method = 16 | max_images = -1 17 | duplicate_train_factor = 1 18 | shuffle_initial_indices = false 19 | shuffle_train_indices = true 20 | split_remove_neighbors = 0 21 | split_index_file_train = 22 | split_index_file_test = 23 | train_on_eval = true 24 | train_factor = 1 25 | num_workers_train = 4 26 | num_workers_eval = 4 27 | scene_dir = 28 | scene_name = 29 | split_name = 30 | optimize_structure_every_epochs = 1 31 | optimize_structure_convergence = 0.9499999881 32 | per_node_batch_size = 256 33 | rays_per_image = 50000 34 | output_volume_size = 256 35 | lr_exex_grid_adam = 0.03999999911 36 | optimize_pose = false 37 | optimize_intrinsics = false 38 | lr_pose = 0.001000000047 39 | lr_intrinsics = 100 40 | lr_decay_factor = 1 41 | optimize_structure_after_epochs = 5 42 | optimize_tree_structure_after_epochs = 1 43 | optimize_tone_mapper_after_epochs = 1 44 | init_bias_with_bg = false 45 | grid_init = uniform 46 | loss_tv = 9.999999747e-05 47 | loss_edge = 0.001 48 | eval_scale = 0.25 49 | experiment_name_str = ct_reco_ 50 | experiment_dir_override = 51 | temp_image_dir = 52 | 53 | 54 | [OctreeParams] 55 | start_layer = 3 56 | tree_depth = 4 57 | optimize_structure = true 58 | max_samples_per_node = 32 59 | culling_start_epoch = 4 60 | node_culling = true 61 | culling_threshold = 0.1000000015 62 | tree_optimizer_params.use_saved_errors = true 63 | tree_optimizer_params.max_active_nodes = 512 64 | tree_optimizer_params.error_merge_factor = 1.1 65 | tree_optimizer_params.error_split_factor = 0.75 66 | tree_optimizer_params.verbose = false 67 | 68 | 69 | [Netparams] 70 | grid_size = 33 71 | grid_features = 4 72 | last_activation_function = softplus 73 | softplus_beta = 2 74 | decoder_skip = false 75 | decoder_lr = 4.999999874e-05 76 | decoder_activation = silu 77 | decoder_hidden_layers = 1 78 | decoder_hidden_features = 64 79 | 80 | 81 | [PhotometricCalibrationParams] 82 | response_enable = true 83 | response_range = 1 84 | response_lr = 0.1000000015 85 | exposure_enable = false 86 | exposure_mult = false 87 | exposure_lr = 0.009999999776 88 | sensor_bias_enable = false 89 | sensor_bias_size_w = 32 90 | sensor_bias_lr = 0.05000000075 91 | -------------------------------------------------------------------------------- /src/lib/opengl/GridGLRenderer.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #include "GridGLRenderer.h" 7 | 8 | struct ocam_model 9 | { 10 | ivec2 image_size; 11 | float c; 12 | float d; 13 | float e; 14 | float cx; 15 | float cy; 16 | int world2cam_size; 17 | float world2cam[20]; 18 | 19 | int cam2world_size; 20 | float cam2world[20]; 21 | }; 22 | 23 | GridGLRenderer::GridGLRenderer(NeuralPointCloudCuda neural_pc) 24 | { 25 | ocam_model_ssbo.createGLBuffer(nullptr, sizeof(ocam_model), GL_DYNAMIC_DRAW); 26 | 27 | std::vector grid_mins = tensor_to_vector(neural_pc->t_cell_bb_min); 28 | // std::cout << neural_pc->t_cell_bb_min.sizes() << ": " << grid_mins.size() << std::endl; 29 | std::cout << "1" << std::endl; 30 | std::vector grid_lengths = tensor_to_vector(neural_pc->t_cell_bb_length); 31 | std::cout << "1" << std::endl; 32 | std::vector grid_values = tensor_to_vector(neural_pc->t_cell_value); 33 | std::cout << "1" << std::endl; 34 | std::vector grid_colors = neural_pc->DebugColorsPerBoxCPU(); 35 | std::cout << "1" << std::endl; 36 | std::vector pos_cube = { 37 | vec3(-1.0f, -1.0f, -1.0f), // triangle 1 : begi 38 | vec3(-1.0f, -1.0f, 1.0f), vec3(-1.0f, 1.0f, 1.0f), // triangle 1 : en)d 39 | vec3(1.0f, 1.0f, -1.0f), // triangle 2 : begi)n 40 | vec3(-1.0f, -1.0f, -1.0f), vec3(-1.0f, 1.0f, -1.0f), // triangle 2 : en)d 41 | vec3(1.0f, -1.0f, 1.0f), vec3(-1.0f, -1.0f, -1.0f), vec3(1.0f, -1.0f, -1.0f), vec3(1.0f, 1.0f, -1.0f), 42 | vec3(1.0f, -1.0f, -1.0f), vec3(-1.0f, -1.0f, -1.0f), vec3(-1.0f, -1.0f, -1.0f), vec3(-1.0f, 1.0f, 1.0f), 43 | vec3(-1.0f, 1.0f, -1.0f), vec3(1.0f, -1.0f, 1.0f), vec3(-1.0f, -1.0f, 1.0f), vec3(-1.0f, -1.0f, -1.0f), 44 | vec3(-1.0f, 1.0f, 1.0f), vec3(-1.0f, -1.0f, 1.0f), vec3(1.0f, -1.0f, 1.0f), vec3(1.0f, 1.0f, 1.0f), 45 | vec3(1.0f, -1.0f, -1.0f), vec3(1.0f, 1.0f, -1.0f), vec3(1.0f, -1.0f, -1.0f), vec3(1.0f, 1.0f, 1.0f), 46 | vec3(1.0f, -1.0f, 1.0f), vec3(1.0f, 1.0f, 1.0f), vec3(1.0f, 1.0f, -1.0f), vec3(-1.0f, 1.0f, -1.0f), 47 | vec3(1.0f, 1.0f, 1.0f), vec3(-1.0f, 1.0f, -1.0f), vec3(-1.0f, 1.0f, 1.0f), vec3(1.0f, 1.0f, 1.0f), 48 | vec3(-1.0f, 1.0f, 1.0f), vec3(1.0f, -1.0f, 1.0f)}; 49 | 50 | std::vector vb_pos; 51 | std::vector idx_buf; 52 | std::vector colors_value; 53 | std::vector colors_index; 54 | sorted_grid_values.clear(); 55 | 56 | for (int i = 0; i < grid_mins.size(); ++i) 57 | { 58 | for (int p = 0; p < pos_cube.size(); ++p) 59 | { 60 | idx_buf.push_back(vb_pos.size()); 61 | PositionIndex pi; 62 | pi.position = 63 | grid_mins[i].array() + (pos_cube[p] * 0.5 + vec3(0.5, 0.5, 0.5)).array() * grid_lengths[i].array(); 64 | pi.index = vb_pos.size(); 65 | vb_pos.push_back(pi); 66 | sorted_grid_values.push_back(grid_values[i]); 67 | colors_value.push_back(vec4(grid_values[i], grid_values[i], grid_values[i], 1)); 68 | colors_index.push_back(vec4(grid_colors[i].x(), grid_colors[i].y(), grid_colors[i].z(), 1)); 69 | } 70 | } 71 | 72 | shader = Saiga::shaderLoader.load("grid_render.glsl"); 73 | gl_grid.setDrawMode(GL_TRIANGLES); 74 | gl_grid.set(vb_pos, GL_STATIC_DRAW); 75 | 76 | gl_color.create(colors_index, GL_STATIC_DRAW); 77 | gl_grid.addExternalBuffer(gl_color, 1, 4, GL_FLOAT, GL_FALSE, sizeof(vec4), 0); 78 | 79 | gl_data.create(colors_value, GL_STATIC_DRAW); 80 | gl_grid.addExternalBuffer(gl_data, 4, 4, GL_FLOAT, GL_FALSE, sizeof(vec4), 0); 81 | std::cout << "s1" << std::endl; 82 | std::sort(sorted_grid_values.begin(), sorted_grid_values.end()); 83 | std::cout << "s1" << std::endl; 84 | } 85 | 86 | void GridGLRenderer::render(const FrameData& fd, float cutoff_val, int mode, bool cutoff_as_percent) 87 | { 88 | float cutoff = cutoff_val; 89 | if (cutoff_as_percent) 90 | { 91 | cutoff_val = 1.f - clamp(cutoff_val, 0.f, 1.0f); 92 | cutoff = sorted_grid_values[int(sorted_grid_values.size() * cutoff_val)]; 93 | } 94 | if (shader->bind()) 95 | { 96 | // glPointSize(1); 97 | // glEnable(GL_PROGRAM_POINT_SIZE); 98 | auto cam = fd.GLCamera(); 99 | 100 | shader->upload(0, mat4(mat4::Identity())); 101 | shader->upload(1, cam.view); 102 | shader->upload(2, cam.proj); 103 | 104 | mat4 v = fd.pose.inverse().matrix().cast().eval(); 105 | mat3 k = fd.K.matrix(); 106 | vec2 s(fd.w, fd.h); 107 | 108 | shader->upload(3, v); 109 | shader->upload(4, k); 110 | shader->upload(5, s); 111 | shader->upload(6, scale); 112 | 113 | vec4 dis1 = fd.distortion.Coeffs().head<4>(); 114 | vec4 dis2 = fd.distortion.Coeffs().tail<4>(); 115 | 116 | shader->upload(7, dis1); 117 | shader->upload(8, dis2); 118 | shader->upload(9, exp2(fd.exposure_value)); 119 | 120 | // render settings 121 | shader->upload(10, mode); 122 | 123 | shader->upload(11, cutoff); 124 | shader->upload(12, (int)fd.camera_model_type); 125 | 126 | ocam_model ocam_mod; 127 | ocam_mod.c = fd.ocam.c; 128 | ocam_mod.d = fd.ocam.d; 129 | ocam_mod.e = fd.ocam.e; 130 | ocam_mod.cx = fd.ocam.cx; 131 | ocam_mod.cy = fd.ocam.cy; 132 | ocam_mod.image_size = ivec2(fd.w, fd.h); 133 | ocam_mod.world2cam_size = fd.ocam.poly_world2cam.size(); 134 | 135 | for (int i = 0; i < fd.ocam.poly_world2cam.size(); ++i) 136 | { 137 | ocam_mod.world2cam[i] = fd.ocam.poly_world2cam[i]; 138 | } 139 | // std::cout << fd.ocam.poly_world2cam.size() << " _ " << fd.ocam.poly_cam2world.size() << std::endl; 140 | ocam_model_ssbo.updateBuffer(&ocam_mod, sizeof(ocam_model), 0); 141 | 142 | ocam_model_ssbo.bind(6); 143 | 144 | 145 | 146 | gl_grid.bindAndDraw(); 147 | 148 | // glDisable(GL_PROGRAM_POINT_SIZE); 149 | shader->unbind(); 150 | // ocam_model_ssbo.unbind(); 151 | } 152 | } 153 | -------------------------------------------------------------------------------- /src/lib/opengl/GridGLRenderer.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | 8 | #include "NeuralPointCloudOpenGL.h" 9 | #include "models/Pipeline.h" 10 | 11 | template 12 | std::vector tensor_to_vector(torch::Tensor t) 13 | { 14 | auto t_x = t.contiguous().cpu(); 15 | SAIGA_ASSERT(t.sizes().size() == 2); 16 | SAIGA_ASSERT(t.dtype() == torch::kFloat); 17 | uint floats_per_elem = sizeof(T) / sizeof(float); 18 | std::vector vec(t_x.numel() / floats_per_elem); 19 | // std::cout << t.sizes() << t.numel()<< std::endl; 20 | std::memcpy((float*)vec.data(), t_x.data_ptr(), t_x.numel() * sizeof(float)); 21 | // std::cout << t_x.numel() << std::endl; 22 | 23 | // std::vector vec = std::vector(t_x.data_ptr(), t_x.data_ptr()+t_x.numel()); 24 | return vec; 25 | } 26 | 27 | class GridGLRenderer : public Saiga::Object3D 28 | { 29 | public: 30 | GridGLRenderer(NeuralPointCloudCuda neural_pc); 31 | 32 | void render(const FrameData& fd, float cutoff_val, int mode, bool cutoff_as_percent = false); 33 | 34 | private: 35 | std::vector sorted_grid_values; 36 | 37 | UniformBuffer ocam_model_ssbo; 38 | 39 | std::shared_ptr shader; 40 | 41 | Saiga::VertexBuffer gl_grid; 42 | 43 | Saiga::TemplatedBuffer gl_normal = {GL_ARRAY_BUFFER}; 44 | Saiga::TemplatedBuffer gl_color = {GL_ARRAY_BUFFER}; 45 | Saiga::TemplatedBuffer gl_data = {GL_ARRAY_BUFFER}; 46 | }; 47 | -------------------------------------------------------------------------------- /src/lib/opengl/NeuralPointCloudOpenGL.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | 8 | #include "NeuralPointCloudOpenGL.h" 9 | 10 | #include "saiga/core/imgui/imgui.h" 11 | #include "saiga/core/util/ProgressBar.h" 12 | 13 | using namespace Saiga; 14 | template <> 15 | void Saiga::VertexBuffer::setVertexAttributes() 16 | { 17 | glEnableVertexAttribArray(0); 18 | // glEnableVertexAttribArray(1); 19 | glEnableVertexAttribArray(2); 20 | // glEnableVertexAttribArray(3); 21 | 22 | glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(PositionIndex), NULL); 23 | // glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, sizeof(NeuralPointVertex), (void*)(4 * sizeof(GLfloat))); 24 | glVertexAttribIPointer(2, 1, GL_INT, sizeof(PositionIndex), (void*)(3 * sizeof(GLfloat))); 25 | // glVertexAttribPointer(3, 3, GL_FLOAT, GL_FALSE, sizeof(NeuralPointVertex), (void*)(9 * sizeof(GLfloat))); 26 | } 27 | 28 | 29 | NeuralPointCloudOpenGL::NeuralPointCloudOpenGL(const UnifiedMesh& model) : NeuralPointCloud(model) 30 | { 31 | shader = Saiga::shaderLoader.load("point_render.glsl"); 32 | gl_points.setDrawMode(GL_POINTS); 33 | gl_points.set(points, GL_STATIC_DRAW); 34 | 35 | if (color.size() == points.size()) 36 | { 37 | gl_color.create(color, GL_STATIC_DRAW); 38 | gl_points.addExternalBuffer(gl_color, 1, 4, GL_FLOAT, GL_FALSE, sizeof(vec4), 0); 39 | } 40 | 41 | 42 | if (normal.size() == points.size()) 43 | { 44 | gl_normal.create(normal, GL_STATIC_DRAW); 45 | gl_points.addExternalBuffer(gl_normal, 3, 4, GL_FLOAT, GL_FALSE, sizeof(vec4), 0); 46 | } 47 | 48 | if (data.size() == points.size()) 49 | { 50 | gl_data.create(data, GL_STATIC_DRAW); 51 | gl_points.addExternalBuffer(gl_data, 4, 4, GL_FLOAT, GL_FALSE, sizeof(vec4), 0); 52 | } 53 | } 54 | 55 | 56 | void NeuralPointCloudOpenGL::render(const FrameData& fd, float scale) 57 | { 58 | if (shader->bind()) 59 | { 60 | glPointSize(1); 61 | glEnable(GL_PROGRAM_POINT_SIZE); 62 | auto cam = fd.GLCamera(); 63 | 64 | shader->upload(0, mat4(mat4::Identity())); 65 | shader->upload(1, cam.view); 66 | shader->upload(2, cam.proj); 67 | 68 | mat4 v = fd.pose.inverse().matrix().cast().eval(); 69 | mat3 k = fd.K.matrix(); 70 | vec2 s(fd.w, fd.h); 71 | 72 | shader->upload(3, v); 73 | shader->upload(4, k); 74 | shader->upload(5, s); 75 | shader->upload(6, scale); 76 | 77 | vec4 dis1 = fd.distortion.Coeffs().head<4>(); 78 | vec4 dis2 = fd.distortion.Coeffs().tail<4>(); 79 | 80 | shader->upload(7, dis1); 81 | shader->upload(8, dis2); 82 | shader->upload(9, exp2(fd.exposure_value)); 83 | 84 | // render settings 85 | shader->upload(10, render_mode); 86 | shader->upload(11, max_value); 87 | shader->upload(12, (int)fd.camera_model_type); 88 | 89 | 90 | gl_points.bindAndDraw(); 91 | 92 | glDisable(GL_PROGRAM_POINT_SIZE); 93 | shader->unbind(); 94 | } 95 | } 96 | void NeuralPointCloudOpenGL::imgui() 97 | { 98 | ImGui::InputFloat("render_max_value", &max_value); 99 | ImGui::SliderInt("Point Render Mode", &render_mode, 0, 4); 100 | ImGui::SliderInt("Point Cam Model", &render_cam_model, 0, 4); 101 | } 102 | -------------------------------------------------------------------------------- /src/lib/opengl/NeuralPointCloudOpenGL.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #include "saiga/opengl/all.h" 10 | 11 | #include "rendering/NeuralPointCloud.h" 12 | 13 | 14 | class NeuralPointCloudOpenGL : public NeuralPointCloud 15 | { 16 | public: 17 | NeuralPointCloudOpenGL(const Saiga::UnifiedMesh& model); 18 | 19 | 20 | 21 | void render(const FrameData& fd, float scale); 22 | 23 | void imgui(); 24 | 25 | std::shared_ptr shader; 26 | 27 | Saiga::VertexBuffer gl_points; 28 | Saiga::TemplatedBuffer gl_normal = {GL_ARRAY_BUFFER}; 29 | Saiga::TemplatedBuffer gl_color = {GL_ARRAY_BUFFER}; 30 | Saiga::TemplatedBuffer gl_data = {GL_ARRAY_BUFFER}; 31 | 32 | // render settings 33 | float max_value = 10; 34 | int render_mode = 0; 35 | int render_cam_model = 0; 36 | }; 37 | -------------------------------------------------------------------------------- /src/lib/opengl/RealTimeRenderer.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/cuda/imgui_cuda.h" 9 | #include "saiga/cuda/interop.h" 10 | #include "saiga/opengl/all.h" 11 | #include "saiga/opengl/rendering/deferredRendering/tone_mapper.h" 12 | #include "saiga/opengl/rendering/lighting/bloom.h" 13 | 14 | #include "config.h" 15 | #include "data/Dataset.h" 16 | #include "models/Pipeline.h" 17 | #include "opengl/SceneViewer.h" 18 | 19 | 20 | 21 | // Helper class which is able to use a pretrained network to render a neural point cloud 22 | // The output is an RGB image of the given viewpoint 23 | class RealTimeRenderer 24 | { 25 | public: 26 | RealTimeRenderer(std::shared_ptr scene); 27 | 28 | void Forward(ImageInfo fd, int rotate_result_90deg, vec3 debug_refl_dir = vec3(0, 0, 0)); 29 | 30 | void Render(ImageInfo fd, vec3 debug_refl_dir = vec3(0, 0, 0)); 31 | 32 | // flags: 33 | // 0: color 34 | // 35 | void RenderColor(ImageInfo fd, int flags, vec3 debug_refl_dir = vec3(0, 0, 0)); 36 | 37 | void SetupRenderedDebugColor(); 38 | 39 | void imgui(); 40 | 41 | void overrideForCustomCam(ImageInfo& fd, float& old_cutoff); 42 | 43 | std::shared_ptr getClosestGTImage(ImageInfo fd); 44 | 45 | struct Experiment 46 | { 47 | // full (absolute) directory of the experimenet folder 48 | std::string dir; 49 | 50 | // only the name 51 | std::string name; 52 | 53 | struct EP 54 | { 55 | // full (absolute) directory of the epxxxx folder 56 | std::string dir; 57 | 58 | // only the name. for example "ep0005" 59 | std::string name; 60 | 61 | // for example "kemenate" 62 | std::string scene_name; 63 | 64 | // ep number 65 | int ep = 0; 66 | }; 67 | std::vector eps; 68 | 69 | Experiment(std::string dir, std::string name, std::string scene_name, bool render_able = true); 70 | }; 71 | 72 | std::string experiments_base = "experiments/"; 73 | std::vector experiments; 74 | int current_ex = 0; 75 | int current_ep = 0; 76 | int current_best_gt = -1; 77 | int best_gt_counter = 0; 78 | 79 | bool mouse_on_view = false; 80 | void LoadNets(); 81 | 82 | TemplatedImage output_image; 83 | TemplatedImage output_image_ldr; 84 | std::shared_ptr output_texture, output_texture_ldr, output_color, best_gt_texture; 85 | std::shared_ptr texure_interop, color_interop; 86 | AABB custom_discard_aabb; 87 | bool reset_new_ex = false; 88 | bool render_env_map = true; 89 | NeuralPointTexture debug_color_texture = nullptr; 90 | torch::Tensor debug_color_texture_texture; 91 | 92 | std::shared_ptr scene; 93 | std::shared_ptr ns; 94 | std::shared_ptr pipeline; 95 | 96 | // The real-time camera parameters for live viewing 97 | IntrinsicsModule rt_intrinsics = nullptr; 98 | PoseModule rt_extrinsics = nullptr; 99 | 100 | torch::Tensor uv_tensor, uv_tensor_center; 101 | bool use_center_tensor = false; 102 | 103 | 104 | bool use_gl_tonemapping = false; 105 | bool use_bloom = false; 106 | bool render_color = true; 107 | 108 | // The custom camera can be controlled by the user and might be a different model 109 | bool use_custom_camera = true; 110 | 111 | float stability_cutoff_value = 0.f; 112 | float stability_cutoff_value_multiplier = 1.f; 113 | 114 | // int current_state_debug_show = 0; 115 | // std::vector debug_states = {"normal", "confidence_value"}; 116 | // int use_which_layer = -1; 117 | 118 | int discard_with_state = 0; 119 | std::vector discard_with_option = {"_via_org_confidence_value", "_via_org_confidence_value"}; 120 | int use_which_layer_discard = -1; 121 | bool invert_colors_and_red = false; 122 | // bool mask_out_original_points = false; 123 | 124 | bool use_discard_in_main_render_window = false; 125 | bool use_visualize_confidence_as_full = false; 126 | 127 | int color_layer = 1; 128 | int color_flags = 0; 129 | float color_scale = 1.f; 130 | 131 | int current_eval_epoch = 600; 132 | // int color_flags = 1; 133 | // float color_scale = 16.f; 134 | 135 | float fixed_confidence = 0.f; 136 | 137 | int manuel_timestep = -1; 138 | 139 | vec3 point_spawn_debug_color = vec3(1, 0, 1); 140 | vec3 prev_debug_col = vec3(1, 1, 1); 141 | 142 | int render_mode = 0; 143 | 144 | torch::DeviceType device = torch::kCUDA; 145 | CUDA::CudaTimerSystem timer_system; 146 | 147 | // Default saiga renderer uses 16-bit float for HDR content 148 | ToneMapper tone_mapper = {GL_RGBA32F}; 149 | Bloom bloom = {GL_RGBA32F}; 150 | std::shared_ptr params; 151 | 152 | TemplatedImage DownloadRender() 153 | { 154 | if (use_gl_tonemapping) 155 | { 156 | SAIGA_ASSERT(output_texture_ldr); 157 | TemplatedImage tmp(output_texture_ldr->getHeight(), output_texture_ldr->getWidth()); 158 | 159 | output_texture_ldr->bind(); 160 | glGetTexImage(output_texture_ldr->getTarget(), 0, GL_RGBA, GL_UNSIGNED_BYTE, tmp.data()); 161 | assert_no_glerror(); 162 | output_texture_ldr->unbind(); 163 | 164 | return tmp; 165 | } 166 | else 167 | { 168 | SAIGA_ASSERT(output_texture); 169 | TemplatedImage tmp(output_texture->getHeight(), output_texture->getWidth()); 170 | 171 | output_texture->bind(); 172 | glGetTexImage(output_texture->getTarget(), 0, GL_RGBA, GL_UNSIGNED_BYTE, tmp.data()); 173 | assert_no_glerror(); 174 | output_texture->unbind(); 175 | 176 | return tmp; 177 | } 178 | } 179 | 180 | TemplatedImage DownloadColor() 181 | { 182 | SAIGA_ASSERT(output_color); 183 | TemplatedImage tmp(output_color->getHeight(), output_color->getWidth()); 184 | 185 | output_color->bind(); 186 | glGetTexImage(output_color->getTarget(), 0, GL_RGBA, GL_UNSIGNED_BYTE, tmp.data()); 187 | assert_no_glerror(); 188 | output_color->unbind(); 189 | return tmp; 190 | } 191 | 192 | TemplatedImage DownloadGt() 193 | { 194 | SAIGA_ASSERT(best_gt_texture); 195 | TemplatedImage tmp(best_gt_texture->getHeight(), best_gt_texture->getWidth()); 196 | 197 | best_gt_texture->bind(); 198 | glGetTexImage(best_gt_texture->getTarget(), 0, GL_RGBA, GL_UNSIGNED_BYTE, tmp.data()); 199 | assert_no_glerror(); 200 | best_gt_texture->unbind(); 201 | return tmp; 202 | } 203 | }; -------------------------------------------------------------------------------- /src/lib/opengl/SceneViewer.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #include "saiga/core/glfw/all.h" 10 | 11 | #include "../config.h" 12 | #include "data/SceneData.h" 13 | #include "opengl/NeuralPointCloudOpenGL.h" 14 | 15 | using namespace Saiga; 16 | 17 | class SceneViewer : public glfw_KeyListener 18 | { 19 | public: 20 | SceneViewer(std::shared_ptr scene); 21 | 22 | void OptimizePoints(); 23 | 24 | void CreateMasks(bool mult_old_mask); 25 | 26 | void RemoveInvalidPoints(); 27 | 28 | void SetRandomPointColor(); 29 | 30 | // Render the shapes for the cameras 31 | void RenderDebug(Camera* cam); 32 | 33 | void RenderMultiscalarSpheres(Camera* cam, bool recompute, int num_spheres = 4, float inner_radius = 10.f, 34 | float radius_factor = 5.f, vec3 mid_point = vec3(0, 0, 0)); 35 | 36 | std::shared_ptr capture_asset; 37 | std::shared_ptr frustum_asset; 38 | 39 | void Select(Ray r); 40 | 41 | void UpdatePointCloudBuffer(); 42 | 43 | FrameData& Current() 44 | { 45 | SAIGA_ASSERT(selected_capture >= 0 && selected_capture < scene->frames.size()); 46 | return scene->frames[selected_capture]; 47 | } 48 | 49 | void DeleteCurrent() 50 | { 51 | if (selected_capture != -1) 52 | { 53 | scene->frames[selected_capture] = scene->frames.back(); 54 | scene->frames.pop_back(); 55 | selected_capture = -1; 56 | } 57 | } 58 | 59 | void imgui(); 60 | 61 | 62 | std::shared_ptr gl_points; 63 | 64 | std::shared_ptr scene; 65 | 66 | UnifiedModel model; 67 | 68 | CameraModel free_view_camera_type = CameraModel::PINHOLE_DISTORTION; 69 | IntrinsicsPinholef default_pinhole = IntrinsicsPinholef(1000.f, 1000.f, 1920.f / 2.f, 1080.f / 2.f, 0.f); 70 | std::vector default_cam_to_world = 71 | std::vector({-1372.79, 0, 0.00036144, -2.81704e-07, 2.64439e-10, -1.07414e-13, 1.76e-17}); 72 | std::vector default_world_to_cam = 73 | std::vector({2180.44, 1372.7, -196.498, -283.005, 139.021, 500.373, -99.8263, -687.294, 46.9305, 777.54, 74 | 120.427, -587.192, -227.447, 242.275, 158.027, -26.5849, -39.8955, -8.11842}); 75 | OCam default_ocam = OCam(5472, 3648, {1.00019, -7.17476e-05, 0.000243079, 1824.93, 2725.09}, 76 | default_cam_to_world, default_world_to_cam); 77 | 78 | IntrinsicsPinholef crop_debug_intrinsics; 79 | ivec2 crop_debug_size; 80 | Matrix crop_debug_rotation; 81 | float rot_angle = 0.f; 82 | bool use_crop_for_rendering = true; 83 | 84 | 85 | enum VIEWER_CAMERA_STATE 86 | { 87 | DEFAULT = 0, 88 | PINHOLE = 1, 89 | OCAM = 2, 90 | SPHERICAL = 3 91 | } viewer_camera = VIEWER_CAMERA_STATE::DEFAULT; 92 | 93 | 94 | ImageInfo CurrentFrameData() const 95 | { 96 | ImageInfo fd; 97 | fd.crop_rotation.setZero(); 98 | fd.crop_rotation(0, 0) = 1; 99 | fd.crop_rotation(1, 1) = 1; 100 | if (viewer_camera == VIEWER_CAMERA_STATE::DEFAULT) 101 | { 102 | fd.w = scene->scene_cameras[0].w; // * scene->dataset_params.render_scale; 103 | fd.h = scene->scene_cameras[0].h; // * scene->dataset_params.render_scale; 104 | fd.K = scene->scene_cameras[0].K; 105 | fd.camera_model_type = CameraModel::PINHOLE_DISTORTION; 106 | if (scene->scene_cameras[0].camera_model_type == CameraModel::OCAM) 107 | { 108 | // cx and cy are "switched"" in the ocam model, as it is computed with x and y flipped 109 | // fd.K = IntrinsicsPinholef(scene->scene_cameras[0].w * scene->dataset_params.render_scale, 110 | // scene->scene_cameras[0].h * scene->dataset_params.render_scale, 111 | // scene->scene_cameras[0].ocam.cy, scene->scene_cameras[0].ocam.cx, 0); 112 | 113 | fd.camera_model_type = CameraModel::OCAM; 114 | } 115 | fd.distortion = scene->scene_cameras[0].distortion; 116 | fd.ocam = scene->scene_cameras[0].ocam.cast(); 117 | } 118 | else if (viewer_camera == VIEWER_CAMERA_STATE::PINHOLE) 119 | { 120 | // fd.w = default_pinhole.cx * 2; // * scene->dataset_params.render_scale; 121 | // fd.h = default_pinhole.cy * 2; // * scene->dataset_params.render_scale; 122 | fd.w = scene->scene_cameras[0].w; 123 | fd.h = scene->scene_cameras[0].h; 124 | fd.K = IntrinsicsPinholef(scene->scene_cameras[0].w / 2, scene->scene_cameras[0].w / 2, 125 | scene->scene_cameras[0].w / 2, scene->scene_cameras[0].h / 2, 0); 126 | // fd.K = default_pinhole; 127 | fd.camera_model_type = CameraModel::PINHOLE_DISTORTION; 128 | fd.distortion = Distortionf(); 129 | } 130 | else if (viewer_camera == VIEWER_CAMERA_STATE::OCAM) 131 | { 132 | fd.w = default_ocam.w * 0.5f; 133 | fd.h = default_ocam.h * 0.5f; 134 | fd.K = default_pinhole; 135 | fd.camera_model_type = CameraModel::OCAM; 136 | // fd.crop_transform = fd.crop_transform.scale(0.5f); 137 | fd.ocam = default_ocam; 138 | } 139 | else if (viewer_camera == VIEWER_CAMERA_STATE::SPHERICAL) 140 | { 141 | fd.camera_model_type = CameraModel::SPHERICAL; 142 | fd.w = 4096; 143 | fd.h = 2048; 144 | } 145 | fd.w *= scene->dataset_params.render_scale; 146 | fd.h *= scene->dataset_params.render_scale; 147 | if (use_crop_for_rendering) 148 | { 149 | fd.crop_transform = crop_debug_intrinsics; 150 | fd.crop_rotation = crop_debug_rotation; 151 | fd.w = crop_debug_size.x(); 152 | fd.h = crop_debug_size.y(); 153 | } 154 | fd.crop_transform = fd.crop_transform.scale(scene->dataset_params.render_scale); 155 | fd.pose = Sophus::SE3f::fitToSE3(scene_camera.model * GL2CVView()).cast(); 156 | return fd; 157 | } 158 | Glfw_Camera scene_camera; 159 | 160 | // temp values 161 | int selected_capture = -1; 162 | bool render_frustums = true; 163 | 164 | bool render_custom_aabb = false; 165 | bool use_custom_aabb_to_cut_out = false; 166 | bool no_env_map = false; 167 | AABB custom_aabb; 168 | // only for ImGui 169 | vec2 max_min_custom_aabb = vec2(-3, 3); 170 | bool custom_aabb_dirty_flag = true; 171 | vec4 custom_aabb_color = vec4(1, 0, 0, 1); 172 | 173 | 174 | bool render_experiment_sphere = true; 175 | bool render_custom_sphere = false; 176 | int custom_num_spheres = 4; 177 | float custom_inner_radius = 10.f; 178 | float custom_radius_factor = 5.f; 179 | bool custom_sphere_dirty_flag = true; 180 | 181 | virtual void keyPressed(int key, int scancode, int mods) override 182 | { 183 | switch (key) 184 | { 185 | case GLFW_KEY_X: 186 | { 187 | DeleteCurrent(); 188 | break; 189 | } 190 | default: 191 | break; 192 | }; 193 | } 194 | }; 195 | -------------------------------------------------------------------------------- /src/lib/rendering/AlphaListSort.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | 9 | #include "saiga/cuda/imgui_cuda.h" 10 | 11 | #include "NeuralPointCloudCuda.h" 12 | 13 | void SegmentedSortBitonicHelper(torch::Tensor counts, torch::Tensor scanned_counts, torch::Tensor data, 14 | CUDA::CudaTimerSystem* timer); 15 | void SegmentedSortBitonicHelper2(torch::Tensor counts, torch::Tensor scanned_counts, torch::Tensor data, 16 | CUDA::CudaTimerSystem* timer = nullptr); 17 | 18 | void SegmentedSortCubHelper(torch::Tensor counts, torch::Tensor scanned_counts, torch::Tensor& data, 19 | CUDA::CudaTimerSystem* timer); -------------------------------------------------------------------------------- /src/lib/rendering/EnvironmentMap.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/cuda/imgui_cuda.h" 9 | #include "saiga/vision/torch/TorchHelper.h" 10 | 11 | #include "config.h" 12 | #include "data/NeuralStructure.h" 13 | #include "data/SceneData.h" 14 | #include "data/Settings.h" 15 | // #define LOW_ENV_MAP 16 | 17 | 18 | // This struct contains a list of rays as tensors. 19 | struct RayList 20 | { 21 | // float [num_rays, D] 22 | torch::Tensor origin; 23 | 24 | // float [num_rays, D] 25 | torch::Tensor direction; 26 | 27 | torch::Tensor alpha_dest_accumulated; 28 | 29 | int width = 0; 30 | int height = 0; 31 | int num_layers = 0; 32 | int num_batches = 0; 33 | RayList() {} 34 | 35 | // Stacks all rays into a single list 36 | RayList(const std::vector& list) 37 | { 38 | std::vector origin_list; 39 | std::vector direction_list; 40 | std::vector alpha_list; 41 | for (auto& l : list) 42 | { 43 | origin_list.push_back(l.origin); 44 | direction_list.push_back(l.direction); 45 | alpha_list.push_back(l.alpha_dest_accumulated); 46 | } 47 | origin = torch::cat(origin_list, 0); 48 | direction = torch::cat(direction_list, 0); 49 | alpha_dest_accumulated = torch::cat(alpha_list, 0); 50 | } 51 | 52 | void Allocate(int num_rays, int D) 53 | { 54 | origin = torch::empty({num_rays, D}, torch::TensorOptions().device(torch::kCUDA)); 55 | direction = torch::empty({num_rays, D}, torch::TensorOptions().device(torch::kCUDA)); 56 | alpha_dest_accumulated = torch::empty({num_rays, 1}, torch::TensorOptions().device(torch::kCUDA)); 57 | } 58 | 59 | void Allocate(int w, int h, int num_b, int num_l, int D) 60 | { 61 | width = w; 62 | height = h; 63 | num_batches = num_b; 64 | num_layers = num_l; 65 | 66 | int num_rays = 0; 67 | for (int i = 0; i < num_layers; ++i) 68 | { 69 | num_rays += w * h * num_batches; 70 | h /= 2; 71 | w /= 2; 72 | } 73 | origin = torch::empty({num_rays, D}, torch::TensorOptions().device(torch::kCUDA)); 74 | direction = torch::empty({num_rays, D}, torch::TensorOptions().device(torch::kCUDA)); 75 | alpha_dest_accumulated = torch::empty({num_rays, 1}, torch::TensorOptions().device(torch::kCUDA)); 76 | } 77 | 78 | 79 | 80 | /* ordering is: [num_batches + layer0][num_batches x layer1] ... 81 | * ____________________________________________________ 82 | * | || || || | | | | | .... 83 | * | || || || | | | | | 84 | * | || || || |-------------------- 85 | * | || || || | 86 | * -------------------------------- 87 | */ 88 | std::pair GetSubViewIndices(int batch_id, int layer_id) 89 | { 90 | std::vector sizes_of_subtensors; 91 | int w = width; 92 | int h = height; 93 | for (int i = 0; i < num_layers; ++i) 94 | { 95 | int size_wh = w * h; 96 | sizes_of_subtensors.push_back(size_wh); 97 | w /= 2; 98 | h /= 2; 99 | } 100 | int start_num_rays = 0; 101 | // get offset in layers 102 | for (int i = 0; i < layer_id; ++i) 103 | { 104 | start_num_rays += sizes_of_subtensors[i] * num_batches; 105 | } 106 | // add offset in batch num 107 | start_num_rays += batch_id * sizes_of_subtensors[layer_id]; 108 | return {start_num_rays, sizes_of_subtensors[layer_id]}; 109 | } 110 | 111 | /* 112 | int size_of_subtensor = (width / std::pow(2, layer_id)) * (height / std::pow(2, layer_id)); 113 | int start_num_rays = 0; 114 | int w = width; 115 | int h = height; 116 | for (int i = 0; i < layer_id + 1; ++i) 117 | { 118 | size_of_subtensor = w * h; 119 | w /= 2; 120 | h /= 2; 121 | } 122 | 123 | w = width; 124 | h = height; 125 | for (int i = 0; i < layer_id; ++i) 126 | { 127 | start_num_rays += num_batches * w * h; 128 | w /= 2; 129 | h /= 2; 130 | } 131 | start_num_rays += batch_id * w * h; 132 | 133 | 134 | int size_of_one_batch = 0; 135 | int size_of_index_to_layer = 0; 136 | 137 | for (int i = 0; i < num_layers; ++i) 138 | { 139 | if(i<=layer_id) 140 | size_of_index_to_layer+=(w*h); 141 | size_of_one_batch+=(w*h); 142 | h /= 2; 143 | w /= 2; 144 | } 145 | start_num_rays = batch_id*size_of_bulk + size_of_index_to_layer; 146 | 147 | for (int i = 0; i < layer_num; ++i) 148 | { 149 | for (int b = 0; b < batch_id; ++b){ 150 | start_num_rays+=(w*h)*batch_id; 151 | } 152 | h /= 2; 153 | w /= 2; 154 | } 155 | w = width; 156 | h = height; 157 | for (int i = 0; i < layer_num; ++i) 158 | { 159 | start_num_rays+=(w*h); 160 | h /= 2; 161 | w /= 2; 162 | } 163 | */ 164 | 165 | torch::Tensor getSubViewDirection(int batch_id, int layer_num) 166 | { 167 | int start, length; 168 | std::tie(start, length) = GetSubViewIndices(batch_id, layer_num); 169 | return direction.slice(0, start, start + length); 170 | } 171 | torch::Tensor getSubViewOrigin(int batch_id, int layer_num) 172 | { 173 | int start, length; 174 | std::tie(start, length) = GetSubViewIndices(batch_id, layer_num); 175 | return origin.slice(0, start, start + length); 176 | } 177 | torch::Tensor getSubViewAlphaDest(int batch_id, int layer_num) 178 | { 179 | int start, length; 180 | std::tie(start, length) = GetSubViewIndices(batch_id, layer_num); 181 | return alpha_dest_accumulated.slice(0, start, start + length); 182 | } 183 | 184 | RayList getSubRayListView(int batch_id, int layer_num) 185 | { 186 | RayList rl = RayList(); 187 | rl.origin = getSubViewOrigin(batch_id, layer_num); 188 | rl.direction = getSubViewDirection(batch_id, layer_num); 189 | rl.alpha_dest_accumulated = getSubViewAlphaDest(batch_id, layer_num); 190 | rl.width = width; 191 | rl.height = height; 192 | rl.num_layers = num_layers; 193 | rl.num_batches = num_batches; 194 | 195 | return rl; 196 | } 197 | void to(torch::Device device) 198 | { 199 | // if(linear_pixel_location.defined()) linear_pixel_location = linear_pixel_location.to(device); 200 | origin = origin.to(device); 201 | direction = direction.to(device); 202 | alpha_dest_accumulated = alpha_dest_accumulated.to(device); 203 | 204 | // if (pixel_uv.defined()) pixel_uv = pixel_uv.to(device); 205 | } 206 | 207 | RayList SubSample(torch::Tensor index) 208 | { 209 | RayList result; 210 | // result.linear_pixel_location = torch::index_select(linear_pixel_location, 0, index); 211 | result.origin = torch::index_select(origin, 0, index); 212 | result.direction = torch::index_select(direction, 0, index); 213 | result.alpha_dest_accumulated = torch::index_select(alpha_dest_accumulated, 0, index); 214 | return result; 215 | } 216 | 217 | size_t Memory() 218 | { 219 | return direction.numel() * sizeof(float) + origin.numel() * sizeof(float); 220 | // +linear_pixel_location.numel() * sizeof(long); 221 | } 222 | 223 | int size() const { return origin.size(0); } 224 | 225 | int Dim() const { return origin.size(1); } 226 | 227 | template 228 | std::pair, Eigen::Vector> GetRay(int i) const 229 | { 230 | SAIGA_ASSERT(i < size()); 231 | SAIGA_ASSERT(D == Dim()); 232 | Eigen::Vector o; 233 | Eigen::Vector d; 234 | for (int k = 0; k < D; ++k) 235 | { 236 | o(k) = origin.template data_ptr()[i * origin.stride(0) + k * origin.stride(1)]; 237 | d(k) = direction.template data_ptr()[i * direction.stride(0) + k * direction.stride(1)]; 238 | } 239 | 240 | return {o, d}; 241 | } 242 | friend std::ostream& operator<<(std::ostream& stream, const RayList& rays); 243 | }; 244 | 245 | 246 | class EnvironmentMapImpl : public torch::nn::Module 247 | { 248 | public: 249 | // up_axis: int: 0,1,2 for the axis used as up. 250 | EnvironmentMapImpl(int channels, int h, int w, bool log_texture, int axis = 0, int num_images = 4, 251 | float inner_radius = 20.f, float radius_factor = 5.f, bool non_subzero_texture = false); 252 | 253 | void CreateEnvMapOptimizer(float env_col_learning_rate = 0.02f, float env_density_learning_rate = 0.001f); 254 | 255 | 256 | // Samples the env. map in all layers and all images of the batch. 257 | // The result is an array of tensor where each element resebles one layer of the stack. 258 | std::vector Sample(torch::Tensor poses, torch::Tensor intrinsics, 259 | ArrayView info_batch, int num_layers, 260 | std::shared_ptr scene, 261 | std::vector> layers_cuda, 262 | CUDA::CudaTimerSystem* timer_system = nullptr); 263 | 264 | std::vector Sample2(torch::Tensor poses, torch::Tensor intrinsics, 265 | ArrayView info_batch, int num_layers, 266 | std::shared_ptr scene, 267 | std::vector> layers_cuda, 268 | CUDA::CudaTimerSystem* timer_system = nullptr); 269 | 270 | 271 | // returns for a give ray the color 272 | // return: [3, num_rays] 273 | torch::Tensor forward_mps(RayList rays, torch::Tensor alpha_dest_weights, 274 | CUDA::CudaTimerSystem* timer_system = nullptr); 275 | torch::Tensor forward_mps2(RayList rays, CUDA::CudaTimerSystem* timer_system = nullptr); 276 | torch::Tensor forward_mps3(RayList rays, CUDA::CudaTimerSystem* timer_system = nullptr); 277 | torch::Tensor forward_mps4(RayList rays, CUDA::CudaTimerSystem* timer_system = nullptr); 278 | torch::Tensor forward_mps5(RayList rays, CUDA::CudaTimerSystem* timer_system = nullptr); 279 | 280 | std::shared_ptr optimizer_adam; 281 | 282 | // [1, 1, num_images, resolution, resolution] 283 | torch::Tensor density; 284 | // [1, num_desc, num_images, resolution, resolution] 285 | torch::Tensor color; 286 | 287 | int NumImages() { return density.size(2); } 288 | 289 | private: 290 | // the radii of the spheres 291 | // length == num_images (from constructor) 292 | std::vector radii; 293 | int up_axis; 294 | int channels; 295 | bool non_subzero_texture; 296 | }; 297 | TORCH_MODULE(EnvironmentMap); 298 | -------------------------------------------------------------------------------- /src/lib/rendering/NeuralPointCloud.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/core/math/math.h" 9 | #include "saiga/core/sophus/Sophus.h" 10 | 11 | #include "config.h" 12 | #include "data/SceneData.h" 13 | 14 | using namespace Saiga; 15 | 16 | struct SAIGA_ALIGN(16) PositionIndex 17 | { 18 | vec3 position; 19 | int index; 20 | }; 21 | 22 | class NeuralPointCloud : public Saiga::Object3D 23 | { 24 | public: 25 | NeuralPointCloud(const Saiga::UnifiedMesh& model) 26 | { 27 | for (int i = 0; i < model.NumVertices(); ++i) 28 | { 29 | PositionIndex npv; 30 | npv.position = model.position[i]; 31 | npv.index = i; 32 | points.push_back(npv); 33 | 34 | if (model.NumVertices() == model.normal.size()) 35 | { 36 | normal.push_back(make_vec4(model.normal[i].normalized(), 0)); 37 | } 38 | if (model.NumVertices() == model.color.size()) 39 | { 40 | color.push_back(model.color[i]); 41 | } 42 | 43 | if (model.NumVertices() == model.data.size()) 44 | { 45 | data.push_back(model.data[i]); 46 | } 47 | } 48 | } 49 | 50 | std::vector points; 51 | std::vector normal; 52 | std::vector color; 53 | std::vector data; 54 | }; 55 | -------------------------------------------------------------------------------- /src/lib/rendering/NeuralPointCloudCuda.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/cuda/cuda.h" 9 | #include "saiga/cuda/thrust_helper.h" 10 | #include "saiga/normal_packing.h" 11 | #include "saiga/vision/torch/TorchHelper.h" 12 | 13 | #include "NeuralPointCloud.h" 14 | 15 | #include 16 | 17 | #include "cuda_fp16.h" 18 | 19 | 20 | 21 | class NeuralPointCloudCudaImpl : public NeuralPointCloud, public torch::nn::Module 22 | { 23 | public: 24 | NeuralPointCloudCudaImpl(const Saiga::UnifiedMesh& model, bool use_grid_loss = false, float cell_ws_size = 1.f, 25 | AABB custom_point_manip_aabb = AABB()); 26 | 27 | void MakeOutlier(int max_index); 28 | 29 | std::vector Indices(); 30 | void SetIndices(std::vector& indices); 31 | 32 | /* Following [Schütz et al. 22], 80*128 conseq points */ 33 | void UpdateCellStructureForRendering(size_t conseq_points = 10240); 34 | 35 | void UpdateCellStructureForPointOptim(float size_of_box_in_ws, AABB custom_aabb = AABB()); 36 | 37 | torch::Tensor GetPerPointBBIndex(); 38 | torch::Tensor GetPerPointBBValue(); 39 | torch::Tensor DebugBBIndexToCol(); 40 | 41 | std::vector DebugColorsPerBoxCPU(); 42 | void ResetCellValues() 43 | { 44 | t_cell_value.set_(torch::zeros_like(t_cell_value)); 45 | t_cell_access_count.set_(torch::zeros_like(t_cell_access_count)); 46 | } 47 | void NormalizeBBCellValue() 48 | { 49 | t_cell_access_count = 50 | torch::where(t_cell_access_count != 0, t_cell_access_count, torch::ones_like(t_cell_access_count)); 51 | t_cell_value /= t_cell_access_count; 52 | } 53 | 54 | void RemoveSelected(torch::Tensor to_keep) 55 | { 56 | torch::NoGradGuard ngg; 57 | std::cout << TensorInfo(to_keep) << std::endl; 58 | std::cout << TensorInfo(t_position) << std::endl; 59 | t_position = t_position.index({to_keep}); 60 | std::cout << TensorInfo(t_position) << std::endl; 61 | std::cout << TensorInfo(t_normal) << std::endl; 62 | t_normal = t_normal.index({to_keep}); 63 | std::cout << TensorInfo(t_normal) << std::endl; 64 | std::cout << TensorInfo(t_original_color) << std::endl; 65 | t_original_color = t_original_color.index({to_keep}); 66 | std::cout << TensorInfo(t_original_color) << std::endl; 67 | std::cout << TensorInfo(t_index) << std::endl; 68 | t_index = t_index.index({to_keep}); 69 | std::cout << TensorInfo(t_index) << std::endl; 70 | std::cout << TensorInfo(t_original_index) << std::endl; 71 | t_original_index = t_original_index.index({to_keep}); 72 | std::cout << TensorInfo(t_original_index) << std::endl; 73 | } 74 | 75 | torch::Tensor DebugColorsPerBox(); 76 | // void Reorder(torch::Tensor indices); 77 | 78 | int Size(); 79 | Saiga::UnifiedMesh Mesh(); 80 | 81 | // [n, 4] 82 | torch::Tensor t_position; 83 | 84 | // [n, 4] 85 | torch::Tensor t_position_displacement; 86 | 87 | // [n, 4] 88 | torch::Tensor t_normal; 89 | // torch::Tensor t_normal_test; 90 | 91 | // [n, 4] 92 | torch::Tensor t_original_color; 93 | 94 | // [n, 1] 95 | torch::Tensor t_index; 96 | // [n, 1] 97 | torch::Tensor t_original_index; 98 | 99 | // [cell_n,3] float 100 | torch::Tensor t_cell_bb_min; 101 | // [cell_n,3] float 102 | torch::Tensor t_cell_bb_length; 103 | // [cell_n,1] float 104 | torch::Tensor t_cell_value; 105 | // [cell_n,1] int 106 | torch::Tensor t_cell_access_count; 107 | 108 | using PointType = vec4; 109 | using NormalType = vec4; 110 | }; 111 | 112 | 113 | TORCH_MODULE(NeuralPointCloudCuda); 114 | 115 | 116 | // A simple helper class to make the kernels more compact. 117 | struct DevicePointCloud 118 | { 119 | int4* position; 120 | int4* position_displacement; 121 | // int4* normal_test; 122 | // half2* normal; 123 | int* normal; 124 | int* index; 125 | 126 | int3* cell_bb_min; 127 | int3* cell_bb_length; 128 | float* cell_value; 129 | int* cell_access_count; 130 | 131 | int n; 132 | int n_cells; 133 | 134 | DevicePointCloud() = default; 135 | 136 | DevicePointCloud(NeuralPointCloudCuda pc) 137 | { 138 | // SAIGA_ASSERT(pc->t_position.size(0) == pc->t_index.size(0)); 139 | SAIGA_ASSERT(pc->t_position.size(0) == pc->Size()); 140 | 141 | position = (int4*)pc->t_position.data_ptr(); 142 | if (pc->t_normal.defined()) 143 | { 144 | SAIGA_ASSERT(pc->t_position.size(0) == pc->t_normal.size(0)); 145 | normal = pc->t_normal.data_ptr(); 146 | } 147 | else 148 | { 149 | normal = nullptr; 150 | } 151 | 152 | if (pc->t_position_displacement.defined()) 153 | { 154 | SAIGA_ASSERT(pc->t_position.size(0) == pc->t_position_displacement.size(0)); 155 | position_displacement = (int4*)pc->t_position_displacement.data_ptr(); 156 | } 157 | else 158 | { 159 | position_displacement = nullptr; 160 | } 161 | 162 | index = (int*)pc->t_index.data_ptr(); 163 | 164 | n = pc->Size(); 165 | 166 | if (pc->t_cell_bb_min.defined()) 167 | { 168 | n_cells = pc->t_cell_bb_min.sizes()[0]; 169 | 170 | cell_bb_min = (int3*)pc->t_cell_bb_min.data_ptr(); 171 | cell_bb_length = (int3*)pc->t_cell_bb_length.data_ptr(); 172 | cell_value = pc->t_cell_value.data_ptr(); 173 | cell_access_count = pc->t_cell_access_count.data_ptr(); 174 | } 175 | 176 | n_cells = pc->t_cell_bb_min.size(0); 177 | } 178 | HD inline thrust::tuple GetCellBB(int cell_id) 179 | { 180 | vec3 bb_min; 181 | vec3 bb_len; 182 | float val; 183 | reinterpret_cast(&bb_min)[0] = cell_bb_min[cell_id]; 184 | reinterpret_cast(&bb_len)[0] = cell_bb_length[cell_id]; 185 | reinterpret_cast(&val)[0] = cell_value[cell_id]; 186 | 187 | return {bb_min, bb_len, val}; 188 | } 189 | HD inline void SetValueForCell(int cell_id, float val) { cell_value[cell_id] = val; } 190 | HD inline float* GetPointerForValueForCell(int cell_id) { return &cell_value[cell_id]; } 191 | HD inline int* GetPointerForAccessCountForCell(int cell_id) { return &cell_access_count[cell_id]; } 192 | 193 | HD inline thrust::tuple GetPoint(int point_index) 194 | { 195 | vec4 p; 196 | vec4 p_displacement = vec4(0, 0, 0, 0); 197 | vec4 n_test; 198 | 199 | // float4 global memory loads are vectorized! 200 | reinterpret_cast(&p)[0] = position[point_index]; 201 | 202 | if (position_displacement) 203 | { 204 | reinterpret_cast(&p_displacement)[0] = position_displacement[point_index]; 205 | } 206 | 207 | vec3 n; 208 | 209 | if (normal) 210 | { 211 | auto enc = normal[point_index]; 212 | n = UnpackNormal10Bit(enc); 213 | } 214 | 215 | float drop_out_radius = p(3); 216 | 217 | vec3 pos = p.head<3>(); 218 | if (position_displacement) 219 | { 220 | pos += p_displacement.head<3>(); 221 | } 222 | 223 | return {pos, n.head<3>(), drop_out_radius}; 224 | } 225 | 226 | HD inline int GetIndex(int tid) { return index[tid]; } 227 | 228 | HD inline void SetIndex(int tid, int value) { index[tid] = value; } 229 | 230 | HD inline int Size() { return n; } 231 | }; -------------------------------------------------------------------------------- /src/lib/rendering/PointBlending.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | #pragma once 7 | #include "saiga/core/math/Types.h" 8 | 9 | #include "config.h" 10 | using namespace Saiga; 11 | 12 | inline HD vec4 compute_blending_fac(vec2 uv_pos, Matrix* J_uv = nullptr) 13 | { 14 | // derivative is 1 in relevant area 15 | // vec2 subpixel_pos = uv_pos - ivec2(__float2int_rd(uv_pos(0)), __float2int_rd(uv_pos(1))); 16 | vec2 subpixel_pos = uv_pos - uv_pos.array().floor(); 17 | 18 | vec4 blend_vec; 19 | blend_vec.setZero(); 20 | blend_vec(0) = (1 - subpixel_pos.x()) * (1 - subpixel_pos.y()); 21 | blend_vec(1) = subpixel_pos.x() * (1 - subpixel_pos.y()); 22 | blend_vec(2) = (1 - subpixel_pos.x()) * subpixel_pos.y(); 23 | blend_vec(3) = subpixel_pos.x() * subpixel_pos.y(); 24 | 25 | if (J_uv) 26 | { 27 | auto& J = *J_uv; 28 | J.setZero(); 29 | J(0, 0) = subpixel_pos.y() - 1; 30 | J(0, 1) = subpixel_pos.x() - 1; 31 | 32 | J(1, 0) = 1 - subpixel_pos.y(); 33 | J(1, 1) = -subpixel_pos.x(); 34 | 35 | 36 | J(2, 0) = -subpixel_pos.y(); 37 | J(2, 1) = 1 - subpixel_pos.x(); 38 | 39 | J(3, 0) = subpixel_pos.y(); 40 | J(3, 1) = subpixel_pos.x(); 41 | } 42 | 43 | return blend_vec; 44 | } 45 | 46 | 47 | // #define CHANNELS 1 48 | template 49 | inline HD desc_vec compute_blend_vec(float alpha_dest, float alpha_s, desc_vec color, desc_vec color_dest, 50 | Saiga::Matrix* J_alphasource = nullptr, 51 | Saiga::Matrix* J_color = nullptr, 52 | Saiga::Matrix* J_alphadest = nullptr, 53 | Saiga::Matrix* J_colordest = nullptr) 54 | { 55 | desc_vec blended_col = alpha_dest * alpha_s * color + color_dest; 56 | 57 | if (J_alphadest) 58 | { 59 | auto& J = *J_alphadest; 60 | for (int i = 0; i < size_of_desc_vec; ++i) 61 | { 62 | J(i, 0) = alpha_s * color[i]; 63 | } 64 | } 65 | if (J_alphasource) 66 | { 67 | auto& J = *J_alphasource; 68 | for (int i = 0; i < size_of_desc_vec; ++i) 69 | { 70 | J(i, 0) = alpha_dest * color[i]; 71 | } 72 | } 73 | if (J_color) 74 | { 75 | auto& J = *J_color; 76 | for (int i = 0; i < size_of_desc_vec; ++i) 77 | { 78 | J(i, i) = alpha_dest * alpha_s; 79 | } 80 | } 81 | if (J_colordest) 82 | { 83 | auto& J = *J_colordest; 84 | for (int i = 0; i < size_of_desc_vec; ++i) J(i, i) = 1; 85 | } 86 | 87 | return blended_col; 88 | } 89 | 90 | inline HD float compute_blend(float alpha_dest, float alpha_s, float color, float color_dest, 91 | float* J_alphasource = nullptr, float* J_color = nullptr, float* J_alphadest = nullptr, 92 | float* J_colordest = nullptr) 93 | { 94 | float blended_col = alpha_dest * alpha_s * color + color_dest; 95 | 96 | if (J_alphadest) 97 | { 98 | *J_alphadest = alpha_s * color; 99 | } 100 | if (J_alphasource) 101 | { 102 | *J_alphasource = alpha_dest * color; 103 | } 104 | if (J_color) 105 | { 106 | *J_color = alpha_dest * alpha_s; 107 | } 108 | if (J_colordest) 109 | { 110 | *J_colordest = 1; 111 | } 112 | 113 | return blended_col; 114 | } 115 | 116 | inline HD float compute_new_alphadest(float alphadest_old, float alpha_s, float* J_alphasource = nullptr, 117 | float* J_alphadest_old = nullptr) 118 | { 119 | float new_alphadest = (1 - alpha_s) * alphadest_old; 120 | if (J_alphadest_old) 121 | { 122 | *J_alphadest_old = (1 - alpha_s); 123 | } 124 | if (J_alphasource) 125 | { 126 | *J_alphasource = -alphadest_old; 127 | } 128 | return new_alphadest; 129 | } 130 | 131 | 132 | 133 | #define CHANNELS 1 134 | inline HD float compute_blend_d(float alpha_dest, float alpha_s, float color, float color_dest, 135 | Saiga::Matrix* J_alphasource = nullptr, 136 | Saiga::Matrix* J_color = nullptr, 137 | Saiga::Matrix* J_alphadest = nullptr, 138 | Saiga::Matrix* J_colordest = nullptr) 139 | { 140 | float blended_col = alpha_dest * alpha_s * color + color_dest; 141 | 142 | if (J_alphadest) 143 | { 144 | auto& J = *J_alphadest; 145 | J(0, 0) = alpha_s * color; 146 | } 147 | if (J_alphasource) 148 | { 149 | auto& J = *J_alphasource; 150 | J(0, 0) = alpha_dest * color; 151 | } 152 | if (J_color) 153 | { 154 | auto& J = *J_color; 155 | J(0, 0) = alpha_dest * alpha_s; 156 | } 157 | if (J_colordest) 158 | { 159 | auto& J = *J_colordest; 160 | J(0, 0) = 1; 161 | } 162 | 163 | return blended_col; 164 | } 165 | 166 | inline HD float compute_new_alphadest_d(float alphadest_old, float alpha_s, 167 | Saiga::Matrix* J_alphasource = nullptr, 168 | Saiga::Matrix* J_alphadest_old = nullptr) 169 | { 170 | float new_alphadest = (1 - alpha_s) * alphadest_old; 171 | if (J_alphadest_old) 172 | { 173 | auto& J = *J_alphadest_old; 174 | J(0, 0) = (1 - alpha_s); 175 | } 176 | if (J_alphasource) 177 | { 178 | auto& J = *J_alphasource; 179 | J(0, 0) = -alphadest_old; 180 | } 181 | return new_alphadest; 182 | } 183 | 184 | 185 | template 186 | inline HD float normalize_by_alphadest(float color, float alphadest, Saiga::Matrix* J_alphadest = nullptr, 187 | Saiga::Matrix* J_color = nullptr) 188 | { 189 | SAIGA_ASSERT(false); // not implemented 190 | float result = color / (1 - alphadest); 191 | if (J_alphadest) 192 | { 193 | auto& J = *J_alphadest; 194 | J(0, 0) = color / ((1 - alphadest) * (1 - alphadest)); 195 | } 196 | if (J_color) 197 | { 198 | auto& J = *J_color; 199 | J(0, 0) = 1.0 / (1 - alphadest); 200 | } 201 | return result; 202 | } -------------------------------------------------------------------------------- /src/lib/rendering/PointRenderer.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/cuda/imageProcessing/image.h" 9 | #include "saiga/cuda/imgui_cuda.h" 10 | 11 | #include "NeuralPointCloudCuda.h" 12 | #include "RenderConstants.h" 13 | #include "RenderInfo.h" 14 | #include "config.h" 15 | #include "data/Dataset.h" 16 | #include "data/NeuralScene.h" 17 | #include "data/Settings.h" 18 | 19 | 20 | 21 | class PointRendererCache; 22 | 23 | class NeuralRenderInfo : public torch::CustomClassHolder 24 | { 25 | public: 26 | NeuralScene* scene; 27 | std::vector images; 28 | RenderParams params; 29 | int num_layers; 30 | bool train; 31 | int current_epoch; 32 | CUDA::CudaTimerSystem* timer_system = nullptr; 33 | PointRendererCache* cache = nullptr; 34 | }; 35 | 36 | 37 | namespace torch::autograd 38 | { 39 | struct PointRender : public Function 40 | { 41 | // returns a tensor for every layer 42 | static variable_list forward(AutogradContext* ctx, Variable texture, Variable background_color, Variable points, 43 | Variable pose_tangents, Variable intrinsics, Variable confidence, 44 | Variable dynamic_refinement, IValue info); 45 | 46 | static variable_list backward(AutogradContext* ctx, variable_list grad_output); 47 | }; 48 | } // namespace torch::autograd 49 | 50 | 51 | // Render the scene into a batch of images 52 | // Every image is a pyramid of layers in different resolutions 53 | std::vector BlendPointCloud(NeuralRenderInfo* info); 54 | 55 | 56 | // ==== Internal ==== 57 | std::pair, std::vector> BlendPointCloudForward( 58 | torch::autograd::AutogradContext* ctx, NeuralRenderInfo* info); 59 | // derivative towards the texture 60 | torch::autograd::variable_list BlendPointCloudBackward(torch::autograd::AutogradContext* ctx, NeuralRenderInfo* info, 61 | torch::autograd::variable_list image_gradients); 62 | 63 | 64 | void ApplyTangentToPose(torch::Tensor tangent, torch::Tensor pose); 65 | 66 | struct LayerCuda 67 | { 68 | ivec2 size = ivec2(0, 0); 69 | float scale = 1; 70 | 71 | 72 | ImageView BatchView(int batch) 73 | { 74 | return ImageView(size(1), size(0), 75 | depth_index_tensor.data_ptr() + batch * depth_index_tensor.stride(0)); 76 | } 77 | 78 | ImageView BatchViewDepth(int batch) 79 | { 80 | return ImageView(size(1), size(0), depth.data_ptr() + batch * depth.stride(0)); 81 | } 82 | ImageView BatchViewCounting(int batch) 83 | { 84 | return ImageView(size(1), size(0), counting.data_ptr() + batch * counting.stride(0)); 85 | } 86 | ImageView BatchViewScannedCounting(int batch) 87 | { 88 | return ImageView(size(1), size(0), scanned_counting.data_ptr() + batch * scanned_counting.stride(0)); 89 | } 90 | ImageView BatchViewWeights(int batch) 91 | { 92 | return ImageView(size(1), size(0), weight.data_ptr() + batch * weight.stride(0)); 93 | } 94 | 95 | 96 | // for new rendering 97 | torch::Tensor depth; 98 | torch::Tensor weight; 99 | 100 | torch::Tensor max_depth; 101 | torch::Tensor counting; 102 | // an ascending count structure, using a scan 103 | torch::Tensor scanned_counting; 104 | 105 | torch::Tensor per_image_atomic_counters; 106 | 107 | 108 | // for old rendering 109 | torch::Tensor depth_index_tensor; 110 | }; 111 | 112 | class PointRendererCache 113 | { 114 | public: 115 | PointRendererCache() {} 116 | 117 | void Build(NeuralRenderInfo* info, bool forward); 118 | 119 | // Allocates cuda memory in tensors. Does not initialize them!!! 120 | // Call InitializeData() below for that 121 | void Allocate(NeuralRenderInfo* info, bool forward); 122 | void InitializeData(bool forward); 123 | 124 | void PushParametersForward(); 125 | void PushParametersBackward(); 126 | DeviceRenderParams PrepareDeviceRenderParams(); 127 | DeviceTexture PrepareDeviceTexture(); 128 | DeviceBackwardParams PrepareDeviceBackwardParams(); 129 | 130 | void ProjectPoints(int batch, NeuralPointCloudCuda point_cloud); 131 | void DepthPrepassMulti(int batch, NeuralPointCloudCuda point_cloud, bool train); 132 | void RenderForwardMulti(int batch, NeuralPointCloudCuda point_cloud, bool train); 133 | 134 | void CountingPrepassMulti(int batch, NeuralPointCloudCuda point_cloud, bool train); 135 | void CollectMulti(int batch, NeuralPointCloudCuda point_cloud, std::vector collection_buffers, 136 | std::vector data_buffer, bool train); 137 | void UploadCollectionBuffers(std::vector buffers, std::vector data_buffers, 138 | int batch_num, std::vector atomics = std::vector()); 139 | void SortMulti(int batch, std::vector collection_buffers, bool train); 140 | void BlendMulti(int batch, DevicePointCloud point_cloud, std::vector collection_buffers, 141 | std::vector data_buffer, torch::Tensor background_color, bool train, 142 | bool use_environment_map); 143 | 144 | void BlendMultiFuzzy(int batch, DevicePointCloud point_cloud, std::vector collection_buffers, 145 | std::vector data_buffer, torch::Tensor background_color, bool train, 146 | bool use_environment_map); 147 | 148 | 149 | void CountingPrepassMultiBilinear(int batch, NeuralPointCloudCuda point_cloud, bool train); 150 | 151 | void CollectMultiBilinear(int batch, NeuralPointCloudCuda point_cloud, 152 | std::vector collection_buffers, std::vector data_buffer, 153 | bool train); 154 | void BlendMultiBilinear(int batch, DevicePointCloud point_cloud, std::vector collection_buffers, 155 | std::vector data_buffer, torch::Tensor background_color, bool train, 156 | bool use_environment_map); 157 | 158 | void CombineAndFillBlend(int batch, torch::Tensor background_color); 159 | 160 | void BlendBackwards(int batch, NeuralPointCloudCuda point_cloud, std::vector collection_buffers, 161 | torch::Tensor background_color, std::vector grad_sum_back_buffers, 162 | bool use_environment_map); 163 | void BlendBackwardsFuzzy(int batch, NeuralPointCloudCuda point_cloud, std::vector collection_buffers, 164 | torch::Tensor background_color, std::vector grad_sum_back_buffers, 165 | bool use_environment_map); 166 | void BlendBackwardsBilinear(int batch, NeuralPointCloudCuda point_cloud, 167 | std::vector collection_buffers, 168 | std::vector per_point_data_buffer, torch::Tensor background_color, 169 | std::vector grad_sum_back_buffers, bool use_environment_map); 170 | 171 | void UploadCollectionBuffersBackwards(std::vector buffers, std::vector data_buffer, 172 | std::vector grad_sum_back_buffers, int batch_num); 173 | 174 | 175 | void FillSceneGridWithLoss(int batch, NeuralPointCloudCuda point_cloud, std::vector images, 176 | torch::Tensor loss_img); 177 | 178 | void RenderBackward(int batch, NeuralPointCloudCuda point_cloud); 179 | 180 | void CreateMask(int batch, float background_value); 181 | 182 | void PartialTextureTransferForwardMulti(int batch, NeuralPointCloudCuda point_cloud, bool train); 183 | 184 | void CombineAndFill(int batch, torch::Tensor background_color); 185 | void CombineAndFillBackward(int batch, torch::Tensor background_color, std::vector gradient); 186 | 187 | enum RenderMode 188 | { 189 | FUZZY_DT = 0, 190 | FULL_BLEND = 1, 191 | FUZZY_BLEND = 2, 192 | BILINEAR_BLEND = 3, 193 | SIZE = 4 194 | } render_mode; 195 | 196 | 197 | 198 | std::vector layers_cuda; 199 | NeuralRenderInfo* info; 200 | int num_batches; 201 | 202 | torch::Tensor cells_to_cull; 203 | torch::Tensor random_numbers; 204 | 205 | torch::Tensor gradient_of_forward_pass_x; 206 | 207 | torch::Tensor l1_error_image; 208 | 209 | 210 | 211 | // [batch, num_points] 212 | torch::Tensor dropout_points; 213 | 214 | std::vector output_forward; 215 | std::vector output_forward_depthbuffer; 216 | std::vector output_forward_blend; 217 | 218 | // [batches, num_points, 2] 219 | torch::Tensor tmp_point_projections; 220 | 221 | torch::Tensor output_gradient_texture; 222 | torch::Tensor output_gradient_confidence; 223 | torch::Tensor output_gradient_background; 224 | torch::Tensor output_gradient_points; 225 | 226 | torch::Tensor output_gradient_pose_tangent; 227 | torch::Tensor output_gradient_pose_tangent_count; 228 | torch::Tensor output_gradient_point_count; 229 | 230 | torch::Tensor output_gradient_intrinsics; 231 | torch::Tensor output_gradient_intrinsics_count; 232 | 233 | torch::Tensor output_gradient_dynamic_points; 234 | torch::Tensor output_gradient_dynamic_point_count; 235 | 236 | 237 | //std::vector point_dynamic_gradients; 238 | 239 | std::vector image_gradients; 240 | 241 | 242 | 243 | // Variables used to check if we have to reallocate 244 | // [num_points, layers, batch_size, h, w] 245 | std::vector cache_size = {0, 0, 0, 0}; 246 | bool cache_has_forward = false; 247 | bool cache_has_backward = false; 248 | }; -------------------------------------------------------------------------------- /src/lib/rendering/RenderConstants.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "saiga/cuda/imageProcessing/image.h" 3 | #include "saiga/cuda/imgui_cuda.h" 4 | #include "saiga/vision/torch/CudaHelper.h" 5 | 6 | #include "NeuralPointCloudCuda.h" 7 | #include "PointRenderer.h" 8 | #include "RenderConstants.h" 9 | #include "RenderInfo.h" 10 | #include "config.h" 11 | #include "data/Dataset.h" 12 | #include "data/NeuralScene.h" 13 | #include "data/Settings.h" 14 | 15 | #include "cooperative_groups.h" 16 | #include 17 | 18 | 19 | 20 | #ifdef CUDA_DEBUG 21 | # define CUDA_DEBUG_ASSERT(_x) CUDA_KERNEL_ASSERT(_x) 22 | #else 23 | # define CUDA_DEBUG_ASSERT(_x) 24 | #endif 25 | using namespace Saiga; 26 | // constants 27 | static constexpr int default_block_size = 128; 28 | static constexpr int POINT_PER_THREAD = 1; 29 | #define POINT_BATCH_SIZE 10240 30 | #define MAX_DEPTH_CONST 100000.f 31 | 32 | using Packtype = unsigned long long; 33 | constexpr int max_layers = 5; 34 | 35 | 36 | struct DeviceRenderParams 37 | { 38 | int num_texture_channels; 39 | bool check_normal; 40 | float dropout; 41 | bool ghost_gradients; 42 | float dist_cutoff; 43 | int num_layers; 44 | float depth_accept; 45 | float depth_accept_blend; 46 | float drop_out_radius_threshold; 47 | bool drop_out_points_by_radius; 48 | int test_backward_mode; 49 | float distortion_gradient_factor; 50 | float K_gradient_factor; 51 | 52 | float debug_test_refl_x; 53 | float debug_test_refl_y; 54 | float debug_test_refl_z; 55 | int current_epoch; 56 | bool use_point_adding_and_removing_module; 57 | float stability_cutoff_value; 58 | // float gradient_spread; 59 | bool viewer_only; 60 | int debug_max_list_length; 61 | // For every layer a batch of images 62 | // [layers, batches, 1, height_of_layer, width_of_layer] 63 | StaticDeviceTensor depth[max_layers]; 64 | StaticDeviceTensor weight[max_layers]; 65 | StaticDeviceTensor max_depth[max_layers]; 66 | 67 | StaticDeviceTensor counting[max_layers]; 68 | StaticDeviceTensor per_image_atomic_counters[max_layers]; 69 | 70 | StaticDeviceTensor gradient_of_forward_pass_x; 71 | 72 | StaticDeviceTensor tmp_projections; 73 | 74 | // for every image one pose 75 | Sophus::SE3d* _poses; 76 | curandState* curand_state; 77 | 78 | HD inline Sophus::SE3f Pose(int image_index) { return _poses[image_index].cast(); } 79 | 80 | // [num_cameras, num_model_params] 81 | StaticDeviceTensor intrinsics; 82 | 83 | HD inline thrust::pair PinholeIntrinsics(int camera_index) 84 | { 85 | float* ptr = &intrinsics(camera_index, 0); 86 | IntrinsicsPinholef K = ((vec5*)ptr)[0]; 87 | Distortionf distortion = ((vec8*)(ptr + 5))[0]; 88 | 89 | return {K, distortion}; 90 | } 91 | 92 | HD inline thrust::pair, ArrayView> OcamIntrinsics(int camera_index) 93 | { 94 | float* ptr = &intrinsics(camera_index, 0); 95 | int count = intrinsics.sizes[1]; 96 | 97 | Vector aff = ((vec5*)ptr)[0]; 98 | ArrayView poly((ptr + 5), count - 5); 99 | 100 | return {aff, poly}; 101 | } 102 | 103 | __host__ __device__ DeviceRenderParams() = default; 104 | 105 | DeviceRenderParams(RenderParams params) 106 | { 107 | num_texture_channels = params.num_texture_channels; 108 | use_point_adding_and_removing_module = params.use_point_adding_and_removing_module; 109 | check_normal = params.check_normal; 110 | dropout = params.dropout; 111 | ghost_gradients = params.ghost_gradients; 112 | dist_cutoff = params.dist_cutoff; 113 | depth_accept = params.depth_accept; 114 | drop_out_points_by_radius = params.drop_out_points_by_radius; 115 | drop_out_radius_threshold = params.drop_out_radius_threshold; 116 | test_backward_mode = params.test_backward_mode; 117 | distortion_gradient_factor = params.distortion_gradient_factor; 118 | K_gradient_factor = params.K_gradient_factor; 119 | debug_test_refl_x = params.test_refl_x; 120 | debug_test_refl_y = params.test_refl_y; 121 | debug_test_refl_z = params.test_refl_z; 122 | stability_cutoff_value = params.stability_cutoff_value; 123 | viewer_only = params.viewer_only; 124 | debug_max_list_length = params.debug_max_list_length; 125 | depth_accept_blend = params.depth_accept_blend; 126 | } 127 | }; 128 | 129 | struct DeviceTexture 130 | { 131 | StaticDeviceTensor in_texture; 132 | 133 | StaticDeviceTensor points_confidence_value; 134 | }; 135 | 136 | struct DeviceAlphaCompositionParams 137 | { 138 | StaticDeviceTensor collections[max_layers]; 139 | StaticDeviceTensor per_point_data[max_layers]; 140 | StaticDeviceTensor gradient_sum_backwards[max_layers]; 141 | StaticDeviceTensor scanned_countings[max_layers]; 142 | StaticDeviceTensor ticket_counter[max_layers]; 143 | }; 144 | // #endif 145 | 146 | struct DeviceForwardParams 147 | { 148 | StaticDeviceTensor neural_out[max_layers]; 149 | // StaticDeviceTensor blend_out[max_layers]; 150 | }; 151 | 152 | struct DeviceBackwardParams 153 | { 154 | Vec6* out_gradient_pose; 155 | float* out_gradient_pose_count; 156 | 157 | vec4* out_gradient_points; 158 | float* out_gradient_points_count; 159 | 160 | 161 | // vec4* out_gradient_dynamic_points; 162 | // float* out_gradient_dynamic_points_count; 163 | StaticDeviceTensor out_gradient_dynamic_points; 164 | StaticDeviceTensor out_gradient_dynamic_points_count; 165 | 166 | StaticDeviceTensor out_gradient_intrinsics; 167 | float* out_gradient_intrinsics_count; 168 | 169 | StaticDeviceTensor out_gradient_texture; 170 | StaticDeviceTensor out_gradient_confidence; 171 | StaticDeviceTensor in_gradient_image[max_layers]; 172 | }; 173 | -------------------------------------------------------------------------------- /src/lib/rendering/RenderInfo.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #include "RenderInfo.h" 8 | 9 | #include "PointRenderer.h" 10 | 11 | 12 | TORCH_LIBRARY(sdfgergsd, m) 13 | { 14 | std::cout << "register neural render info" << std::endl; 15 | m.class_("NeuralRenderInfo").def(torch::init()); 16 | } 17 | 18 | std::vector BlendPointCloud(NeuralRenderInfo* info) 19 | { 20 | torch::intrusive_ptr render_data = torch::make_intrusive(*info); 21 | 22 | SAIGA_ASSERT(render_data->scene); 23 | SAIGA_ASSERT(render_data->scene->texture); 24 | 25 | return torch::autograd::PointRender::apply( 26 | render_data->scene->texture->texture, render_data->scene->texture->background_color, 27 | render_data->scene->point_cloud_cuda->t_position, render_data->scene->poses->tangent_poses, 28 | render_data->scene->intrinsics->intrinsics, render_data->scene->texture->confidence_value_of_point, 29 | render_data->scene->dynamic_refinement_t, torch::IValue(render_data)); 30 | } 31 | 32 | namespace torch::autograd 33 | { 34 | variable_list PointRender::forward(AutogradContext* ctx, Variable texture, Variable background_color, Variable points, 35 | Variable pose_tangents, Variable intrinsics, Variable confidence, 36 | Variable dynamic_refinement, IValue info) 37 | { 38 | ctx->saved_data["render_info"] = info; 39 | 40 | auto [color, mask] = BlendPointCloudForward(ctx, info.toCustomClass().get()); 41 | color.insert(color.end(), mask.begin(), mask.end()); 42 | return color; 43 | } 44 | 45 | variable_list PointRender::backward(AutogradContext* ctx, variable_list grad_output) 46 | { 47 | IValue info = ctx->saved_data["render_info"]; 48 | auto in = info.toCustomClass().get(); 49 | auto derives = BlendPointCloudBackward(ctx, in, grad_output); 50 | 51 | // With enviroment map we have additonal inputs because we render them first 52 | int expected_output = 5; 53 | expected_output = 6; 54 | expected_output = 7; 55 | SAIGA_ASSERT(derives.size() == expected_output); 56 | 57 | // std::cout << "BLEND OUTPUT" << std::endl; 58 | // for (int i = 0; i < expected_output; ++i) 59 | //{ 60 | // std::cout << TensorInfo(derives[i]) << std::endl; 61 | // } 62 | 63 | // The last empty grad is for the 'IValue info' of the forward function 64 | derives.push_back(Variable()); 65 | 66 | return derives; 67 | } 68 | // namespace torch::autograd 69 | } // namespace torch::autograd -------------------------------------------------------------------------------- /src/lib/rendering/RenderInfo.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "saiga/cuda/imageProcessing/image.h" 9 | 10 | #include "NeuralPointCloudCuda.h" 11 | #include "config.h" 12 | #include "data/SceneData.h" 13 | 14 | #include 15 | 16 | 17 | 18 | struct TorchFrameData 19 | { 20 | ReducedImageInfo img; 21 | 22 | torch::Tensor target; 23 | 24 | // binary float tensor, where the loss should only be taken if target_mask == 1 25 | torch::Tensor target_mask; 26 | 27 | torch::Tensor uv, uv_local; 28 | 29 | // long index for the camera 30 | // used for training camera specific parameters 31 | torch::Tensor camera_index; 32 | torch::Tensor scale; 33 | torch::Tensor timestep; 34 | 35 | int scene_id; 36 | void to(torch::Device device) 37 | { 38 | if (target.defined()) target = target.to(device); 39 | if (target_mask.defined()) target_mask = target_mask.to(device); 40 | if (uv.defined()) uv = uv.to(device); 41 | if (camera_index.defined()) camera_index = camera_index.to(device); 42 | } 43 | }; 44 | -------------------------------------------------------------------------------- /src/lib/rendering/RenderModule.cpp: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #include "RenderModule.h" 8 | 9 | #include 10 | PointRenderModuleImpl::PointRenderModuleImpl(std::shared_ptr params) 11 | : params(params), num_layers(params->net_params.num_input_layers) 12 | { 13 | cache = std::make_shared(); 14 | } 15 | 16 | std::pair, std::vector> PointRenderModuleImpl::forward( 17 | NeuralScene& scene, const std::vector& batch, int current_epoch, 18 | CUDA::CudaTimerSystem* timer_system) 19 | { 20 | NeuralRenderInfo render_data; 21 | render_data.scene = &scene; 22 | render_data.num_layers = num_layers; 23 | render_data.params = params->render_params; 24 | render_data.timer_system = timer_system; 25 | render_data.current_epoch = current_epoch; 26 | 27 | if (!this->is_training()) 28 | { 29 | render_data.params.dropout = 0; 30 | render_data.train = false; 31 | } 32 | else 33 | { 34 | render_data.train = true; 35 | } 36 | 37 | for (auto& b : batch) 38 | { 39 | render_data.images.push_back(b->img); 40 | } 41 | 42 | return forward(&render_data); 43 | } 44 | 45 | std::pair, std::vector> PointRenderModuleImpl::forward(NeuralRenderInfo* nri) 46 | { 47 | if (0) 48 | { 49 | auto poses = nri->scene->poses->Download(); 50 | auto ks = nri->scene->intrinsics->DownloadK(); 51 | for (auto i : nri->images) 52 | { 53 | std::cout << "Render (" << i.camera_index << ", " << i.image_index << ") Pose: " << poses[i.image_index] 54 | << " K: " << ks[i.camera_index] << std::endl; 55 | } 56 | } 57 | 58 | nri->cache = cache.get(); 59 | 60 | if (params->render_params.super_sampling) 61 | { 62 | for (auto& i : nri->images) 63 | { 64 | i.w *= 2; 65 | i.h *= 2; 66 | i.crop_transform = i.crop_transform.scale(2); 67 | } 68 | } 69 | 70 | 71 | auto combined_images_masks = BlendPointCloud(nri); 72 | 73 | 74 | std::vector images(combined_images_masks.begin(), combined_images_masks.begin() + nri->num_layers); 75 | 76 | std::vector depths; 77 | if (params->render_params.add_depth_to_network) 78 | depths = std::vector(combined_images_masks.begin() + nri->num_layers, 79 | combined_images_masks.begin() + 2 * nri->num_layers); 80 | 81 | std::vector blendPointMasks; 82 | 83 | if (params->render_params.output_background_mask) 84 | { 85 | if (params->render_params.add_depth_to_network) 86 | { 87 | SAIGA_ASSERT(combined_images_masks.size() == nri->num_layers * 3); 88 | blendPointMasks = std::vector(combined_images_masks.begin() + 2 * nri->num_layers, 89 | combined_images_masks.end()); 90 | } 91 | else 92 | { 93 | SAIGA_ASSERT(combined_images_masks.size() == nri->num_layers * 2); 94 | blendPointMasks = std::vector(combined_images_masks.begin() + nri->num_layers, 95 | combined_images_masks.end()); 96 | } 97 | for (auto& m : blendPointMasks) 98 | { 99 | m.detach_(); 100 | } 101 | SAIGA_ASSERT(!blendPointMasks.front().requires_grad()); 102 | } 103 | 104 | if (params->render_params.super_sampling) 105 | { 106 | for (auto& img : images) 107 | { 108 | img = torch::avg_pool2d(img, {2, 2}); 109 | } 110 | 111 | for (auto& img : depths) 112 | { 113 | img = torch::avg_pool2d(img, {2, 2}); 114 | std::cout << "think about this: " << __LINE__ << std::endl; 115 | } 116 | 117 | for (auto& img : blendPointMasks) 118 | { 119 | img = torch::avg_pool2d(img, {2, 2}); 120 | } 121 | 122 | for (auto& i : nri->images) 123 | { 124 | i.w /= 2; 125 | i.h /= 2; 126 | i.crop_transform = i.crop_transform.scale(0.5); 127 | } 128 | } 129 | 130 | if (nri->scene->environment_map && params->pipeline_params.enable_environment_map) 131 | { 132 | // SAIGA_OPTIONAL_TIME_MEASURE("Environment Map", nri->timer_system); 133 | std::vector> layers(cache->layers_cuda.size()); 134 | for (int i = 0; i < cache->layers_cuda.size(); ++i) 135 | { 136 | layers[i] = std::vector(cache->num_batches); 137 | std::vector dt_masking_tensors; 138 | for (int b = 0; b < cache->num_batches; ++b) 139 | { 140 | layers[i][b] = cache->layers_cuda[i].weight.slice(0, b, b + 1).clone(); 141 | 142 | // for depth testing, either weight is zero, thus alpha_dest should be 1 or weight is >0 thus 143 | // accumulation is complete; also mask out background color 144 | if (cache->render_mode == PointRendererCache::RenderMode::FUZZY_DT) 145 | { 146 | dt_masking_tensors.push_back(torch::where(layers[i][b] > 0.5, torch::ones_like(layers[i][b]), 147 | torch::zeros_like(layers[i][b]))); 148 | layers[i][b] = torch::where(layers[i][b] > 0.5, torch::zeros_like(layers[i][b]), 149 | torch::ones_like(layers[i][b])); 150 | } 151 | } 152 | if (cache->render_mode == PointRendererCache::RenderMode::FUZZY_DT) 153 | images[i] = images[i] * torch::cat(dt_masking_tensors, 0); 154 | } 155 | 156 | SAIGA_ASSERT(params->render_params.output_background_mask); 157 | 158 | static bool use_new_sampling = true; 159 | // ImGui::Checkbox("use new sampling", &use_new_sampling); 160 | 161 | std::vector env_maps; 162 | if (use_new_sampling) 163 | { 164 | env_maps = nri->scene->environment_map->Sample2( 165 | nri->scene->poses->poses_se3, nri->scene->intrinsics->intrinsics, nri->images, nri->num_layers, 166 | nri->scene->scene, layers, nri->timer_system); 167 | } 168 | else 169 | { 170 | env_maps = nri->scene->environment_map->Sample( 171 | nri->scene->poses->poses_se3, nri->scene->intrinsics->intrinsics, nri->images, nri->num_layers, 172 | nri->scene->scene, layers, nri->timer_system); 173 | } 174 | 175 | if (params->pipeline_params.cat_env_to_color) 176 | { 177 | for (int i = 0; i < nri->num_layers; ++i) 178 | { 179 | images[i] = torch::cat({images[i], env_maps[i]}, 1); 180 | } 181 | } 182 | else 183 | { 184 | for (int i = 0; i < nri->num_layers; ++i) 185 | { 186 | torch::Tensor env_map_masking = torch::ones_like(env_maps[i]); 187 | if (params->render_params.no_envmap_at_points) 188 | { 189 | std::vector masking; 190 | for (int b = 0; b < cache->num_batches; ++b) 191 | { 192 | layers[i][b] = cache->layers_cuda[i].weight.slice(0, b, b + 1).clone(); 193 | 194 | masking.push_back(torch::where(layers[i][b] > 0.001, torch::zeros_like(layers[i][b]), 195 | torch::ones_like(layers[i][b]))); 196 | } 197 | env_map_masking = 198 | torch::cat(masking, 0).repeat({1, params->render_params.num_texture_channels, 1, 1}); 199 | } 200 | images[i] = images[i] + env_map_masking * env_maps[i]; 201 | } 202 | } 203 | } 204 | 205 | if (params->pipeline_params.cat_masks_to_color) 206 | { 207 | for (int i = 0; i < nri->num_layers; ++i) 208 | { 209 | images[i] = torch::cat({images[i], blendPointMasks[i]}, 1); 210 | } 211 | } 212 | 213 | #if 0 214 | for (int i = 0; i < nri->num_layers; ++i) 215 | { 216 | using namespace torch::indexing; 217 | std::cout << images[i].sizes() << ", " << depths[i].sizes() << std::endl; 218 | auto img = Saiga::TensorToImage(images[i].index({0,Slice(0,3),Slice(),Slice()})); 219 | img.save("./img_"+ std::to_string(i) + ".png"); 220 | } 221 | #endif 222 | 223 | if (params->render_params.add_depth_to_network) 224 | { 225 | for (int i = 0; i < nri->num_layers; ++i) 226 | { 227 | #if 0 228 | 229 | auto dep = Saiga::TensorToImage(depths[i].expand({-1,3,-1,-1})/10.0); 230 | dep.save("./dep_"+ std::to_string(i) + ".png"); 231 | #endif 232 | images[i] = torch::cat({images[i], depths[i]}, 1); 233 | } 234 | } 235 | 236 | 237 | 238 | // SAIGA_ASSERT(images.front().size(0) == nri->images.size()); 239 | SAIGA_ASSERT(images.front().size(2) == nri->images.front().h); 240 | SAIGA_ASSERT(images.front().size(3) == nri->images.front().w); 241 | 242 | 243 | 244 | return {images, blendPointMasks}; 245 | } 246 | -------------------------------------------------------------------------------- /src/lib/rendering/RenderModule.h: -------------------------------------------------------------------------------- 1 | /** 2 | * Copyright (c) 2023 Darius Rückert and Linus Franke 3 | * Licensed under the MIT License. 4 | * See LICENSE file for more information. 5 | */ 6 | 7 | #pragma once 8 | #include "PointRenderer.h" 9 | 10 | class PointRenderModuleImpl : public torch::nn::Module 11 | { 12 | public: 13 | PointRenderModuleImpl(std::shared_ptr params); 14 | 15 | // Renders the point cloud image in multiple scale levels and returns all of them. 16 | // If masks are required, the first n images are the color images followed by n mask images 17 | // 18 | std::pair, std::vector> forward(NeuralRenderInfo* nri); 19 | 20 | std::pair, std::vector> forward( 21 | NeuralScene& scene, const std::vector& batch, int current_epoch = -1, 22 | CUDA::CudaTimerSystem* timer_system = nullptr); 23 | 24 | std::shared_ptr params; 25 | int num_layers; 26 | std::shared_ptr cache; 27 | }; 28 | 29 | 30 | TORCH_MODULE(PointRenderModule); --------------------------------------------------------------------------------