├── .gitignore ├── .gitmodules ├── CMakeLists.txt ├── LICENSE ├── README.md ├── config ├── ROG-Map │ └── rog.yaml ├── base.yaml ├── fast_livo │ ├── 2_dao_fast.yaml │ ├── campus.yaml │ ├── cbd.yaml │ ├── culture01.yaml │ ├── drive.yaml │ ├── fast_livo.yaml │ ├── red_bird.yaml │ ├── station.yaml │ └── sysu.yaml ├── replica │ └── replica.yaml └── rviz │ └── vis_thirdperson.rviz ├── docs └── FAST_LIVO2_Data_Training.md ├── eval ├── draw.py ├── draw_loss.py ├── draw_mem.py ├── image_metrics │ ├── eval_image.sh │ ├── image_utils.py │ ├── loss_utils.py │ ├── lpipsPyTorch │ │ ├── __init__.py │ │ └── modules │ │ │ ├── lpips.py │ │ │ ├── networks.py │ │ │ └── utils.py │ ├── metrics.py │ ├── metrics2.py │ └── metrics_single.py └── structure_metrics │ ├── cull_mesh.py │ ├── eval_utils.py │ └── evaluator.py ├── include ├── data_loader │ ├── data_loader.cpp │ ├── data_loader.h │ └── data_parsers │ │ ├── base_parser.cpp │ │ ├── base_parser.h │ │ ├── fastlivo_parser.hpp │ │ ├── kitti_parser.hpp │ │ ├── ncd_parser.hpp │ │ ├── neuralrgbd_parser.hpp │ │ ├── oxford_spires_parser.hpp │ │ ├── r3live_parser.hpp │ │ ├── replica_parser.hpp │ │ └── rosbag_parser.hpp ├── mesher │ ├── cumcubes │ │ ├── __init__.py │ │ ├── include │ │ │ ├── cumcubes.hpp │ │ │ └── utils.cuh │ │ ├── src │ │ │ ├── bindings.cpp │ │ │ ├── cumcubes.cpp │ │ │ └── cumcubes_kernel.cu │ │ ├── utils.py │ │ └── version.py │ ├── mesher.cpp │ └── mesher.h ├── neural_mapping │ ├── neural_mapping.cpp │ └── neural_mapping.h ├── neural_net │ ├── encoding_map.cpp │ ├── encoding_map.h │ ├── encodings │ │ ├── encodings.cpp │ │ └── encodings.h │ ├── local_map.cpp │ ├── local_map.h │ ├── sub_map.cpp │ └── sub_map.h ├── optimizer │ ├── loss.cpp │ ├── loss.h │ └── loss_utils │ │ ├── loss_utils.cpp │ │ └── loss_utils.h ├── params │ ├── params.cpp │ └── params.h ├── tracer │ ├── sphere_trace │ │ ├── sphere_trace.cpp │ │ ├── sphere_trace.h │ │ └── sphere_trace_cuda.cu │ ├── tracer.cpp │ └── tracer.h └── utils │ ├── bin_utils │ └── endian.h │ ├── coordinates.cpp │ ├── coordinates.h │ ├── loss_utils.hpp │ ├── ray_utils │ ├── ray_utils.cpp │ └── ray_utils.h │ ├── sensor_utils │ ├── cameras.hpp │ └── sensors.hpp │ ├── tqdm.hpp │ ├── utils.cpp │ └── utils.h ├── launch └── rviz.launch ├── package.xml ├── pics └── pipeline_h.jpg ├── src ├── backward.hpp └── neural_mapping_node.cpp ├── submodules └── .gitignore └── supplement.pdf /.gitignore: -------------------------------------------------------------------------------- 1 | **/build/** 2 | **/.vscode/** 3 | **/__pycache__/** 4 | **/output/** 5 | 6 | *.pyc 7 | output 8 | data 9 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "submodules/llog"] 2 | path = submodules/llog 3 | url = https://github.com/jianhengLiu/llog.git 4 | [submodule "submodules/kaolin_wisp_cpp"] 5 | path = submodules/kaolin_wisp_cpp 6 | url = https://github.com/jianhengLiu/kaolin_wisp_cpp.git 7 | [submodule "submodules/tcnn_binding"] 8 | path = submodules/tcnn_binding 9 | url = https://github.com/jianhengLiu/tcnn_binding.git 10 | [submodule "submodules/nerfacc_cpp"] 11 | path = submodules/nerfacc_cpp 12 | url = https://github.com/jianhengLiu/nerfacc_cpp.git 13 | [submodule "submodules/ROG-Map"] 14 | path = submodules/ROG-Map 15 | url = https://github.com/jianhengLiu/ROG-Map.git 16 | branch = neural_mapping 17 | [submodule "include/utils/ply_utils"] 18 | path = include/utils/ply_utils 19 | url = https://github.com/jianhengLiu/ply_utils.git 20 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0.2) 2 | project(neural_mapping LANGUAGES CXX CUDA) 3 | 4 | # set(CUDA_TOOLKIT_ROOT_DIR /usr/local/cuda-11.8/) 5 | set(CMAKE_CXX_STANDARD 17) 6 | set(CMAKE_BUILD_TYPE "RelWithDebInfo") 7 | set(CMAKE_EXPORT_COMPILE_COMMANDS ON) # for clangd 8 | 9 | # set(CMAKE_BUILD_TYPE "Release") set(CMAKE_BUILD_TYPE "Debug") 10 | set(CMAKE_CXX_FLAGS "-fPIC") 11 | 12 | add_definitions(-O3 -DWITH_CUDA -DTHRUST_IGNORE_CUB_VERSION_CHECK) 13 | 14 | option(ENABLE_ROS "Enable ROS support" OFF) 15 | 16 | # Define installation directories based on build type 17 | if(ENABLE_ROS) 18 | message(STATUS "ROS ENABLED - Using catkin build system") 19 | add_definitions(-DENABLE_ROS) 20 | 21 | # Use catkin build system 22 | find_package( 23 | catkin REQUIRED 24 | COMPONENTS roscpp 25 | rosbag 26 | roslib 27 | std_msgs 28 | geometry_msgs 29 | nav_msgs 30 | mesh_msgs 31 | cv_bridge 32 | tf) 33 | set(ROS_LIBRARIES ${catkin_LIBRARIES}) 34 | else() 35 | message(STATUS "ROS DISABLED - Using standard CMake build system") 36 | # No catkin dependency 37 | set(ROS_LIBRARIES "") 38 | endif() 39 | 40 | find_package(OpenMP REQUIRED) 41 | 42 | if(OPENMP_FOUND) 43 | message("OPENMP FOUND") 44 | set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") 45 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") 46 | endif() 47 | 48 | find_package(OpenCV REQUIRED) 49 | include_directories(${OpenCV_INCLUDE_DIRS}) 50 | find_package(Eigen3 REQUIRED) 51 | message(STATUS "Eigen: ${EIGEN3_INCLUDE_DIR}") 52 | 53 | set(CMAKE_SUPPRESS_DEVELOPER_WARNINGS 54 | 1 55 | CACHE INTERNAL "No dev warnings") # 关闭pcl烦人的警告 56 | find_package(PCL REQUIRED) 57 | 58 | # llog 59 | add_subdirectory(submodules/llog) 60 | include_directories(submodules/llog/include) 61 | 62 | # 指定libTorch位置 63 | find_package(Torch REQUIRED) 64 | 65 | # ROG-Map 66 | add_subdirectory(submodules/ROG-Map) 67 | include_directories(submodules/ROG-Map/include) 68 | 69 | include_directories(${EIGEN3_INCLUDE_DIR} ${PCL_INCLUDE_DIRS} include) 70 | 71 | # Add include directories for ROS if enabled 72 | if(ENABLE_ROS) 73 | include_directories(${catkin_INCLUDE_DIRS}) 74 | 75 | catkin_package() 76 | endif() 77 | 78 | # kaolin_wisp_cpp 79 | add_subdirectory(submodules/kaolin_wisp_cpp) 80 | include_directories(submodules/kaolin_wisp_cpp 81 | submodules/kaolin_wisp_cpp/submodules/kaolin) 82 | 83 | # tcnn_binding 84 | add_subdirectory(submodules/tcnn_binding) 85 | include_directories( 86 | submodules/tcnn_binding 87 | submodules/tcnn_binding/submodules/tiny-cuda-nn/include 88 | submodules/tcnn_binding/submodules/tiny-cuda-nn/dependencies) 89 | 90 | # nerfacc_cpp 91 | add_subdirectory(submodules/nerfacc_cpp) 92 | include_directories(submodules/nerfacc_cpp 93 | submodules/nerfacc_cpp/submodules/nerfacc) 94 | 95 | add_library(ply_utils include/utils/ply_utils/ply_utils_pcl.cpp 96 | include/utils/ply_utils/ply_utils_torch.cpp) 97 | target_link_libraries(ply_utils ${ROS_LIBRARIES} ${PCL_LIBRARIES} 98 | ${TORCH_LIBRARIES}) 99 | 100 | add_library(cumcubes include/mesher/cumcubes/src/cumcubes.cpp 101 | include/mesher/cumcubes/src/cumcubes_kernel.cu) 102 | target_link_libraries(cumcubes ${ROS_LIBRARIES} ${TORCH_LIBRARIES}) 103 | target_include_directories(cumcubes PUBLIC include/mesher/cumcubes/include) 104 | 105 | add_library(mesher include/utils/utils.cpp include/mesher/mesher.cpp) 106 | target_link_libraries(mesher ply_utils cumcubes) 107 | 108 | add_library(data_parser include/data_loader/data_parsers/base_parser.cpp) 109 | target_link_libraries(data_parser ${TORCH_LIBRARIES} ply_utils) 110 | 111 | add_library(ray_utils include/utils/ray_utils/ray_utils.cpp) 112 | target_link_libraries(ray_utils ${TORCH_LIBRARIES}) 113 | 114 | add_library(data_loader include/data_loader/data_loader.cpp 115 | include/utils/coordinates.cpp) 116 | target_link_libraries(data_loader data_parser ray_utils) 117 | 118 | add_library( 119 | neural_net 120 | include/params/params.cpp include/neural_net/sub_map.cpp 121 | include/neural_net/encoding_map.cpp include/neural_net/local_map.cpp) 122 | target_link_libraries(neural_net mesher kaolin_wisp_cpp tcnn_binding llog 123 | ray_utils) 124 | 125 | add_library( 126 | tracer include/tracer/tracer.cpp include/tracer/sphere_trace/sphere_trace.cpp 127 | include/tracer/sphere_trace/sphere_trace_cuda.cu) 128 | target_link_libraries(tracer ${TORCH_LIBRARIES} nerfacc kaolin_wisp_cpp) 129 | 130 | add_library(loss_utils include/optimizer/loss.cpp 131 | include/optimizer/loss_utils/loss_utils.cpp) 132 | target_link_libraries(loss_utils ${TORCH_LIBRARIES}) 133 | 134 | add_library(neural_mapping_lib include/neural_mapping/neural_mapping.cpp 135 | include/params/params.cpp) 136 | target_link_libraries( 137 | neural_mapping_lib 138 | ${OpenCV_LIBS} 139 | dw 140 | data_loader 141 | neural_net 142 | rog_map_cuda 143 | tracer 144 | loss_utils) 145 | 146 | add_executable(neural_mapping_node src/neural_mapping_node.cpp) 147 | target_link_libraries(neural_mapping_node neural_mapping_lib) 148 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # M2Mapping 2 | 3 | ## Neural Surface Reconstruction and Rendering for LiDAR-Visual Systems 4 | 5 | ### 1. Introduction 6 | 7 | ![alt text](pics/pipeline_h.jpg) 8 | A MultiModal Mapping (M2Mapping) Framework for LiDAR-Visual Systems. 9 | 10 | [Project page](https://jianhengliu.github.io/Projects/M2Mapping/) | [Paper](https://arxiv.org/pdf/2409.05310) | [Video](https://www.youtube.com/watch?v=XFzzAGVbzek) | [M2Mapping Datasets](https://furtive-lamprey-00b.notion.site/M2Mapping-Datasets-e6318dcd710e4a9d8a4f4b3fbe176764) | [M2Mapping Results](https://furtive-lamprey-00b.notion.site/M2Mapping-Results-bf02b1b1ebbd443099e2f076019d1c51) 11 | 12 | Our paper is accepted by ICRA 2025. If you use M2Mapping for your academic research, please cite the following paper: 13 | 14 | ``` 15 | @article{liu2024neural, 16 | title={Neural Surface Reconstruction and Rendering for LiDAR-Visual Systems}, 17 | author={Jianheng Liu and Chunran Zheng and Yunfei Wan and Bowen Wang and Yixi Cai and Fu Zhang}, 18 | journal={arXiv preprint arXiv:2108.10470}, 19 | year={2024}, 20 | } 21 | ``` 22 | 23 | ### 2. Installation 24 | 25 | - Tested on Ubuntu 20.04, cuda 11.8 26 | 27 | > The software not relies on ROS, but under ROS noetic installed, the installation should be easier. 28 | > And if real-time visualization is needed, ROS is required and refer to the [Visualization](#3-visualization) section. 29 | 30 | ```bash 31 | pip install open3d==0.18.0 32 | # Libtorch 33 | wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.4.1%2Bcu118.zip 34 | apt install zip 35 | unzip libtorch-cxx11-abi-shared-with-deps-*.zip 36 | rm *.zip 37 | echo "export Torch_DIR=$PWD/libtorch/share/cmake/Torch" >> ~/.bashrc # ~/.zshrc if you use zsh 38 | source ~/.bashrc # .zshrc if you use zsh 39 | 40 | mkdir -p m2mapping_ws/src 41 | cd m2mapping_ws/src 42 | apt install git libdw-dev 43 | git clone https://github.com/hku-mars/M2Mapping.git --recursive 44 | cd .. 45 | 46 | sudo apt install python3-catkin-tools # install catkin 47 | ``` 48 | 49 | #### Build with ROS for visualization 50 | 51 | ```bash 52 | catkin_make -j8 -DENABLE_ROS=ON 53 | ``` 54 | 55 | #### (Alternative) Build without ROS 56 | 57 | ```bash 58 | # Instead of build with catkin_make, you can also build with cmake 59 | cd m2mapping_ws/src/m2mapping 60 | mkdir build 61 | cd build 62 | cmake .. 63 | make -j8 64 | ``` 65 | 66 | ### 3. Data Preparation 67 | 68 | - The processed FAST-LIVO2 Datasets and Replica Extrapolation Datasets are available at [M2Mapping Datasets](https://furtive-lamprey-00b.notion.site/M2Mapping-Datasets-e6318dcd710e4a9d8a4f4b3fbe176764) 69 | 70 | #### 3.1. Replica 71 | 72 | - Download the Replica dataset from [M2Mapping Datasets](https://furtive-lamprey-00b.notion.site/M2Mapping-Datasets-e6318dcd710e4a9d8a4f4b3fbe176764) and unzip it to `src/M2Mapping/data`: 73 | ```bash 74 | wget https://cvg-data.inf.ethz.ch/nice-slam/data/Replica.zip 75 | # Replica.zip, cull_replica_mesh.zip, and replica_extra_eval.zip are supposed under m2mapping_ws 76 | unzip -d src/M2Mapping/data Replica.zip 77 | unzip -d src/M2Mapping/data/Replica cull_replica_mesh.zip 78 | unzip -d src/M2Mapping/data replica_extra_eval.zip 79 | ``` 80 | - Arrange the data as follows: 81 | ```bash 82 | ├── Replica 83 | │   ├── cull_replica_mesh 84 | │   │   ├── *.ply 85 | │   ├── room2 86 | │   │   ├── eval 87 | │   │   │   └── results 88 | │   │   │   │   └── *.jpg 89 | │   │   │   │   └── *.png 90 | │   │   │   └── traj.txt 91 | │   │   └── results 92 | │   │   │   └── *.jpg 93 | │   │   │   └── *.png 94 | │   │   └── traj.txt 95 | ``` 96 | 97 | #### 3.2. FAST-LIVO2 Datasets 98 | 99 | - Download either Rosbag or Parsered Data in [M2Mapping Datasets](https://furtive-lamprey-00b.notion.site/M2Mapping-Datasets-e6318dcd710e4a9d8a4f4b3fbe176764). 100 | - Arrange the data as follows: 101 | 102 | - For Rosbag: 103 | ```bash 104 | ├── data 105 | │   ├── FAST_LIVO2_Datasets 106 | │   ├── campus 107 | │   │   │   ├── fast_livo2_campus.bag 108 | ``` 109 | - For Parsered Data: 110 | ```bash 111 | ├── data 112 | │   ├── FAST_LIVO2_Datasets 113 | │   │   ├── campus 114 | │   │   │   ├── images 115 | │   │   │   ├── depths 116 | │   │   │   ├── color_poses.txt 117 | │   │   │   ├── depth_poses.txt 118 | ``` 119 | 120 | #### 3.3. Custom FAST-LIVO2 Datasets 121 | 122 | - Clone the [modified-FAST-LIVO2](https://github.com/jianhengLiu/FAST-LIVO2) repo; install and run FAST-LIVO2 as the official instruction. The overall pipeline as: 123 | ```bash 124 | # 1. open a terminal to start LIVO 125 | roslaunch fast_livo mapping_avia.launch 126 | # 2. open another terminal to get ready for bag recording 127 | rosbag record /aft_mapped_to_init /origin_img /cloud_registered_body /tf /tf_static /path -O "fast_livo2_YOUR_DOWNLOADED" -b 2048 128 | # 3. open another terminal to play your downloaded/collected bag 129 | rosbag play YOUR_DOWNLOADED.bag 130 | ``` 131 | > Warning: The modified-FAST-LIVO2 outputs the original size of the image, therefore if you try to reproduce the FAST-LIVO2 results from the original FAST-LIVO2 repo, you need to change the intrinsic parameters in cofig files into the original intrinsic parameters of the camera. 132 | 133 | ### 4. Run 134 | 135 | ```bash 136 | source devel/setup.bash # or setup.zsh 137 | 138 | # Replica 139 | ./src/M2Mapping/build/neural_mapping_node train src/M2Mapping/config/replica/replica.yaml src/M2Mapping/data/Replica/room2 140 | # If ROS is installed, you can also run the following command: 141 | # rosrun neural_mapping neural_mapping_node train src/M2Mapping/config/replica/replica.yaml src/M2Mapping/data/Replica/room2 142 | 143 | # FAST-LIVO2 (ROS installed & ROS bag) 144 | ./src/M2Mapping/build/neural_mapping_node train src/M2Mapping/config/fast_livo/campus.yaml src/M2Mapping/data/FAST_LIVO2_RIM_Datasets/campus/fast_livo2_campus.bag 145 | # If ROS is installed, you can also run the following command: 146 | # rosrun neural_mapping neural_mapping_node train src/M2Mapping/config/fast_livo/campus.yaml src/M2Mapping/data/FAST_LIVO2_RIM_Datasets/campus/fast_livo2_campus.bag 147 | 148 | # FAST-LIVO2 (Parsered ROS bag format) 149 | ./src/M2Mapping/build/neural_mapping_node train src/M2Mapping/config/fast_livo/campus.yaml src/M2Mapping/data/FAST_LIVO2_RIM_Datasets/campus/color_poses.txt 150 | # If ROS is installed, you can also run the following command: 151 | # rosrun neural_mapping neural_mapping_node train src/M2Mapping/config/fast_livo/campus.yaml src/M2Mapping/data/FAST_LIVO2_RIM_Datasets/campus/color_poses.txt 152 | ``` 153 | 154 | After running, the training and evaluation results will be saved in the `src/M2Mapping/output` directory. 155 | 156 | For afterward visualization/evaluation, you can use the following command: 157 | 158 | ```bash 159 | source devel/setup.bash # or setup.zsh 160 | ./src/M2Mapping/build/neural_mapping_node view src/M2Mapping/output/(your_output_folder) 161 | # If ROS is installed, you can also run the following command: 162 | # rosrun neural_mapping neural_mapping_node view src/M2Mapping/output/(your_output_folder) 163 | ``` 164 | 165 | Input `h` + `Enter` to see the help message. 166 | 167 | ### 5. Visualization 168 | 169 | - Tested on Ubuntu 20.04, cuda 11.8, ROS Noetic 170 | - We use RVIZ for visualization for now. Please install ROS Noetic following the [official guide](http://wiki.ros.org/noetic/Installation/Ubuntu) or refer to the [Docker](#6-docker) 'ROS Installation' section. 171 | - Re-build the packege: 172 | 173 | ```bash 174 | cd src 175 | git clone https://github.com/jianhengLiu/rviz_map_plugin.git 176 | git clone https://github.com/jianhengLiu/rviz_cinematographer.git 177 | sudo apt install ros-noetic-mesh-msgs ros-noetic-rviz-animated-view-controller ros-noetic-hdf5-map-io 178 | catkin_make -DENABLE_ROS=ON 179 | ``` 180 | - Run the following command to visualize the map in real-time: 181 | 182 | ```bash 183 | source devel/setup.bash # or setup.zsh 184 | roslaunch neural_mapping rviz.launch 185 | ``` 186 | 187 | Drag the view to activate and control the view with the mouse. 188 | - For post-training visualization, you can use the following command: 189 | 190 | ```bash 191 | ./src/M2Mapping/build/neural_mapping_node view src/M2Mapping/output/(your_output_folder) 192 | # If ROS is installed, you can also run the following command: 193 | # rosrun neural_mapping neural_mapping_node view src/M2Mapping/output/(your_output_folder) 194 | 195 | roslaunch neural_mapping rviz.launch 196 | ``` 197 | 198 | ### 6. Docker 199 | 200 | - We provide a [enroot](https://github.com/NVIDIA/enroot) docker image for testing. 201 | ```bash 202 | # https://github.com/NVIDIA/enroot 203 | enroot import docker://nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04 204 | enroot create --name m2mapping ~/nvidia+cuda+11.8.0-cudnn8-devel-ubuntu20.04.sqsh 205 | # check if create right 206 | enroot list 207 | enroot start --root --rw m2mapping 208 | # ctrl + d to return 209 | 210 | cd ~ 211 | # ROS Installation 212 | apt update 213 | apt install lsb-release 214 | sh -c 'echo "deb http://packages.ros.org/ros/ubuntu $(lsb_release -sc) main" > /etc/apt/sources.list.d/ros-latest.list' 215 | apt install curl # if you haven't already installed curl 216 | curl -s https://raw.githubusercontent.com/ros/rosdistro/master/ros.asc | apt-key add - 217 | apt update 218 | apt install ros-noetic-desktop-full 219 | # Dependencies 220 | apt install ros-noetic-mesh-msgs ros-noetic-rviz-animated-view-controller ros-noetic-hdf5-map-io 221 | echo "source /opt/ros/noetic/setup.bash" >> ~/.bashrc 222 | 223 | # Libtorch 224 | wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.4.1%2Bcu118.zip 225 | apt install zip 226 | unzip libtorch-cxx11-abi-shared-with-deps-*.zip 227 | rm *.zip 228 | echo "export Torch_DIR=$PWD/libtorch/share/cmake/Torch" >> ~/.bashrc 229 | source ~/.bashrc 230 | 231 | # upgrad cmake 232 | wget https://github.com/Kitware/CMake/releases/download/v3.23.0/cmake-3.23.0-linux-x86_64.sh 233 | bash ./cmake-3.23.0-linux-x86_64.sh --skip-licence --prefix=/usr 234 | # opt1: y; opt2: n 235 | 236 | mkdir -p m2mapping_ws/src 237 | cd m2mapping_ws/src 238 | apt install git libdw-dev 239 | git clone https://github.com/hku-mars/M2Mapping.git --recursive 240 | cd .. 241 | catkin_make -DENABLE_ROS=ON # if lacking memory try restricting number of cores: catkin_make -j8 242 | 243 | # Image export 244 | enroot export --output m2mapping.sqsh m2mapping 245 | ``` 246 | 247 | ### 7. Acknowledgement 248 | 249 | Thanks for the excellent open-source projects that we rely on: 250 | [RIM](https://github.com/HITSZ-NRSL/RIM), [FAST-LIVO2](https://github.com/hku-mars/FAST-LIVO2), [ROG-Map](https://github.com/hku-mars/ROG-Map), [nerfacc](https://github.com/nerfstudio-project/nerfacc), [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn), [kaolin-wisp](https://github.com/NVIDIAGameWorks/kaolin-wisp) 251 | -------------------------------------------------------------------------------- /config/ROG-Map/rog.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | ############################################# 3 | # Mapping # 4 | ############################################# 5 | # The points [out of local map] is considered to be [OUT_OF_MAP] 6 | # The points [out of virtual ceil or ground] is considered to be [OCCUPIED] 7 | fsm_node: 8 | rog_map: 9 | map_sliding: 10 | enable: false 11 | # The minimum distance [m] to slide the map. 12 | threshold: 0.01 13 | 14 | # If the map rolling is disable, the map origin [m] should be set. 15 | fix_map_origin: [0, 0, 0] 16 | 17 | frontier_extraction_en: false 18 | 19 | resolution: 0.1 20 | inflation_en: false 21 | inflation_resolution: 0.3 22 | inflation_step: 1 23 | map_size: [25, 25, 25] 24 | 25 | # Virtual ceil and ground 26 | virtual_ceil_height: 9999 27 | virtual_ground_height: -9999 28 | 29 | # Intensity filter for input point cloud. If the point intensity is less than this value, 30 | # it will be ignored (to avoid noise like dusk). 31 | intensity_thresh: -1 32 | # The temporal downsample rate for input point cloud. 33 | # e.g. if set to k, the point cloud will be downsampled to 1/k. 34 | point_filt_num: 1 35 | 36 | # Probabilistc update 37 | raycasting: 38 | # if disable, the map will only maintain occupied information, and all other grid 39 | # will be considered as unknown. 40 | enable: true 41 | batch_update_size: 3 42 | local_update_box: !!opencv-matrix 43 | rows: 1 44 | cols: 3 45 | dt: d 46 | data: [30, 30, 30] 47 | # The range of raycasting [m]. 48 | ray_range: !!opencv-matrix 49 | rows: 1 50 | cols: 2 51 | dt: d 52 | data: [0.0, 40] 53 | 54 | # logit: log((x)/(1-(x))) 55 | p_min: 0.12 56 | p_miss: 0.49 57 | p_free: 0.499 58 | p_occ: 0.6 # we set it to 0.6 to make the map more conservative 59 | p_hit: 0.9 60 | p_max: 0.98 61 | # The threshold of a non-occupied inf map grid to be considered as free. [0.0-1.0] 62 | # for counter map, if each cell contains N cells in prob map, then only when 63 | # N_unk > N * unk_thresh cells are unknown, the cell will be considered as unknown. 64 | # example: 1) if unk_thresh is set 0.0, then once at least one cell in prob map is unknown, 65 | # the cell will be considered as unknown. 66 | # 2) if unk_thresh is set to 1.0, then only all cells in prob map are unknown, 67 | # the cell will be considered as unknown. 68 | # in all, 1.0 is more aggressive while 0.0 is more conservative. 69 | inf_map_known_free_thresh: 0.2 70 | 71 | gpu: 72 | GPU_BLOCKSIZE: 512 73 | CLOUD_BUFFER_SIZE: 100000 74 | 75 | astar: 76 | map_voxel_num: [50, 50, 20] 77 | visual_process: false 78 | -------------------------------------------------------------------------------- /config/base.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | 3 | debug: 0 # 0: disable; 1: enable 4 | 5 | device_param: 1 # 0: cpu; 1: gpu 6 | 7 | # tcnn encoder params 8 | n_levels: 16 9 | n_features_per_level: 2 10 | log2_hashmap_size: 19 11 | 12 | # tcnn decoder params 13 | hidden_dim: 64 14 | geo_feat_dim: 14 # geo_feat_dim + k_strc_dim <= 16 / 8 / 4 / 2 or tcnn decoder will become cutlass and crash 15 | geo_num_layer: 3 16 | color_hidden_dim: 64 17 | color_num_layer: 3 18 | 19 | # rim params 20 | trunc_sdf: 1 21 | surface_sample_num: 3 22 | free_sample_num: 3 23 | color_batch_pt_num: 256000 # color render pt batch size 24 | 25 | lr: 5e-3 26 | lr_end: 1e-4 27 | 28 | sdf_weight: 1.0 29 | eikonal_weight: 1e-1 # it will greatly affect structure 30 | curvate_weight: 5e-4 # should be the same loss level to eikonal loss 31 | dist_weight: 1e-1 32 | 33 | rgb_weight: 10.0 34 | 35 | # visualization 36 | vis_frame_step: 10 37 | 38 | export_interval: 1000 # every export_interval frames, the test will be exported 39 | export_colmap_format: 0 # 0: disable; 1: for 3dgs; 2: for nerfstudio colmap 40 | export_train_pcl: 0 # 0: disable; 1: enable -------------------------------------------------------------------------------- /config/fast_livo/2_dao_fast.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.2 12 | fill_level: 1 13 | bce_sigma: 1.0 # if get floaters, increase this value 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 646.78472 22 | fy: 646.65775 23 | cx: 313.456795 24 | cy: 261.399612 25 | 26 | d0: -0.076160 27 | d1: 0.123001 28 | d2: -0.00113 29 | d3: 0.000251 30 | d4: 0.0 31 | 32 | extrinsic: 33 | # lidar to camera 34 | T_C_L: !!opencv-matrix 35 | rows: 4 36 | cols: 4 37 | dt: d 38 | data: [ 0.00610193,-0.999863,-0.0154172, 0.0194384, 39 | -0.00615449,0.0153796,-0.999863, 0.104689, 40 | 0.999962,0.00619598,-0.0060598, -0.0251952, 41 | 0.0, 0.0, 0.0, 1.0 ] 42 | 43 | # lidar to base(imu) 44 | T_B_L: !!opencv-matrix 45 | rows: 4 46 | cols: 4 47 | dt: d 48 | data: [ 1.0, 0.0, 0.0, 0.04165, 49 | 0.0, 1.0, 0.0, 0.02326, 50 | 0.0, 0.0, 1.0, -0.0284, 51 | 0.0, 0.0, 0.0, 1.0 ] 52 | 53 | map: 54 | map_origin: !!opencv-matrix 55 | rows: 1 56 | cols: 3 57 | dt: d 58 | data: [ 0, -10, 20 ] 59 | 60 | map_size: 60 -------------------------------------------------------------------------------- /config/fast_livo/campus.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.2 12 | fill_level: 1 13 | bce_sigma: 1.0 # if get floaters, increase this value 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 588.143714607 22 | fy: 588.107927227 23 | cx: 296.059369138 24 | cy: 254.543215481 25 | 26 | d0: -0.13218037625958456 27 | d1: 0.15360732717073536 28 | d2: 0.00036918417348059815 29 | d3: -0.00031715324469463964 30 | d4: 0.0 31 | 32 | extrinsic: 33 | # lidar to camera 34 | T_C_L: !!opencv-matrix 35 | rows: 4 36 | cols: 4 37 | dt: d 38 | data: [ -0.00200, -0.99975, -0.02211, 0.00260, 39 | -0.00366, 0.02212, -0.99975, 0.05057, 40 | 0.99999, -0.00192, -0.00371, -0.00587, 41 | 0.0, 0.0, 0.0, 1.0 ] 42 | 43 | # lidar to base(imu) 44 | T_B_L: !!opencv-matrix 45 | rows: 4 46 | cols: 4 47 | dt: d 48 | data: [ 1.0, 0.0, 0.0, 0.04165, 49 | 0.0, 1.0, 0.0, 0.02326, 50 | 0.0, 0.0, 1.0, -0.0284, 51 | 0.0, 0.0, 0.0, 1.0 ] 52 | 53 | map: 54 | map_origin: !!opencv-matrix 55 | rows: 1 56 | cols: 3 57 | dt: d 58 | data: [ 20, 0, 20 ] 59 | 60 | map_size: 50 -------------------------------------------------------------------------------- /config/fast_livo/cbd.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.2 12 | fill_level: 1 13 | bce_sigma: 1.0 # if get floaters, increase this value 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 588.143714607 22 | fy: 588.107927227 23 | cx: 296.059369138 24 | cy: 254.543215481 25 | 26 | d0: -0.13218037625958456 27 | d1: 0.15360732717073536 28 | d2: 0.00036918417348059815 29 | d3: -0.00031715324469463964 30 | d4: 0.0 31 | 32 | extrinsic: 33 | # lidar to camera 34 | T_C_L: !!opencv-matrix 35 | rows: 4 36 | cols: 4 37 | dt: d 38 | data: [ -0.00200, -0.99975, -0.02211, 0.00260, 39 | -0.00366, 0.02212, -0.99975, 0.05057, 40 | 0.99999, -0.00192, -0.00371, -0.00587, 41 | 0.0, 0.0, 0.0, 1.0 ] 42 | 43 | # lidar to base(imu) 44 | T_B_L: !!opencv-matrix 45 | rows: 4 46 | cols: 4 47 | dt: d 48 | data: [ 1.0, 0.0, 0.0, 0.04165, 49 | 0.0, 1.0, 0.0, 0.02326, 50 | 0.0, 0.0, 1.0, -0.0284, 51 | 0.0, 0.0, 0.0, 1.0 ] 52 | 53 | map: 54 | map_origin: !!opencv-matrix 55 | rows: 1 56 | cols: 3 57 | dt: d 58 | data: [ 0, -25, 0 ] 59 | 60 | map_size: 120 -------------------------------------------------------------------------------- /config/fast_livo/culture01.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.4 12 | fill_level: 1 13 | bce_sigma: 1.0 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 588.143714607 22 | fy: 588.107927227 23 | cx: 296.059369138 24 | cy: 254.543215481 25 | 26 | d0: -0.13218037625958456 27 | d1: 0.15360732717073536 28 | d2: 0.00036918417348059815 29 | d3: -0.00031715324469463964 30 | d4: 0.0 31 | 32 | 33 | extrinsic: 34 | # lidar to camera 35 | T_C_L: !!opencv-matrix 36 | rows: 4 37 | cols: 4 38 | dt: d 39 | data: [ -0.00200, -0.99975, -0.02211, 0.00260, 40 | -0.00366, 0.02212, -0.99975, 0.05057, 41 | 0.99999, -0.00192, -0.00371, -0.00587, 42 | 0.0, 0.0, 0.0, 1.0 ] 43 | 44 | # lidar to base(imu) 45 | T_B_L: !!opencv-matrix 46 | rows: 4 47 | cols: 4 48 | dt: d 49 | data: [ 1.0, 0.0, 0.0, 0.04165, 50 | 0.0, 1.0, 0.0, 0.02326, 51 | 0.0, 0.0, 1.0, -0.0284, 52 | 0.0, 0.0, 0.0, 1.0 ] 53 | 54 | map: 55 | map_origin: !!opencv-matrix 56 | rows: 1 57 | cols: 3 58 | dt: d 59 | data: [ 5, 20, 20 ] 60 | 61 | map_size: 110 -------------------------------------------------------------------------------- /config/fast_livo/drive.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.4 12 | fill_level: 1 13 | bce_sigma: 1.0 # if get floaters, increase this value 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 588.143714607 22 | fy: 588.107927227 23 | cx: 296.059369138 24 | cy: 254.543215481 25 | 26 | d0: -0.13218037625958456 27 | d1: 0.15360732717073536 28 | d2: 0.00036918417348059815 29 | d3: -0.00031715324469463964 30 | d4: 0.0 31 | 32 | extrinsic: 33 | # lidar to camera 34 | T_C_L: !!opencv-matrix 35 | rows: 4 36 | cols: 4 37 | dt: d 38 | data: [ -0.00200, -0.99975, -0.02211, 0.00260, 39 | -0.00366, 0.02212, -0.99975, 0.05057, 40 | 0.99999, -0.00192, -0.00371, -0.00587, 41 | 0.0, 0.0, 0.0, 1.0 ] 42 | 43 | # lidar to base(imu) 44 | T_B_L: !!opencv-matrix 45 | rows: 4 46 | cols: 4 47 | dt: d 48 | data: [ 1.0, 0.0, 0.0, 0.04165, 49 | 0.0, 1.0, 0.0, 0.02326, 50 | 0.0, 0.0, 1.0, -0.0284, 51 | 0.0, 0.0, 0.0, 1.0 ] 52 | 53 | map: 54 | map_origin: !!opencv-matrix 55 | rows: 1 56 | cols: 3 57 | dt: d 58 | data: [ 80, 0, 0 ] 59 | 60 | map_size: 180 -------------------------------------------------------------------------------- /config/fast_livo/fast_livo.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | 3 | base_config: "../base.yaml" 4 | 5 | preload: 1 # 0: disable; 1: enable # accelerate the loading process but will cost more memory 6 | llff: 1 # 0: disable; 1: enable; every 8 frame will be used for evaluation 7 | prob_map_en: 1 # 0: disable; 1: enable; Whether the map is a probabilistic map 8 | 9 | # dataset_type 10 | # Replica = 0, 11 | # R3live = 1, 12 | # NeuralRGBD = 2, 13 | # KITTI = 3, 14 | # FastLivo = 4, 15 | dataset_type: 4 16 | dir_embedding_degree: 4 17 | 18 | min_range: 0.01 19 | max_range: 100 20 | ds_pt_num: 10000 # downsampled point number 21 | max_pt_num: 1000000 22 | 23 | outlier_remove: 1 # unnecessary for static scenes 24 | outlier_dist: 0.05 25 | outlier_removal_interval: 2000 26 | 27 | # visualization 28 | vis_attribute: 2 # 0: disable to save storage; 1: normal; 2: color 29 | vis_resolution: 0.1 # better no more than leaf_sizes or will miss faces 30 | export_resolution: 0.04 31 | fps: 10 -------------------------------------------------------------------------------- /config/fast_livo/red_bird.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.4 12 | fill_level: 1 13 | bce_sigma: 1.0 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 646.294856087 22 | fy: 646.155780233 23 | cx: 313.423980506 24 | cy: 262.903356268 25 | 26 | d0: -0.07581043635239915 27 | d1: 0.1276695531950238 28 | d2: -0.0004921162736736936 29 | d3: 3.265638595122261e-05 30 | d4: 0.0 31 | 32 | extrinsic: 33 | # lidar to camera 34 | T_C_L: !!opencv-matrix 35 | rows: 4 36 | cols: 4 37 | dt: d 38 | data: [ 0.000825355,-0.999798,-0.0200863, 0.0147334, 39 | -0.00363991,0.0200832,-0.999792, 0.0700588, 40 | 0.999993,0.000898296,-0.0036226, -0.0209133, 41 | 0.0, 0.0, 0.0, 1.0 ] 42 | 43 | # lidar to base(imu) 44 | T_B_L: !!opencv-matrix 45 | rows: 4 46 | cols: 4 47 | dt: d 48 | data: [ 1.0, 0.0, 0.0, 0.04165, 49 | 0.0, 1.0, 0.0, 0.02326, 50 | 0.0, 0.0, 1.0, -0.0284, 51 | 0.0, 0.0, 0.0, 1.0 ] 52 | 53 | map: 54 | map_origin: !!opencv-matrix 55 | rows: 1 56 | cols: 3 57 | dt: d 58 | data: [ 20, 0, 20 ] 59 | map_size: 100 -------------------------------------------------------------------------------- /config/fast_livo/station.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.2 12 | fill_level: 1 13 | bce_sigma: 1.0 # if get floaters, increase this value 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | width: 640 19 | height: 512 20 | 21 | fx: 646.78472 22 | fy: 646.65775 23 | cx: 313.456795 24 | cy: 261.399612 25 | 26 | d0: -0.076160 27 | d1: 0.123001 28 | d2: -0.00113 29 | d3: 0.000251 30 | d4: 0.0 31 | 32 | extrinsic: 33 | # lidar to camera 34 | T_C_L: !!opencv-matrix 35 | rows: 4 36 | cols: 4 37 | dt: d 38 | data: [ 0.00610193,-0.999863,-0.0154172, 0.0194384, 39 | -0.00615449,0.0153796,-0.999863, 0.104689, 40 | 0.999962,0.00619598,-0.0060598, -0.0251952, 41 | 0.0, 0.0, 0.0, 1.0 ] 42 | 43 | # lidar to base(imu) 44 | T_B_L: !!opencv-matrix 45 | rows: 4 46 | cols: 4 47 | dt: d 48 | data: [ 1.0, 0.0, 0.0, 0.04165, 49 | 0.0, 1.0, 0.0, 0.02326, 50 | 0.0, 0.0, 1.0, -0.0284, 51 | 0.0, 0.0, 0.0, 1.0 ] 52 | 53 | map: 54 | map_origin: !!opencv-matrix 55 | rows: 1 56 | cols: 3 57 | dt: d 58 | data: [ 0, -10, 20 ] 59 | 60 | map_size: 60 -------------------------------------------------------------------------------- /config/fast_livo/sysu.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/rim/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | scene_config: "fast_livo.yaml" 6 | 7 | iter_step: 20000 8 | 9 | trace_iter: 100 # sphere tracing iteration times 10 | 11 | leaf_sizes: 0.2 12 | fill_level: 1 13 | bce_sigma: 1.0 # if get floaters, increase this value 14 | sphere_trace_thr: 1e-3 15 | 16 | camera: 17 | model: 0 # 0: pinhole; 1: equidistant 18 | 19 | width: 1280 20 | height: 1024 21 | 22 | fx: 1311.89517127580 23 | fy: 1311.36488586115 24 | cx: 656.523841857393 25 | cy: 504.136322840350 26 | 27 | d0: -0.0780830982640722 28 | d1: 0.146382433670493 29 | d2: -0.00110111633050301 30 | d3: -0.00110752991013068 31 | d4: 0.0 32 | 33 | extrinsic: 34 | # lidar to camera 35 | T_C_L: !!opencv-matrix 36 | rows: 4 37 | cols: 4 38 | dt: d 39 | data: [ -0.0036250, -0.9998907, -0.0143360,0.00549469, 40 | 0.0075568, 0.0143083, -0.9998690,0.0712101, 41 | 0.9999649, -0.0037329, 0.0075041, 0.0322054, 42 | 0.0, 0.0, 0.0, 1.0 ] 43 | 44 | # lidar to base(imu) 45 | T_B_L: !!opencv-matrix 46 | rows: 4 47 | cols: 4 48 | dt: d 49 | data: [ 1.0, 0.0, 0.0, 0.04165, 50 | 0.0, 1.0, 0.0, 0.02326, 51 | 0.0, 0.0, 1.0, -0.0284, 52 | 0.0, 0.0, 0.0, 1.0 ] 53 | 54 | map: 55 | map_origin: !!opencv-matrix 56 | rows: 1 57 | cols: 3 58 | dt: d 59 | data: [ 9, -2, 20 ] 60 | 61 | map_size: 60 -------------------------------------------------------------------------------- /config/replica/replica.yaml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | # rosrun neural_mapping neural_mapping_node src/RIM2/neural_mapping/config/neural_rgbd/neural_rgbd.yaml src/RIM2/data/neural_rgbd_data/kitchen_kitti_format 3 | 4 | base_config: "../base.yaml" 5 | 6 | iter_step: 20000 7 | 8 | trace_iter: 50 # sphere tracing iteration times 9 | 10 | preload: 1 # 0: disable; 1: enable # accelerate the loading process but will cost more memory 11 | llff: 0 # 0: disable; 1: enable; every 8 frame will be used for evaluation 12 | cull_mesh: 1 13 | prob_map_en: 0 # 0: disable; 1: enable; Whether the map is a probabilistic map 14 | 15 | # dataset_type 16 | # Replica = 0, 17 | # R3live = 1, 18 | # NeuralRGBD = 2, 19 | # KITTI = 3, 20 | # FastLivo = 4, 21 | dataset_type: 0 22 | 23 | dir_embedding_degree: 0 24 | 25 | map: 26 | map_origin: !!opencv-matrix 27 | rows: 1 28 | cols: 3 29 | dt: d 30 | data: [ 0, 0, 0 ] 31 | map_size: 14 32 | 33 | min_range: 0.0 34 | max_range: 100 35 | ds_pt_num: 10000 # downsampled point number 36 | max_pt_num: 1000000 37 | 38 | leaf_sizes: 0.05 39 | fill_level: 1 40 | bce_sigma: 0.02 41 | sphere_trace_thr: 1e-3 42 | 43 | outlier_removal_interval: 2000 44 | outlier_remove: 0 # unnecessary for static scenes 45 | outlier_dist: 0.05 46 | 47 | # visualization 48 | vis_attribute: 2 # 0: disable to save storage; 1: normal; 2: color 49 | vis_resolution: 0.1 # better no more than leaf_sizes or will miss faces 50 | export_resolution: 0.01 51 | fps: 30 -------------------------------------------------------------------------------- /docs/FAST_LIVO2_Data_Training.md: -------------------------------------------------------------------------------- 1 | Here is a instrution for the FAST LIVO2 data format traning. 2 | 3 | 1. Build the [modified FAST-LIVO2](https://github.com/jianhengLiu/FAST-LIVO2.git) with the same instruction as the official repo. 4 | 2. Collect or Download datasets from [FAST-LIVO2-Dataset](https://connecthkuhk-my.sharepoint.com/personal/zhengcr_connect_hku_hk/_layouts/15/onedrive.aspx?id=%2Fpersonal%2Fzhengcr%5Fconnect%5Fhku%5Fhk%2FDocuments%2Ffast%2Dlivo2%2Ddataset&ga=1). 5 | 3. https://github.com/jianhengLiu/FAST-LIVO2.git 6 | 4. open a terminal to start LIVO 7 | ```bash 8 | roslaunch fast_livo mapping_avia.launch 9 | ``` 10 | 5. open another terminal to get ready for bag recording 11 | ```bash 12 | rosbag record /aft_mapped_to_init /origin_img /cloud_registered_body /tf /tf_static /path -O "fast_livo2_YOUR_DOWNLOADED" -b 2048 13 | ``` 14 | 6. open another terminal to play your downloaded/collected bag 15 | ```bash 16 | rosbag play YOUR_DOWNLOADED.bag 17 | ``` 18 | 7. `Ctrl+C` to stop recording when you finish the bag recording. 19 | 8. Train M2Mapping model with the following command: 20 | ```bash 21 | source devel/setup.bash # or setup.zsh 22 | # ./src/M2Mapping/build/neural_mapping_node train [config_path] [dataset_path] 23 | # For example: 24 | ./src/M2Mapping/build/neural_mapping_node train src/M2Mapping/config/fast_livo/campus.yaml src/M2Mapping/data/FAST_LIVO2_RIM_Datasets/campus/fast_livo2_campus.bag 25 | ``` 26 | After the training the results will be saved in the `src/M2Mapping/output/dae-dataset_name` directory. -------------------------------------------------------------------------------- /eval/draw.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | linestyle = '-' 6 | linewidth = 2 7 | marker = 'o' 8 | markersize = 3 9 | markersize = 5 10 | 11 | # recon quantity maicity 12 | x = np.array([10, 15, 20, 25, 30, 35, 40, 45, 50]) 13 | chamfer_loss = np.array([1.740137573,1.652615321,1.626417989,1.598619474,1.580639635,1.57094054,1.568409511,1.559219683,1.562328578]) 14 | f1_score = np.array([98.747425629,98.884420986,98.911179736,98.923287838,98.95279367,98.942050761,98.924957865,98.928133584,98.923456473]) 15 | timing = np.array([85.01,128.591,163.551,214.483,252.910333333,298.99,332.655,368.883,394.071]) 16 | 17 | a = 4 18 | fig, ax = plt.subplots(figsize=(8, 5)) 19 | plt.rcParams['ytick.direction'] = 'in' # 刻度线显示在内部 20 | plt.rcParams['xtick.direction'] = 'in' # 刻度线显示在内部 21 | plt.subplots_adjust(left=None, bottom=None, right=0.8, 22 | top=0.95, wspace=0.3, hspace=0.3) 23 | 24 | 25 | ax.plot(x, chamfer_loss, label="C-l1", 26 | linestyle=linestyle, linewidth=linewidth, marker=marker, markersize=markersize, color='black') 27 | ax.set_ylim(ymin=1.5, ymax=1.75) 28 | ax.set_ylabel("Chamfer-L1(cm)", fontdict={ 29 | 'family': 'Times New Roman', 'size': 16}) 30 | ax.set_xlabel("Reconstruction Iter", fontdict={ 31 | 'family': 'Times New Roman', 'size': 16}) 32 | ax.spines['right'].set_visible(False) # ax右轴隐藏 33 | legend_font = {"family": "Times New Roman", "size": 11} 34 | # ax.legend(prop=legend_font) 35 | 36 | bx_color = 'red' 37 | bx = ax.twinx() 38 | bx.plot(x, f1_score, 's--', color=bx_color) 39 | bx.set_ylim(ymin=98.7, ymax=99) 40 | bx.set_ylabel("F1-Score(%)", fontdict={ 41 | 'family': 'Times New Roman', 'size': 16}, color=bx_color) 42 | bx.tick_params(length=6, width=2, color=bx_color, labelcolor=bx_color) 43 | bx.spines['right'].set(color=bx_color, linewidth=2.0, linestyle=':') 44 | # bx.legend(prop=legend_font) 45 | 46 | cx_color = 'blue' 47 | cx = ax.twinx() 48 | cx.plot(x, timing, 'o-.', color=cx_color) 49 | cx.set_ylim(ymin=50, ymax=450) 50 | cx.set_ylabel("Timing(ms)", fontdict={ 51 | 'family': 'Times New Roman', 'size': 16}, color=cx_color) 52 | cx.tick_params(length=6, width=2, color=cx_color, labelcolor=cx_color) 53 | cx.spines['right'].set(color=cx_color, linewidth=2.0, linestyle='-.') 54 | cx.spines['right'].set_position(('axes', 1.15)) 55 | # cx.legend(prop=legend_font) 56 | 57 | # plt.plot(x, f1_score, label="F1-Score", 58 | # linestyle=linestyle, linewidth=linewidth, marker=marker, markersize=markersize) 59 | 60 | # plt.xlim(0, 11) 61 | # plt.xticks(np.linspace(0, 11, 12, endpoint=True)) 62 | # plt.ylim(0.0, 1.0) 63 | # plt.yticks(np.linspace(0, 1, 5, endpoint=True)) 64 | 65 | 66 | # plt.title('Surface Coverage Evaluation') 67 | 68 | plt.show() 69 | -------------------------------------------------------------------------------- /eval/draw_loss.py: -------------------------------------------------------------------------------- 1 | # read output/mem_usage.txt and draw the memory usage curve 2 | # Usage: python3 draw_mem.py 3 | 4 | from argparse import ArgumentParser 5 | import matplotlib.pyplot as plt 6 | from pathlib import Path 7 | 8 | 9 | def draw(args): 10 | log_file = args.log_file 11 | log_file = Path(log_file[0]) 12 | log_file_name = log_file.name.replace(".txt", "") 13 | # make directory to save the loss curve 14 | save_fig_paths = log_file.parent / log_file_name 15 | save_fig_paths.mkdir(exist_ok=True) 16 | 17 | linestyle = "-" 18 | linewidth = 1 19 | marker = "o" 20 | markersize = 1 21 | 22 | # read output/mem_usage.txt and 23 | titles = [] 24 | x = [] 25 | with open(log_file, "r") as f: 26 | for title in f.readline().split(): 27 | titles.append(title) 28 | x.append([]) 29 | 30 | for line in f: 31 | line = line.strip() 32 | if line and (len(line.split()) == len(titles)): 33 | j = 0 34 | for attr in line.split(): 35 | x[j].append(float(attr)) 36 | j = j + 1 37 | 38 | iter_range = list(range(len(x[0]))) 39 | for j in range(len(titles)): 40 | # find the min max to normalize the loss 41 | # min_loss = min(x[j + 1]) 42 | # max_loss = max(x[j + 1]) 43 | # loss = [((i - min_loss) / (max_loss - min_loss)) for i in x[j + 1]] 44 | 45 | loss = x[j] 46 | plt.cla() 47 | plt.plot(iter_range, loss, label=titles[j]) 48 | 49 | plt.legend(prop={"size": 11}) 50 | if args.vis: 51 | plt.show() 52 | # plt.show() 53 | plt.savefig(log_file.parent / log_file_name / (titles[j] + ".png")) 54 | # plt.plot(x, cpu_mem, "-.", label="System Memory Usage", color="red") 55 | 56 | # plt.plot(x, gpu_mem, "-.", label="Video Memory Usage", color="blue") 57 | # plt.xlabel("Frame Index", fontdict={"size": 13}) 58 | # plt.ylabel("Memory Usage (GB)", fontdict={"size": 13}) 59 | # plt.ylim(ymin=0, ymax=12) 60 | # plt.xlim(xmin=0, xmax=4540) 61 | 62 | # plt.legend(prop={"size": 11}) 63 | # plt.show() 64 | # plt.savefig(model_paths / "loss.png", dpi=300) 65 | 66 | 67 | if __name__ == "__main__": 68 | 69 | # Set up command line argument parser 70 | parser = ArgumentParser(description="Training script parameters") 71 | parser.add_argument( 72 | "--log_file", "-l", required=True, nargs="+", type=str, default=[] 73 | ) 74 | parser.add_argument("--vis", "-v", type=bool, default=False) 75 | args = parser.parse_args() 76 | draw(args) 77 | -------------------------------------------------------------------------------- /eval/draw_mem.py: -------------------------------------------------------------------------------- 1 | # read output/mem_usage.txt and draw the memory usage curve 2 | # Usage: python3 draw_mem.py 3 | 4 | from argparse import ArgumentParser 5 | import matplotlib.pyplot as plt 6 | from pathlib import Path 7 | 8 | 9 | def draw(model_paths): 10 | model_paths = Path(model_paths[0]) 11 | linestyle = "-" 12 | linewidth = 1 13 | marker = "o" 14 | markersize = 1 15 | 16 | # read output/mem_usage.txt and 17 | x = [] 18 | cpu_mem = [] 19 | gpu_mem = [] 20 | with open(model_paths / "mem_usage.txt", "r") as f: 21 | f.readline() 22 | for line in f: 23 | line = line.strip() 24 | if line: 25 | x.append(int(line.split()[0])) 26 | cpu_mem.append(float(line.split()[1])) 27 | gpu_mem.append(float(line.split()[2])) 28 | 29 | plt.plot(x, cpu_mem, "-.", label="System Memory Usage", color="red") 30 | 31 | plt.plot(x, gpu_mem, "-.", label="Video Memory Usage", color="blue") 32 | plt.xlabel("Frame Index", fontdict={"size": 13}) 33 | plt.ylabel("Memory Usage (GB)", fontdict={"size": 13}) 34 | # plt.ylim(ymin=0, ymax=12) 35 | # plt.xlim(xmin=0, xmax=4540) 36 | 37 | plt.legend(prop={"size": 11}) 38 | # plt.show() 39 | plt.savefig(model_paths / "mem_usage.png", dpi=300) 40 | 41 | 42 | if __name__ == "__main__": 43 | 44 | # Set up command line argument parser 45 | parser = ArgumentParser(description="Training script parameters") 46 | parser.add_argument( 47 | "--model_paths", "-m", required=True, nargs="+", type=str, default=[] 48 | ) 49 | args = parser.parse_args() 50 | draw(args.model_paths) 51 | -------------------------------------------------------------------------------- /eval/image_metrics/eval_image.sh: -------------------------------------------------------------------------------- 1 | if [ "$#" -ne 1 ]; then 2 | echo "Usage: $0 " 3 | exit 1 4 | fi 5 | 6 | model=$1 7 | 8 | # find parent_path according to $0 9 | parent_path=$(dirname "$0") 10 | python $parent_path/metrics.py -m $model/train/color 11 | python $parent_path/metrics.py -m $model/eval/color 12 | -------------------------------------------------------------------------------- /eval/image_metrics/image_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | 14 | def mse(img1, img2): 15 | return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 16 | 17 | def psnr(img1, img2): 18 | mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 19 | return 20 * torch.log10(1.0 / torch.sqrt(mse)) 20 | -------------------------------------------------------------------------------- /eval/image_metrics/loss_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import torch.nn.functional as F 14 | from torch.autograd import Variable 15 | from math import exp 16 | 17 | 18 | def l1_loss(network_output, gt): 19 | return torch.abs((network_output - gt)).mean() 20 | 21 | 22 | def l2_loss(network_output, gt): 23 | return ((network_output - gt) ** 2).mean() 24 | 25 | 26 | def gaussian(window_size, sigma): 27 | gauss = torch.Tensor( 28 | [ 29 | exp(-((x - window_size // 2) ** 2) / float(2 * sigma**2)) 30 | for x in range(window_size) 31 | ] 32 | ) 33 | return gauss / gauss.sum() 34 | 35 | 36 | def create_window(window_size, channel): 37 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 38 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 39 | window = Variable( 40 | _2D_window.expand(channel, 1, window_size, window_size).contiguous() 41 | ) 42 | return window 43 | 44 | 45 | def ssim(img1, img2, window_size=11, size_average=True): 46 | # data range: [0,1] 47 | channel = img1.size(-3) 48 | window = create_window(window_size, channel) 49 | 50 | if img1.is_cuda: 51 | window = window.cuda(img1.get_device()) 52 | window = window.type_as(img1) 53 | 54 | return _ssim(img1, img2, window, window_size, channel, size_average) 55 | 56 | 57 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 58 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 59 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 60 | 61 | mu1_sq = mu1.pow(2) 62 | mu2_sq = mu2.pow(2) 63 | mu1_mu2 = mu1 * mu2 64 | 65 | sigma1_sq = ( 66 | F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 67 | ) 68 | sigma2_sq = ( 69 | F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 70 | ) 71 | sigma12 = ( 72 | F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) 73 | - mu1_mu2 74 | ) 75 | 76 | C1 = 0.01**2 77 | C2 = 0.03**2 78 | 79 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ( 80 | (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2) 81 | ) 82 | 83 | if size_average: 84 | return ssim_map.mean() 85 | else: 86 | return ssim_map.mean(1).mean(1).mean(1) 87 | -------------------------------------------------------------------------------- /eval/image_metrics/lpipsPyTorch/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .modules.lpips import LPIPS 4 | 5 | 6 | class Lpips: 7 | def __init__(self, net_type: str = "alex", version: str = "0.1"): 8 | self.criterion = LPIPS(net_type, version).cuda() 9 | 10 | def forward(self, x: torch.Tensor, y: torch.Tensor): 11 | r"""Function that measures 12 | Learned Perceptual Image Patch Similarity (LPIPS). 13 | 14 | Arguments: 15 | x, y (torch.Tensor): the input tensors to compare. 16 | net_type (str): the network type to compare the features: 17 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 18 | version (str): the version of LPIPS. Default: 0.1. 19 | """ 20 | return self.criterion(x, y) 21 | -------------------------------------------------------------------------------- /eval/image_metrics/lpipsPyTorch/modules/lpips.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from .networks import get_network, LinLayers 5 | from .utils import get_state_dict 6 | 7 | 8 | class LPIPS(nn.Module): 9 | r"""Creates a criterion that measures 10 | Learned Perceptual Image Patch Similarity (LPIPS). 11 | 12 | Arguments: 13 | net_type (str): the network type to compare the features: 14 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 15 | version (str): the version of LPIPS. Default: 0.1. 16 | """ 17 | def __init__(self, net_type: str = 'alex', version: str = '0.1'): 18 | 19 | assert version in ['0.1'], 'v0.1 is only supported now' 20 | 21 | super(LPIPS, self).__init__() 22 | 23 | # pretrained network 24 | self.net = get_network(net_type) 25 | 26 | # linear layers 27 | self.lin = LinLayers(self.net.n_channels_list) 28 | self.lin.load_state_dict(get_state_dict(net_type, version)) 29 | 30 | def forward(self, x: torch.Tensor, y: torch.Tensor): 31 | feat_x, feat_y = self.net(x), self.net(y) 32 | 33 | diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)] 34 | res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)] 35 | 36 | return torch.sum(torch.cat(res, 0), 0, True) 37 | -------------------------------------------------------------------------------- /eval/image_metrics/lpipsPyTorch/modules/networks.py: -------------------------------------------------------------------------------- 1 | from typing import Sequence 2 | 3 | from itertools import chain 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torchvision import models 8 | 9 | from .utils import normalize_activation 10 | 11 | 12 | def get_network(net_type: str): 13 | if net_type == 'alex': 14 | return AlexNet() 15 | elif net_type == 'squeeze': 16 | return SqueezeNet() 17 | elif net_type == 'vgg': 18 | return VGG16() 19 | else: 20 | raise NotImplementedError('choose net_type from [alex, squeeze, vgg].') 21 | 22 | 23 | class LinLayers(nn.ModuleList): 24 | def __init__(self, n_channels_list: Sequence[int]): 25 | super(LinLayers, self).__init__([ 26 | nn.Sequential( 27 | nn.Identity(), 28 | nn.Conv2d(nc, 1, 1, 1, 0, bias=False) 29 | ) for nc in n_channels_list 30 | ]) 31 | 32 | for param in self.parameters(): 33 | param.requires_grad = False 34 | 35 | 36 | class BaseNet(nn.Module): 37 | def __init__(self): 38 | super(BaseNet, self).__init__() 39 | 40 | # register buffer 41 | self.register_buffer( 42 | 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) 43 | self.register_buffer( 44 | 'std', torch.Tensor([.458, .448, .450])[None, :, None, None]) 45 | 46 | def set_requires_grad(self, state: bool): 47 | for param in chain(self.parameters(), self.buffers()): 48 | param.requires_grad = state 49 | 50 | def z_score(self, x: torch.Tensor): 51 | return (x - self.mean) / self.std 52 | 53 | def forward(self, x: torch.Tensor): 54 | x = self.z_score(x) 55 | 56 | output = [] 57 | for i, (_, layer) in enumerate(self.layers._modules.items(), 1): 58 | x = layer(x) 59 | if i in self.target_layers: 60 | output.append(normalize_activation(x)) 61 | if len(output) == len(self.target_layers): 62 | break 63 | return output 64 | 65 | 66 | class SqueezeNet(BaseNet): 67 | def __init__(self): 68 | super(SqueezeNet, self).__init__() 69 | 70 | self.layers = models.squeezenet1_1(True).features 71 | self.target_layers = [2, 5, 8, 10, 11, 12, 13] 72 | self.n_channels_list = [64, 128, 256, 384, 384, 512, 512] 73 | 74 | self.set_requires_grad(False) 75 | 76 | 77 | class AlexNet(BaseNet): 78 | def __init__(self): 79 | super(AlexNet, self).__init__() 80 | 81 | self.layers = models.alexnet(True).features 82 | self.target_layers = [2, 5, 8, 10, 12] 83 | self.n_channels_list = [64, 192, 384, 256, 256] 84 | 85 | self.set_requires_grad(False) 86 | 87 | 88 | class VGG16(BaseNet): 89 | def __init__(self): 90 | super(VGG16, self).__init__() 91 | 92 | self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features 93 | self.target_layers = [4, 9, 16, 23, 30] 94 | self.n_channels_list = [64, 128, 256, 512, 512] 95 | 96 | self.set_requires_grad(False) 97 | -------------------------------------------------------------------------------- /eval/image_metrics/lpipsPyTorch/modules/utils.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | 3 | import torch 4 | 5 | 6 | def normalize_activation(x, eps=1e-10): 7 | norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True)) 8 | return x / (norm_factor + eps) 9 | 10 | 11 | def get_state_dict(net_type: str = 'alex', version: str = '0.1'): 12 | # build url 13 | url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ 14 | + f'master/lpips/weights/v{version}/{net_type}.pth' 15 | 16 | # download 17 | old_state_dict = torch.hub.load_state_dict_from_url( 18 | url, progress=True, 19 | map_location=None if torch.cuda.is_available() else torch.device('cpu') 20 | ) 21 | 22 | # rename keys 23 | new_state_dict = OrderedDict() 24 | for key, val in old_state_dict.items(): 25 | new_key = key 26 | new_key = new_key.replace('lin', '') 27 | new_key = new_key.replace('model.', '') 28 | new_state_dict[new_key] = val 29 | 30 | return new_state_dict 31 | -------------------------------------------------------------------------------- /eval/image_metrics/metrics.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from pathlib import Path 13 | import os 14 | from PIL import Image 15 | import torch 16 | import torchvision.transforms.functional as tf 17 | from loss_utils import ssim 18 | from lpipsPyTorch import Lpips 19 | import json 20 | from tqdm import tqdm 21 | from image_utils import psnr 22 | from argparse import ArgumentParser 23 | 24 | 25 | def evaluate_render(renders_dir, gt_dir, fname, lpips_inferer: Lpips): 26 | render = Image.open(renders_dir / fname) 27 | gt = Image.open(gt_dir / fname) 28 | render = tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda() 29 | gt = tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda() 30 | 31 | # [1, chaneel, H, W] 32 | return ( 33 | fname, 34 | ssim(render, gt), 35 | psnr(render, gt), 36 | lpips_inferer.forward(render, gt), 37 | ) 38 | 39 | 40 | def evaluate(model_paths): 41 | 42 | full_dict = {} 43 | per_view_dict = {} 44 | full_dict_polytopeonly = {} 45 | per_view_dict_polytopeonly = {} 46 | 47 | for scene_dir in model_paths: 48 | try: 49 | full_dict[scene_dir] = {} 50 | per_view_dict[scene_dir] = {} 51 | full_dict_polytopeonly[scene_dir] = {} 52 | per_view_dict_polytopeonly[scene_dir] = {} 53 | 54 | # color_dir = Path(scene_dir) / "color" 55 | scene_dir_path = Path(scene_dir) 56 | 57 | gt_color_dir = scene_dir_path / "gt" 58 | renders_color_dir = scene_dir_path / "renders" 59 | 60 | image_names = [] 61 | ssims = [] 62 | psnrs = [] 63 | lpipss = [] 64 | 65 | lpips_inferer = Lpips(net_type="vgg") 66 | list = os.listdir(renders_color_dir) 67 | if len(list) == 0: 68 | print("No renders found in", renders_color_dir) 69 | continue 70 | for fname in tqdm(os.listdir(renders_color_dir)): 71 | fname, ssim, psnr, lpips = evaluate_render( 72 | renders_color_dir, gt_color_dir, fname, lpips_inferer 73 | ) 74 | 75 | image_names.append(fname) 76 | ssims.append(ssim) 77 | psnrs.append(psnr) 78 | lpipss.append(lpips) 79 | 80 | print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 81 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 82 | print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 83 | # print max PSNR and min PSNR and their filename 84 | max_psnr = torch.tensor(psnrs).max() 85 | min_psnr = torch.tensor(psnrs).min() 86 | max_psnr_index = torch.tensor(psnrs).argmax() 87 | min_psnr_index = torch.tensor(psnrs).argmin() 88 | 89 | print( 90 | " Max PSNR: {:>12.7f} for {}".format( 91 | max_psnr, image_names[max_psnr_index] 92 | ) 93 | ) 94 | print( 95 | " Min PSNR: {:>12.7f} for {}".format( 96 | min_psnr, image_names[min_psnr_index] 97 | ) 98 | ) 99 | print("") 100 | 101 | full_dict[scene_dir].update( 102 | { 103 | "SSIM": torch.tensor(ssims).mean().item(), 104 | "PSNR": torch.tensor(psnrs).mean().item(), 105 | "LPIPS": torch.tensor(lpipss).mean().item(), 106 | "Max PSNR {}:".format(image_names[max_psnr_index]): max_psnr.item(), 107 | "Min PSNR {}:".format(image_names[min_psnr_index]): min_psnr.item(), 108 | } 109 | ) 110 | per_view_dict[scene_dir].update( 111 | { 112 | "SSIM": { 113 | name: ssim 114 | for ssim, name in zip(torch.tensor(ssims).tolist(), image_names) 115 | }, 116 | "PSNR": { 117 | name: psnr 118 | for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names) 119 | }, 120 | "LPIPS": { 121 | name: lp 122 | for lp, name in zip(torch.tensor(lpipss).tolist(), image_names) 123 | }, 124 | } 125 | ) 126 | 127 | with open(scene_dir_path / "../render_eval.json", "w") as fp: 128 | json.dump(full_dict, fp, indent=True) 129 | with open(scene_dir_path / "../render_eval_per_view.json", "w") as fp: 130 | json.dump(per_view_dict[scene_dir], fp, indent=True) 131 | with open(scene_dir_path / "../../evaluation_results.json", "a") as fp: 132 | full_dict[scene_dir] = { 133 | k: round(v, 3) for k, v in full_dict[scene_dir].items() 134 | } 135 | json.dump(full_dict, fp, indent=True) 136 | fp.write("\n") 137 | 138 | with open(scene_dir_path / "../../../all_evaluation_results.json", "a") as fp: 139 | # add dataset name 140 | json.dump(full_dict, fp, indent=True) 141 | fp.write("\n") 142 | except Exception as e: 143 | print("Unable to compute metrics for model", scene_dir_path) 144 | print(e) 145 | 146 | 147 | if __name__ == "__main__": 148 | device = torch.device("cuda:0") 149 | torch.cuda.set_device(device) 150 | 151 | # Set up command line argument parser 152 | parser = ArgumentParser(description="Training script parameters") 153 | parser.add_argument( 154 | "--model_paths", "-m", required=True, nargs="+", type=str, default=[] 155 | ) 156 | args = parser.parse_args() 157 | evaluate(args.model_paths) 158 | -------------------------------------------------------------------------------- /eval/image_metrics/metrics2.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from pathlib import Path 13 | import os 14 | from PIL import Image 15 | import torch 16 | import torchvision.transforms.functional as tf 17 | from loss_utils import ssim 18 | from lpipsPyTorch import Lpips 19 | import json 20 | from tqdm import tqdm 21 | from image_utils import psnr 22 | from argparse import ArgumentParser 23 | 24 | def evaluate_render(renders_dir, gt_dir, fname, lpips_inferer: Lpips): 25 | render = Image.open(renders_dir / fname) 26 | gt = Image.open(gt_dir / fname) 27 | render = tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda() 28 | gt = tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda() 29 | 30 | # [1, chaneel, H, W] 31 | return ( 32 | fname, 33 | ssim(render, gt), 34 | psnr(render, gt), 35 | lpips_inferer.forward(render, gt), 36 | ) 37 | 38 | 39 | def evaluate(args): 40 | 41 | full_dict = {} 42 | per_view_dict = {} 43 | full_dict_polytopeonly = {} 44 | per_view_dict_polytopeonly = {} 45 | 46 | full_dict[args.renders_color_dir] = {} 47 | per_view_dict[args.renders_color_dir] = {} 48 | full_dict_polytopeonly[args.renders_color_dir] = {} 49 | per_view_dict_polytopeonly[args.renders_color_dir] = {} 50 | 51 | # color_dir = Path(scene_dir) / "color" 52 | scene_dir_path = Path(args.renders_color_dir) 53 | 54 | gt_color_dir = Path(args.gt_color_dir) 55 | renders_color_dir = Path(args.renders_color_dir) 56 | 57 | image_names = [] 58 | ssims = [] 59 | psnrs = [] 60 | lpipss = [] 61 | 62 | lpips_inferer = Lpips(net_type="vgg") 63 | 64 | for fname in tqdm(os.listdir(renders_color_dir)): 65 | fname, ssim, psnr, lpips = evaluate_render( 66 | renders_color_dir, gt_color_dir, fname, lpips_inferer 67 | ) 68 | 69 | image_names.append(fname) 70 | ssims.append(ssim) 71 | psnrs.append(psnr) 72 | lpipss.append(lpips) 73 | 74 | print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 75 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 76 | print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 77 | # print max PSNR and min PSNR and their filename 78 | max_psnr = torch.tensor(psnrs).max() 79 | min_psnr = torch.tensor(psnrs).min() 80 | max_psnr_index = torch.tensor(psnrs).argmax() 81 | min_psnr_index = torch.tensor(psnrs).argmin() 82 | 83 | print(" Max PSNR: {:>12.7f} for {}".format(max_psnr, image_names[max_psnr_index])) 84 | print(" Min PSNR: {:>12.7f} for {}".format(min_psnr, image_names[min_psnr_index])) 85 | print("") 86 | 87 | full_dict[args.renders_color_dir].update( 88 | { 89 | "SSIM": torch.tensor(ssims).mean().item(), 90 | "PSNR": torch.tensor(psnrs).mean().item(), 91 | "LPIPS": torch.tensor(lpipss).mean().item(), 92 | "Max PSNR {}:".format(image_names[max_psnr_index]): max_psnr.item(), 93 | "Min PSNR {}:".format(image_names[min_psnr_index]): min_psnr.item(), 94 | } 95 | ) 96 | per_view_dict[args.renders_color_dir].update( 97 | { 98 | "SSIM": { 99 | name: ssim 100 | for ssim, name in zip(torch.tensor(ssims).tolist(), image_names) 101 | }, 102 | "PSNR": { 103 | name: psnr 104 | for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names) 105 | }, 106 | "LPIPS": { 107 | name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names) 108 | }, 109 | } 110 | ) 111 | 112 | with open(scene_dir_path / "../render_eval.json", "w") as fp: 113 | json.dump(full_dict, fp, indent=True) 114 | with open(scene_dir_path / "../render_eval_per_view.json", "w") as fp: 115 | json.dump(per_view_dict[args.renders_color_dir], fp, indent=True) 116 | with open(scene_dir_path / "../../evaluation_results.json", "a") as fp: 117 | full_dict[args.renders_color_dir] = { 118 | k: round(v, 3) for k, v in full_dict[args.renders_color_dir].items() 119 | } 120 | json.dump(full_dict, fp, indent=True) 121 | fp.write("\n") 122 | 123 | with open(scene_dir_path / "../../../all_evaluation_results.json", "a") as fp: 124 | # add dataset name 125 | json.dump(full_dict, fp, indent=True) 126 | fp.write("\n") 127 | 128 | 129 | if __name__ == "__main__": 130 | device = torch.device("cuda:0") 131 | torch.cuda.set_device(device) 132 | 133 | # Set up command line argument parser 134 | parser = ArgumentParser(description="Training script parameters") 135 | parser.add_argument("--gt_color_dir", type=str, default="") 136 | parser.add_argument("--renders_color_dir", type=str, default="") 137 | args = parser.parse_args() 138 | evaluate(args) 139 | -------------------------------------------------------------------------------- /eval/image_metrics/metrics_single.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from pathlib import Path 13 | import os 14 | from PIL import Image 15 | import torch 16 | import torchvision.transforms.functional as tf 17 | from loss_utils import ssim 18 | from lpipsPyTorch import Lpips 19 | import json 20 | from tqdm import tqdm 21 | from image_utils import psnr 22 | from argparse import ArgumentParser 23 | 24 | 25 | def evaluate_render(renders, gt): 26 | render = Image.open(renders) 27 | gt = Image.open(gt) 28 | render = tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda() 29 | gt = tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda() 30 | 31 | # [1, chaneel, H, W] 32 | return ( 33 | render, 34 | psnr(render, gt), 35 | ) 36 | 37 | 38 | def evaluate(args): 39 | 40 | full_dict = {} 41 | per_view_dict = {} 42 | full_dict_polytopeonly = {} 43 | per_view_dict_polytopeonly = {} 44 | 45 | full_dict[args.renders_color] = {} 46 | per_view_dict[args.renders_color] = {} 47 | full_dict_polytopeonly[args.renders_color] = {} 48 | per_view_dict_polytopeonly[args.renders_color] = {} 49 | 50 | scene_dir_path = Path(args.renders_color) 51 | 52 | gt_color_dir = Path(args.gt_color) 53 | renders_color = Path(args.renders_color) 54 | 55 | image_names = [] 56 | psnrs = [] 57 | 58 | fname, psnr = evaluate_render(renders_color, gt_color_dir) 59 | psnrs.append(psnr) 60 | 61 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 62 | 63 | full_dict[args.renders_color].update( 64 | { 65 | "PSNR": torch.tensor(psnrs).mean().item(), 66 | } 67 | ) 68 | per_view_dict[args.renders_color].update( 69 | { 70 | "PSNR": { 71 | name: psnr 72 | for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names) 73 | }, 74 | } 75 | ) 76 | 77 | with open(scene_dir_path.parent / "../render_eval.json", "w") as fp: 78 | json.dump(full_dict, fp, indent=True) 79 | with open(scene_dir_path.parent / "../../evaluation_results.json", "a") as fp: 80 | full_dict[args.renders_color] = { 81 | k: round(v, 3) for k, v in full_dict[args.renders_color].items() 82 | } 83 | json.dump(full_dict, fp, indent=True) 84 | fp.write("\n") 85 | 86 | 87 | if __name__ == "__main__": 88 | device = torch.device("cuda:0") 89 | torch.cuda.set_device(device) 90 | 91 | # Set up command line argument parser 92 | parser = ArgumentParser(description="Training script parameters") 93 | parser.add_argument("--gt_color", type=str, default="") 94 | parser.add_argument("--renders_color", type=str, default="") 95 | args = parser.parse_args() 96 | evaluate(args) 97 | -------------------------------------------------------------------------------- /eval/structure_metrics/cull_mesh.py: -------------------------------------------------------------------------------- 1 | # This file is a part of ESLAM. 2 | # 3 | # ESLAM is a NeRF-based SLAM system. It utilizes Neural Radiance Fields (NeRF) 4 | # to perform Simultaneous Localization and Mapping (SLAM) in real-time. 5 | # This software is the implementation of the paper "ESLAM: Efficient Dense SLAM 6 | # System Based on Hybrid Representation of Signed Distance Fields" by 7 | # Mohammad Mahdi Johari, Camilla Carta, and Francois Fleuret. 8 | # 9 | # Copyright 2023 ams-OSRAM AG 10 | # 11 | # Author: Mohammad Mahdi Johari 12 | # 13 | # Licensed under the Apache License, Version 2.0 (the "License"); 14 | # you may not use this file except in compliance with the License. 15 | # You may obtain a copy of the License at 16 | # 17 | # http://www.apache.org/licenses/LICENSE-2.0 18 | # 19 | # Unless required by applicable law or agreed to in writing, software 20 | # distributed under the License is distributed on an "AS IS" BASIS, 21 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 22 | # See the License for the specific language governing permissions and 23 | # limitations under the License. 24 | # 25 | # This file is a modified version of https://github.com/cvg/nice-slam/blob/master/src/tools/cull_mesh.py 26 | # which is covered by the following copyright and permission notice: 27 | # 28 | # Copyright 2022 Zihan Zhu, Songyou Peng, Viktor Larsson, Weiwei Xu, Hujun Bao, Zhaopeng Cui, Martin R. Oswald, Marc Pollefeys 29 | # 30 | # Licensed under the Apache License, Version 2.0 (the "License"); 31 | # you may not use this file except in compliance with the License. 32 | # You may obtain a copy of the License at 33 | # 34 | # http://www.apache.org/licenses/LICENSE-2.0 35 | # 36 | # Unless required by applicable law or agreed to in writing, software 37 | # distributed under the License is distributed on an "AS IS" BASIS, 38 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 39 | # See the License for the specific language governing permissions and 40 | # limitations under the License. 41 | 42 | import argparse 43 | 44 | import numpy as np 45 | import torch 46 | import torch.nn.functional as F 47 | import trimesh 48 | from tqdm import tqdm 49 | 50 | import sys 51 | 52 | sys.path.append(".") 53 | from src.utils.datasets import get_dataset 54 | from src import config 55 | 56 | 57 | def cull_mesh(mesh_file, cfg, args, device, estimate_c2w_list=None): 58 | """ 59 | Cull the mesh by removing the points that are not visible in any of the frames. 60 | The output mesh file will be saved in the same directory as the input mesh file. 61 | Args: 62 | mesh_file (str): path to the mesh file 63 | cfg (dict): configuration 64 | args (argparse.Namespace): arguments 65 | device (torch.device): device 66 | estimate_c2w_list (list): list of estimated camera poses, if None, it uses the ground truth camera poses 67 | Returns: 68 | None 69 | 70 | """ 71 | frame_reader = get_dataset(cfg, args, 1, device=device) 72 | 73 | eval_rec = cfg["meshing"]["eval_rec"] 74 | truncation = cfg["model"]["truncation"] 75 | H, W, fx, fy, cx, cy = ( 76 | cfg["cam"]["H"], 77 | cfg["cam"]["W"], 78 | cfg["cam"]["fx"], 79 | cfg["cam"]["fy"], 80 | cfg["cam"]["cx"], 81 | cfg["cam"]["cy"], 82 | ) 83 | 84 | if estimate_c2w_list is not None: 85 | n_imgs = len(estimate_c2w_list) 86 | else: 87 | n_imgs = len(frame_reader) 88 | 89 | mesh = trimesh.load(mesh_file, process=False) 90 | pc = mesh.vertices 91 | 92 | whole_mask = np.ones(pc.shape[0]).astype("bool") 93 | for i in tqdm(range(0, n_imgs, 1)): 94 | _, _, depth, c2w = frame_reader[i] 95 | depth, c2w = depth.to(device), c2w.to(device) 96 | 97 | if not estimate_c2w_list is None: 98 | c2w = estimate_c2w_list[i].to(device) 99 | 100 | points = pc.copy() 101 | points = torch.from_numpy(points).to(device) 102 | 103 | w2c = torch.inverse(c2w) 104 | K = torch.from_numpy( 105 | np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]]).reshape(3, 3) 106 | ).to(device) 107 | ones = torch.ones_like(points[:, 0]).reshape(-1, 1).to(device) 108 | homo_points = ( 109 | torch.cat([points, ones], dim=1).reshape(-1, 4, 1).to(device).float() 110 | ) 111 | cam_cord_homo = w2c @ homo_points 112 | cam_cord = cam_cord_homo[:, :3] 113 | 114 | cam_cord[:, 0] *= -1 115 | uv = K.float() @ cam_cord.float() 116 | z = uv[:, -1:] + 1e-5 117 | uv = uv[:, :2] / z 118 | uv = uv.squeeze(-1) 119 | 120 | grid = uv[None, None].clone() 121 | grid[..., 0] = grid[..., 0] / W 122 | grid[..., 1] = grid[..., 1] / H 123 | grid = 2 * grid - 1 124 | depth_samples = F.grid_sample( 125 | depth[None, None], grid, padding_mode="zeros", align_corners=True 126 | ).squeeze() 127 | 128 | edge = 0 129 | if eval_rec: 130 | mask = ( 131 | (depth_samples + truncation >= -z[:, 0, 0]) 132 | & (0 <= -z[:, 0, 0]) 133 | & (uv[:, 0] < W - edge) 134 | & (uv[:, 0] > edge) 135 | & (uv[:, 1] < H - edge) 136 | & (uv[:, 1] > edge) 137 | ) 138 | else: 139 | mask = ( 140 | (0 <= -z[:, 0, 0]) 141 | & (uv[:, 0] < W - edge) 142 | & (uv[:, 0] > edge) 143 | & (uv[:, 1] < H - edge) 144 | & (uv[:, 1] > edge) 145 | ) 146 | 147 | mask = mask.cpu().numpy() 148 | 149 | whole_mask &= ~mask 150 | 151 | face_mask = whole_mask[mesh.faces].all(axis=1) 152 | mesh.update_faces(~face_mask) 153 | mesh.remove_unreferenced_vertices() 154 | mesh.process(validate=False) 155 | 156 | mesh_ext = mesh_file.split(".")[-1] 157 | output_file = mesh_file[: -len(mesh_ext) - 1] + "_culled." + mesh_ext 158 | 159 | mesh.export(output_file) 160 | 161 | 162 | ## It is also possible to use the cull_mesh function in the following way, where the ground truth camera poses are used. 163 | if __name__ == "__main__": 164 | parser = argparse.ArgumentParser(description="Arguments to cull the mesh.") 165 | 166 | parser.add_argument("config", type=str, help="path to the config file") 167 | parser.add_argument("--input_mesh", type=str, help="path to the mesh to be culled") 168 | 169 | args = parser.parse_args() 170 | args.input_folder = None 171 | 172 | cfg = config.load_config(args.config, "configs/ESLAM.yaml") 173 | 174 | cull_mesh(args.input_mesh, cfg, args, "cuda") 175 | -------------------------------------------------------------------------------- /eval/structure_metrics/eval_utils.py: -------------------------------------------------------------------------------- 1 | # This file is derived from [Atlas](https://github.com/magicleap/Atlas). 2 | # Originating Author: Zak Murez (zak.murez.com) 3 | # Modified for [SHINEMapping] by Yue Pan. 4 | 5 | # Original header: 6 | # Copyright 2020 Magic Leap, Inc. 7 | 8 | # Licensed under the Apache License, Version 2.0 (the "License"); 9 | # you may not use this file except in compliance with the License. 10 | # You may obtain a copy of the License at 11 | 12 | # http://www.apache.org/licenses/LICENSE-2.0 13 | 14 | # Unless required by applicable law or agreed to in writing, software 15 | # distributed under the License is distributed on an "AS IS" BASIS, 16 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 17 | # See the License for the specific language governing permissions and 18 | # limitations under the License. 19 | 20 | 21 | import open3d as o3d 22 | import numpy as np 23 | 24 | 25 | def eval_mesh( 26 | file_pred, 27 | file_trgt, 28 | down_sample_res=0.02, 29 | threshold=0.05, 30 | truncation_acc=0.50, 31 | truncation_com=0.50, 32 | gt_bbx_mask_on=True, 33 | mesh_sample_point=10000000, 34 | possion_sample_init_factor=5, 35 | ): 36 | """Compute Mesh metrics between prediction and target. 37 | Opens the Meshs and runs the metrics 38 | Args: 39 | file_pred: file path of prediction (should be mesh) 40 | file_trgt: file path of target (shoud be point cloud) 41 | down_sample_res: use voxel_downsample to uniformly sample mesh points 42 | threshold: distance threshold used to compute precision/recall 43 | truncation_acc: points whose nearest neighbor is farther than the distance would not be taken into account (take pred as reference) 44 | truncation_com: points whose nearest neighbor is farther than the distance would not be taken into account (take trgt as reference) 45 | gt_bbx_mask_on: use the bounding box of the trgt as a mask of the pred mesh 46 | mesh_sample_point: number of the sampling points from the mesh 47 | possion_sample_init_factor: used for possion uniform sampling, check open3d for more details (deprecated) 48 | Returns: 49 | 50 | Returns: 51 | Dict of mesh metrics (chamfer distance, precision, recall, f1 score, etc.) 52 | """ 53 | o3d.utility.random.seed(0) 54 | 55 | mesh_pred = o3d.io.read_triangle_mesh(file_pred) 56 | 57 | pcd_trgt = o3d.io.read_point_cloud(file_trgt) 58 | 59 | # (optional) filter the prediction outside the gt bounding box (since gt sometimes is not complete enough) 60 | if gt_bbx_mask_on: 61 | # trgt_bbx = pcd_trgt.get_axis_aligned_bounding_box() 62 | trgt_bbx = pcd_trgt.get_minimal_oriented_bounding_box() 63 | mesh_pred = mesh_pred.crop(trgt_bbx) 64 | 65 | # min_bound = trgt_bbx.get_min_bound() 66 | # min_bound[2] -= down_sample_res 67 | # max_bound = trgt_bbx.get_max_bound() 68 | # max_bound[2] += down_sample_res 69 | # trgt_bbx = o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound) 70 | # mesh_pred = mesh_pred.crop(trgt_bbx) 71 | # # pcd_sample_pred = pcd_sample_pred.crop(trgt_bbx) 72 | 73 | # pcd_sample_pred = mesh_pred.sample_points_poisson_disk(number_of_points=mesh_sample_point, init_factor=possion_sample_init_factor) 74 | # mesh uniform sampling 75 | pcd_sample_pred = mesh_pred.sample_points_uniformly( 76 | number_of_points=mesh_sample_point 77 | ) 78 | 79 | if down_sample_res > 0: 80 | pred_pt_count_before = len(pcd_sample_pred.points) 81 | pcd_pred = pcd_sample_pred.voxel_down_sample(down_sample_res) 82 | pcd_trgt = pcd_trgt.voxel_down_sample(down_sample_res) 83 | pred_pt_count_after = len(pcd_pred.points) 84 | print( 85 | "Predicted mesh unifrom sample: ", 86 | pred_pt_count_before, 87 | " --> ", 88 | pred_pt_count_after, 89 | " (", 90 | down_sample_res, 91 | "m)", 92 | ) 93 | 94 | verts_pred = np.asarray(pcd_pred.points) 95 | verts_trgt = np.asarray(pcd_trgt.points) 96 | 97 | _, dist_p = nn_correspondance( 98 | verts_trgt, verts_pred, truncation_acc, True 99 | ) # find nn in ground truth samples for each predict sample -> precision related 100 | _, dist_r = nn_correspondance( 101 | verts_pred, verts_trgt, truncation_com, False 102 | ) # find nn in predict samples for each ground truth sample -> recall related 103 | 104 | dist_p = np.array(dist_p) 105 | dist_r = np.array(dist_r) 106 | 107 | dist_p_s = np.square(dist_p) 108 | dist_r_s = np.square(dist_r) 109 | 110 | dist_p_mean = np.mean(dist_p) 111 | dist_r_mean = np.mean(dist_r) 112 | 113 | dist_p_s_mean = np.mean(dist_p_s) 114 | dist_r_s_mean = np.mean(dist_r_s) 115 | 116 | chamfer_l1 = 0.5 * (dist_p_mean + dist_r_mean) 117 | chamfer_l2 = np.sqrt(0.5 * (dist_p_s_mean + dist_r_s_mean)) 118 | 119 | precision = np.mean((dist_p < threshold).astype("float")) * 100.0 # % 120 | recall = np.mean((dist_r < threshold).astype("float")) * 100.0 # % 121 | fscore = 2 * precision * recall / (precision + recall) # % 122 | 123 | metrics = { 124 | "MAE_accuracy (cm)": dist_p_mean * 100.0, 125 | "MAE_completeness (cm)": dist_r_mean * 100.0, 126 | "Chamfer_L1 (cm)": chamfer_l1 * 100.0, 127 | "Chamfer_L2 (cm)": chamfer_l2 * 100.0, 128 | "Precision [Accuracy] (%)": precision, 129 | "Recall [Completeness] (%)": recall, 130 | "F-score (%)": fscore, 131 | "Spacing (cm)": down_sample_res * 100.0, # evlaution setup 132 | "Inlier_threshold (cm)": threshold * 100.0, # evlaution setup 133 | "Outlier_truncation_acc (cm)": truncation_acc * 100.0, # evlaution setup 134 | "Outlier_truncation_com (cm)": truncation_com * 100.0, # evlaution setup 135 | } 136 | return metrics 137 | 138 | 139 | def nn_correspondance(verts1, verts2, truncation_dist, ignore_outlier=True): 140 | """for each vertex in verts2 find the nearest vertex in verts1 141 | Args: 142 | nx3 np.array's 143 | scalar truncation_dist: points whose nearest neighbor is farther than the distance would not be taken into account 144 | Returns: 145 | ([indices], [distances]) 146 | """ 147 | 148 | indices = [] 149 | distances = [] 150 | if len(verts1) == 0 or len(verts2) == 0: 151 | return indices, distances 152 | 153 | pcd = o3d.geometry.PointCloud() 154 | pcd.points = o3d.utility.Vector3dVector(verts1) 155 | kdtree = o3d.geometry.KDTreeFlann(pcd) 156 | 157 | truncation_dist_square = truncation_dist**2 158 | 159 | for vert in verts2: 160 | _, inds, dist_square = kdtree.search_knn_vector_3d(vert, 1) 161 | 162 | if dist_square[0] < truncation_dist_square: 163 | indices.append(inds[0]) 164 | distances.append(np.sqrt(dist_square[0])) 165 | else: 166 | if not ignore_outlier: 167 | indices.append(inds[0]) 168 | distances.append(truncation_dist) 169 | 170 | return indices, distances 171 | 172 | 173 | def eval_depth(depth_pred, depth_trgt): 174 | """Computes 2d metrics between two depth maps 175 | Args: 176 | depth_pred: mxn np.array containing prediction 177 | depth_trgt: mxn np.array containing ground truth 178 | Returns: 179 | Dict of metrics 180 | """ 181 | mask1 = depth_pred > 0 # ignore values where prediction is 0 (% complete) 182 | mask = (depth_trgt < 10) * (depth_trgt > 0) * mask1 183 | 184 | depth_pred = depth_pred[mask] 185 | depth_trgt = depth_trgt[mask] 186 | abs_diff = np.abs(depth_pred - depth_trgt) 187 | abs_rel = abs_diff / depth_trgt 188 | sq_diff = abs_diff**2 189 | sq_rel = sq_diff / depth_trgt 190 | sq_log_diff = (np.log(depth_pred) - np.log(depth_trgt)) ** 2 191 | thresh = np.maximum((depth_trgt / depth_pred), (depth_pred / depth_trgt)) 192 | r1 = (thresh < 1.25).astype("float") 193 | r2 = (thresh < 1.25**2).astype("float") 194 | r3 = (thresh < 1.25**3).astype("float") 195 | 196 | metrics = {} 197 | metrics["AbsRel"] = np.mean(abs_rel) 198 | metrics["AbsDiff"] = np.mean(abs_diff) 199 | metrics["SqRel"] = np.mean(sq_rel) 200 | metrics["RMSE"] = np.sqrt(np.mean(sq_diff)) 201 | metrics["LogRMSE"] = np.sqrt(np.mean(sq_log_diff)) 202 | metrics["r1"] = np.mean(r1) 203 | metrics["r2"] = np.mean(r2) 204 | metrics["r3"] = np.mean(r3) 205 | metrics["complete"] = np.mean(mask1.astype("float")) 206 | 207 | return metrics 208 | 209 | 210 | def crop_intersection( 211 | file_gt, files_pred, out_file_crop, dist_thre=0.1, mesh_sample_point=1000000 212 | ): 213 | """Get the cropped ground truth point cloud according to the intersection of the predicted 214 | mesh by different methods 215 | Args: 216 | file_gt: file path of the ground truth (shoud be point cloud) 217 | files_pred: a list of the paths of different methods's reconstruction (shoud be mesh) 218 | out_file_crop: output path of the cropped ground truth point cloud 219 | dist_thre: nearest neighbor distance threshold in meter 220 | mesh_sample_point: number of the sampling points from the mesh 221 | """ 222 | print("Load the original ground truth point cloud from:", file_gt) 223 | pcd_gt = o3d.io.read_point_cloud(file_gt) 224 | pcd_gt_pts = np.asarray(pcd_gt.points) 225 | dist_square_thre = dist_thre**2 226 | for i in range(len(files_pred)): 227 | cur_file_pred = files_pred[i] 228 | print("Process", cur_file_pred) 229 | cur_mesh_pred = o3d.io.read_triangle_mesh(cur_file_pred) 230 | 231 | cur_sample_pred = cur_mesh_pred.sample_points_uniformly( 232 | number_of_points=mesh_sample_point 233 | ) 234 | 235 | cur_kdtree = o3d.geometry.KDTreeFlann(cur_sample_pred) 236 | 237 | crop_pcd_gt_pts = [] 238 | for pt in pcd_gt_pts: 239 | _, _, dist_square = cur_kdtree.search_knn_vector_3d(pt, 1) 240 | 241 | if dist_square[0] < dist_square_thre: 242 | crop_pcd_gt_pts.append(pt) 243 | 244 | pcd_gt_pts = np.asarray(crop_pcd_gt_pts) 245 | 246 | crop_pcd_gt = o3d.geometry.PointCloud() 247 | crop_pcd_gt.points = o3d.utility.Vector3dVector(pcd_gt_pts) 248 | 249 | print("Output the croped ground truth to:", out_file_crop) 250 | o3d.io.write_point_cloud(out_file_crop, crop_pcd_gt) 251 | -------------------------------------------------------------------------------- /eval/structure_metrics/evaluator.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | from eval_utils import eval_mesh 4 | from pathlib import Path 5 | 6 | if __name__ == "__main__": 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument("--gt_pcd", help="folder containing the depth images") 9 | parser.add_argument("--pred_mesh", help="folder containing the rgb images") 10 | # evaluation parameters 11 | parser.add_argument("--down_sample_vox", type=float, default=0.01) 12 | parser.add_argument("--dist_thre", type=float, default=0.02) 13 | parser.add_argument("--truncation_dist_acc", type=float, default=0.2) 14 | parser.add_argument("--truncation_dist_com", type=float, default=0.2) 15 | args = parser.parse_args() 16 | 17 | output_json_path = Path(args.pred_mesh).parent 18 | 19 | # evaluation 20 | eval_metric = eval_mesh( 21 | args.pred_mesh, 22 | args.gt_pcd, 23 | down_sample_res=args.down_sample_vox, 24 | threshold=args.dist_thre, 25 | truncation_acc=args.truncation_dist_acc, 26 | truncation_com=args.truncation_dist_com, 27 | gt_bbx_mask_on=True, 28 | ) 29 | 30 | try: 31 | with open(output_json_path / "structure_eval.json", "w") as fp: 32 | json.dump(args.pred_mesh, fp, indent=True) 33 | json.dump(eval_metric, fp, indent=True) 34 | for key in eval_metric: 35 | print(key + ": " + str(eval_metric[key])) 36 | print(f"Structure evaluation results are written into {output_json_path}") 37 | 38 | with open(output_json_path / "evaluation_results.json", "a") as fp: 39 | json.dump(args.pred_mesh, fp, indent=True) 40 | eval_metric = {k: round(v, 3) for k, v in eval_metric.items()} 41 | json.dump(eval_metric, fp, indent=True) 42 | fp.write("\n") 43 | 44 | with open(output_json_path / "../all_evaluation_results.json", "a") as fp: 45 | # add dataset name 46 | json.dump(args.pred_mesh, fp, indent=True) 47 | eval_metric = {k: round(v, 3) for k, v in eval_metric.items()} 48 | json.dump(eval_metric, fp, indent=True) 49 | fp.write("\n") 50 | except IOError: 51 | print("I/O error") 52 | -------------------------------------------------------------------------------- /include/data_loader/data_loader.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #include "data_parsers/base_parser.h" 6 | #include "utils/sensor_utils/sensors.hpp" 7 | 8 | namespace dataloader { 9 | class DataLoader { 10 | public: 11 | typedef std::shared_ptr Ptr; 12 | explicit DataLoader(const std::string &dataset_path, 13 | const int &_dataset_type = 0, 14 | const torch::Device &_device = torch::kCPU, 15 | const bool &_preload = false, 16 | const float &_res_scale = 1.0, 17 | const sensor::Sensors &_sensor = sensor::Sensors()); 18 | 19 | // Base paths. 20 | // std::string calib_path_, poses_path_; 21 | torch::Device device_ = torch::kCPU; 22 | 23 | dataparser::DataParser::Ptr dataparser_ptr_; 24 | 25 | bool get_next_data(int idx, torch::Tensor &_pose, PointCloudT &_points); 26 | 27 | bool get_next_data(int idx, torch::Tensor &_pose, DepthSamples &_depth_rays, 28 | ColorSamples &_color_rays, 29 | const torch::Device &_device = torch::Device(torch::kCPU)); 30 | 31 | bool get_next_data(int idx, torch::Tensor &_pose, DepthSamples &_depth_rays, 32 | const torch::Device &_device = torch::Device(torch::kCPU)); 33 | bool get_next_data(int idx, torch::Tensor &_pose, ColorSamples &_color_rays, 34 | const torch::Device &_device = torch::Device(torch::kCPU)); 35 | 36 | torch::Tensor get_pose(int idx, const int &pose_type = 0); 37 | 38 | int export_image(std::filesystem::path &gs_sparse_path, 39 | std::filesystem::path &gs_path, int data_type, 40 | std::string filename, bool bin, uint32_t camera_id, 41 | bool eval, bool llff, std::string prefix = "", 42 | int prefix_num = 0); 43 | void export_as_colmap_format(bool bin = true, bool llff = false); 44 | void export_as_colmap_format_for_nerfstudio(bool bin = true); 45 | 46 | private: 47 | // torch::Tensor Tr; 48 | 49 | // std::vector poses_vec_; 50 | // std::vector depth_name_vec_; 51 | 52 | int dataset_type_; 53 | 54 | // bool load_calib(); 55 | // bool load_poses(); 56 | 57 | // bool load_depth_list(); 58 | }; 59 | } // namespace dataloader -------------------------------------------------------------------------------- /include/data_loader/data_parsers/base_parser.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "utils/sensor_utils/cameras.hpp" 4 | #include "utils/sensor_utils/sensors.hpp" 5 | #include "utils/utils.h" 6 | #include 7 | #include 8 | 9 | #include 10 | #include 11 | 12 | namespace dataparser { 13 | 14 | enum DataType { 15 | RawColor = 0, 16 | RawDepth = 1, 17 | TrainColor = 2, 18 | TrainDepth = 3, 19 | EvalColor = 4, 20 | EvalDepth = 5, 21 | TestColor = 6, 22 | TestDepth = 7 23 | }; 24 | 25 | enum DepthType { Image = 0, PLY = 1, BIN = 2, PCD = 3 }; 26 | 27 | std::vector 28 | read_filelists(const std::filesystem::path &directory, 29 | const std::string &prefix, 30 | const std::string &extension = ".ply"); 31 | 32 | void sort_filelists(std::vector &filists); 33 | 34 | void load_file_list(const std::string &dir_path, 35 | std::vector &out_filelsits, 36 | const std::string &prefix, 37 | const std::string &file_extension = ".ply"); 38 | 39 | struct DataParser { 40 | typedef std::shared_ptr Ptr; 41 | explicit DataParser(const std::string &_dataset_path, 42 | const torch::Device &_device = torch::kCPU, 43 | const bool &_preload = true, 44 | const float &_res_scale = 1.0, 45 | const int &_dataset_system_type = 0, 46 | const sensor::Sensors &_sensor = sensor::Sensors()) 47 | : dataset_path_(_dataset_path), device_(_device), preload_(_preload), 48 | res_scale_(_res_scale), dataset_system_type_(_dataset_system_type), 49 | sensor_(_sensor) { 50 | // you are supposed to call load_data() in the derived class 51 | }; 52 | 53 | std::filesystem::path dataset_path_, dataset_name_; 54 | std::filesystem::path pose_path_, calib_path_, color_path_, depth_path_; 55 | 56 | std::filesystem::path eval_pose_path_, eval_color_path_, eval_depth_path_; 57 | 58 | torch::Device device_ = torch::kCPU; 59 | torch::Tensor train_color_; // [N, H, W, 3] 60 | DepthSamples train_depth_pack_; // [N] 61 | torch::Tensor depth_poses_, train_depth_poses_, color_poses_, 62 | train_color_poses_; // [N, 3, 4] 63 | torch::Tensor time_stamps_; // [N] 64 | torch::Tensor eval_color_poses_, eval_depth_poses_; // [N, 4, 4] 65 | int dataset_system_type_; 66 | sensor::Sensors sensor_; 67 | bool preload_; 68 | float res_scale_; 69 | int depth_type_ = DepthType::Image; // 0: image; 1: ply; 2: bin; 3: pcd 70 | torch::Tensor 71 | T_B_S_; // [4, 4]; extrinsic param, transformation from sensor to body 72 | torch::Tensor T_S_B_; 73 | 74 | std::vector raw_color_filelists_, raw_depth_filelists_, 75 | train_color_filelists_, train_depth_filelists_, eval_color_filelists_, 76 | eval_depth_filelists_; 77 | 78 | float color_scale_inv_ = 1.0 / 255.0; 79 | float depth_scale_inv_ = 1e-3; 80 | 81 | virtual torch::Tensor get_pose(const int &idx, const int &pose_type) const; 82 | 83 | virtual torch::Tensor get_pose(const torch::Tensor &idx, 84 | const int &pose_type) const; 85 | 86 | virtual torch::Tensor get_image(const int &idx, const int &image_type) const; 87 | virtual torch::Tensor get_image(const torch::Tensor &idx, 88 | const int &image_type) const; 89 | 90 | size_t size(const int &data_type) const { 91 | switch (data_type) { 92 | case DataType::RawColor: { 93 | return raw_color_filelists_.size(); 94 | } 95 | case DataType::RawDepth: { 96 | return raw_depth_filelists_.size(); 97 | } 98 | case DataType::TrainColor: { 99 | return train_color_filelists_.size(); 100 | } 101 | case DataType::TrainDepth: { 102 | return train_depth_filelists_.size(); 103 | ; 104 | } 105 | case DataType::EvalColor: { 106 | return eval_color_filelists_.size(); 107 | } 108 | case DataType::EvalDepth: { 109 | return eval_depth_filelists_.size(); 110 | } 111 | case DataType::TestColor: 112 | case DataType::TestDepth: 113 | default: 114 | throw std::runtime_error("Invalid image type"); 115 | } 116 | } 117 | 118 | std::filesystem::path get_file(const int &idx, const int &image_type = 0); 119 | cv::Mat get_image_cv_mat(const int &idx, const int &image_type = 0) const; 120 | torch::Tensor get_color_image(const int &idx, const int &image_type = 0, 121 | const float &scale = 1.0) const; 122 | torch::Tensor get_depth_image(const int &idx) const; 123 | 124 | virtual std::vector get_depth_zdir(const int &idx); 125 | virtual std::vector get_distance_ndir_zdirn(const int &idx); 126 | virtual at::Tensor get_points(const int &idx); 127 | virtual std::vector get_points_dist_ndir_zdirn(const int &idx); 128 | 129 | virtual void load_data() { throw std::runtime_error("Not implemented"); } 130 | 131 | virtual void load_eval_data() { throw std::runtime_error("Not implemented"); } 132 | 133 | virtual bool load_calib() { 134 | throw std::runtime_error("Not implemented"); 135 | 136 | if (!std::filesystem::exists(calib_path_)) { 137 | std::cerr << "Calibration file does not exist\n"; 138 | T_B_S_ = torch::eye(4); 139 | } 140 | /* else { 141 | std::ifstream import_file(calib_path_, std::ios::in); 142 | if (!import_file) { 143 | std::cerr << "Could not open calibration file: " << calib_path_ << 144 | "\n"; return false; 145 | } 146 | 147 | std::string line; 148 | while (std::getline(import_file, line)) { 149 | std::stringstream line_stream(line); 150 | 151 | // Check what the header is. Each line consists of two parts: 152 | // a header followed by a ':' followed by space-separated data. 153 | std::string header; 154 | std::getline(line_stream, header, ':'); 155 | std::string data; 156 | std::getline(line_stream, data, ':'); 157 | 158 | std::vector parsed_floats; 159 | if (header == "Tr") { 160 | // Parse the translation matrix. 161 | if (parseVectorOfFloats(data, &parsed_floats)) { 162 | Tr = 163 | torch::from_blob(parsed_floats.data(), {3, 4}, 164 | torch::kFloat32).clone(); Tr = torch::cat({Tr, torch::tensor({{0, 0, 0, 165 | 1}})}, 0); 166 | } 167 | } 168 | } 169 | } */ 170 | std::cout << "T_B_S:" << "\n" << T_B_S_ << "\n"; 171 | T_S_B_ = T_B_S_.inverse(); 172 | 173 | return true; 174 | } 175 | 176 | virtual std::vector load_poses(const std::string &pose_path, 177 | const bool &with_head, 178 | const int &pose_type = 0); 179 | 180 | virtual void load_intrinsics() { 181 | throw std::runtime_error("Not implemented"); 182 | } 183 | 184 | virtual void 185 | align_pose_sensor(std::vector &out_filelsits, 186 | torch::Tensor &sensor_poses, 187 | float max_time_diff_sensor_and_pose); 188 | 189 | virtual void load_colors(const std::string &file_extension, 190 | const std::string &prefix = "", 191 | const bool eval = false, const bool &llff = false); 192 | 193 | virtual void load_depths(const std::string &file_extension, 194 | const std::string &prefix = "", 195 | const bool eval = false, const bool &llff = false); 196 | 197 | virtual std::string get_gt_mesh_path() { return ""; } 198 | }; 199 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/fastlivo_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "data_loader/data_parsers/rosbag_parser.hpp" 3 | #include "utils/coordinates.h" 4 | #include "utils/sensor_utils/cameras.hpp" 5 | #include "utils/sensor_utils/sensors.hpp" 6 | #include 7 | 8 | namespace dataparser { 9 | struct Fastlivo : Rosbag { 10 | explicit Fastlivo(const std::filesystem::path &_bag_path, 11 | const torch::Device &_device = torch::kCPU, 12 | const bool &_preload = true, const float &_res_scale = 1.0, 13 | const sensor::Sensors &_sensor = sensor::Sensors()) 14 | : Rosbag(_bag_path, _device, _preload, _res_scale, 15 | coords::SystemType::OpenCV, _sensor) { 16 | depth_type_ = DepthType::PLY; 17 | 18 | dataset_name_ = bag_path_.filename(); 19 | dataset_name_ = dataset_name_.replace_extension(); 20 | 21 | pose_topic = "/aft_mapped_to_init"; 22 | // pose_topic = "/Odometry"; 23 | color_topic = "/origin_img"; 24 | depth_topic = "/cloud_registered_body"; 25 | 26 | load_intrinsics(); 27 | 28 | load_data(); 29 | if (std::filesystem::exists(pose_path_)) { 30 | } else { 31 | std::cout << "pose_path_ does not exist: " << pose_path_ << std::endl; 32 | } 33 | } 34 | 35 | void load_intrinsics() override { 36 | // auto scale = 0.5f; // output image from Fastlivo is 640x512 37 | auto scale = 1.0f; // output image from Fastlivo is 1280x1024 38 | sensor_.camera.width = scale * sensor_.camera.width; 39 | sensor_.camera.height = scale * sensor_.camera.height; 40 | // HKU 41 | sensor_.camera.fx = scale * sensor_.camera.fx; 42 | sensor_.camera.fy = scale * sensor_.camera.fy; 43 | sensor_.camera.cx = scale * sensor_.camera.cx; 44 | sensor_.camera.cy = scale * sensor_.camera.cy; 45 | } 46 | }; 47 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/kitti_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "base_parser.h" 3 | #include "utils/coordinates.h" 4 | #include "utils/sensor_utils/cameras.hpp" 5 | #include 6 | 7 | namespace dataparser { 8 | struct Kitti : DataParser { 9 | explicit Kitti(const std::filesystem::path &_dataset_lidar_path, 10 | const torch::Device &_device = torch::kCPU, 11 | const bool &_preload = true, const float &_res_scale = 1.0) 12 | : DataParser(_dataset_lidar_path, _device, _preload, _res_scale, 13 | coords::SystemType::OpenCV) { 14 | depth_type_ = DepthType::BIN; 15 | 16 | // KITTI/data_odometry_velodyne/dataset/sequences/00 17 | dataset_name_ = _dataset_lidar_path.filename(); 18 | // /media/chrisliu/T9/Datasets/KITTI/data_odometry_poses 19 | 20 | auto base_path = (dataset_path_ / "../../../..").lexically_normal(); 21 | calib_path_ = base_path / "data_odometry_calib/dataset/sequences" / 22 | dataset_name_ / "calib.txt"; 23 | pose_path_ = base_path / "data_odometry_poses/dataset/poses" / 24 | (dataset_name_.string() + ".txt"); 25 | color_path_ = base_path / "data_odometry_color/dataset/sequences" / 26 | dataset_name_ / "image_2"; 27 | depth_path_ = base_path / "data_odometry_velodyne/dataset/sequences" / 28 | dataset_name_ / "velodyne"; 29 | load_intrinsics(); 30 | } 31 | 32 | torch::Tensor T_C0_L, T_C0_C2; 33 | 34 | void load_data() override { 35 | 36 | auto T_C0_C0 = load_poses(pose_path_, false, 2)[0]; 37 | TORCH_CHECK(T_C0_C0.size(0) > 0); 38 | auto T_W_C0 = 39 | coords::change_world_system(T_C0_C0, coords::SystemType::Kitti); 40 | color_poses_ = T_W_C0.matmul(T_C0_C2); 41 | depth_poses_ = T_W_C0.matmul(T_C0_L); 42 | 43 | load_colors(".png", "", false, true); 44 | TORCH_CHECK(color_poses_.size(0) == raw_color_filelists_.size()); 45 | load_depths(".bin", "", false, true); 46 | TORCH_CHECK(depth_poses_.size(0) == raw_depth_filelists_.size()); 47 | } 48 | 49 | void load_intrinsics() override { 50 | // open calibration file 51 | std::ifstream calib_file(calib_path_); 52 | if (!calib_file) { 53 | throw std::runtime_error("Could not open calibration file"); 54 | } 55 | 56 | std::string line; 57 | T_C0_L = torch::eye(4); 58 | T_C0_C2 = torch::eye(4); 59 | while (std::getline(calib_file, line)) { 60 | std::istringstream iss(line); 61 | std::string token; 62 | iss >> token; 63 | if (token == "P2:") { 64 | int index = 0; 65 | torch::Tensor baseline = torch::zeros(3); 66 | while (iss >> token) { 67 | if (index == 0) { 68 | sensor_.camera.fx = std::stof(token); 69 | } else if (index == 2) { 70 | sensor_.camera.cx = std::stof(token); 71 | } else if (index == 3) { 72 | T_C0_C2[0][3] = -std::stof(token) / sensor_.camera.fx; 73 | } else if (index == 5) { 74 | sensor_.camera.fy = std::stof(token); 75 | } else if (index == 6) { 76 | sensor_.camera.cy = std::stof(token); 77 | } else if (index == 7) { 78 | T_C0_C2[1][3] = -std::stof(token) / sensor_.camera.fy; 79 | } else if (index == 11) { 80 | T_C0_C2[2][3] = -std::stof(token); 81 | } 82 | index++; 83 | } 84 | } else if (token == "Tr:") { 85 | int i = 0; 86 | int j = 0; 87 | float value; 88 | while (iss >> value) { 89 | T_C0_L[i][j] = value; 90 | j++; 91 | if ((j % 4) == 0) { 92 | i++; 93 | j = 0; 94 | } 95 | } 96 | } 97 | } 98 | sensor_.camera.width = 1241; 99 | sensor_.camera.height = 376; 100 | depth_scale_inv_ = 1.0; 101 | // print out cameras 102 | std::cout << "fx: " << sensor_.camera.fx << ", fy: " << sensor_.camera.fy 103 | << ", cx: " << sensor_.camera.cx << ", cy: " << sensor_.camera.cy 104 | << "\n"; 105 | std::cout << "T_C0_L:\n" << T_C0_L << "\n"; 106 | } 107 | 108 | std::vector get_distance_ndir_zdirn(const int &idx) override { 109 | /** 110 | * @description: 111 | * @return {distance, ndir, dir_norm}, where ndir.norm = 1; 112 | {[height width 1], [height width 3], [height width 1]} 113 | */ 114 | 115 | auto pointcloud = get_depth_image(idx); 116 | // [height width 1] 117 | auto distance = pointcloud.norm(2, -1, true); 118 | auto ndir = pointcloud / distance; 119 | return {distance, ndir, distance}; 120 | } 121 | }; 122 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/ncd_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "base_parser.h" 3 | #include "nerfacc_cpp/cameras.hpp" 4 | #include "utils/coordinates.h" 5 | 6 | namespace dataparser { 7 | struct NewerCollege : DataParser { 8 | explicit NewerCollege( 9 | const std::string &_dataset_path, 10 | const torch::Device &_device = torch::kCPU, const bool &_preload = true, 11 | const int &_dataset_system_type = coords::SystemType::OpenCV) 12 | : DataParser(_dataset_path, _device, _preload, _dataset_system_type) { 13 | 14 | // Deprecated: no color images 15 | 16 | std::cout << dataset_path_ << std::endl; 17 | // input _dataset_path should be end with .bag: 18 | // NewerCollegeDataset/2021-ouster-os0-128-alphasense/collection 1 - newer 19 | // college/2021-07-01-10-37-38-quad-easy.bag 20 | auto bag_name = dataset_path_.filename().string(); 21 | // extract quad-easy from 2021-07-01-10-37-38-quad-easy.bag 22 | bag_name = bag_name.substr(0, bag_name.find(".bag")); 23 | bag_name = bag_name.substr(bag_name.find_last_of("0123456789") + 2); 24 | std::cout << bag_name << std::endl; 25 | 26 | dataset_path_ = dataset_path_.parent_path(); 27 | std::cout << dataset_path_ << std::endl; 28 | 29 | pose_path_ = dataset_path_ / "ground_truth"; 30 | color_path_ = dataset_path_ / "results"; 31 | depth_path_ = color_path_; 32 | 33 | gt_mesh_path_ = dataset_path_.parent_path() / "cull_replica_mesh" / 34 | (dataset_path_.filename().string() + ".ply"); 35 | // gt_mesh_path_ = dataset_path_.parent_path() / 36 | // (dataset_path_.filename().string() + "_mesh.ply"); 37 | 38 | load_intrinsics(); 39 | } 40 | 41 | std::filesystem::path gt_mesh_path_; 42 | 43 | void load_data() override { 44 | depth_poses_ = load_poses(pose_path_, false, 1)[0]; 45 | TORCH_CHECK(depth_poses_.size(0) > 0); 46 | 47 | load_colors(".jpg", "frame", false, true); 48 | TORCH_CHECK(depth_poses_.size(0) == color_filelists_.size()); 49 | load_depths(".png", "depth", false, true); 50 | TORCH_CHECK(color_filelists_.size() == depth_filelists_.size()); 51 | } 52 | 53 | void load_intrinsics() override { 54 | // Replica/cam_params.json 55 | sensor_.camera.width = 1200; 56 | sensor_.camera.height = 680; 57 | sensor_.camera.fx = 600.0; 58 | sensor_.camera.fy = 600.0; 59 | sensor_.camera.cx = 599.5; 60 | sensor_.camera.cy = 339.5; 61 | depth_scale_inv_ = 1.0 / 6553.5; 62 | } 63 | 64 | std::string get_gt_mesh_path() override { return gt_mesh_path_.string(); } 65 | }; 66 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/neuralrgbd_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "base_parser.h" 3 | #include "utils/sensor_utils/cameras.hpp" 4 | #include "utils/coordinates.h" 5 | 6 | namespace dataparser { 7 | struct NeuralRGBD : DataParser { 8 | explicit NeuralRGBD(const std::string &_dataset_path, 9 | const torch::Device &_device = torch::kCPU, 10 | const bool &_preload = true, 11 | const float &_res_scale = 1.0, const int &_depth_type = 0, 12 | const int &_dataset_system_type = 2) 13 | : DataParser(_dataset_path, _device, _preload, _res_scale, 14 | _dataset_system_type) { 15 | pose_path_ = dataset_path_ / "poses.txt"; 16 | color_path_ = dataset_path_ / "images"; 17 | switch (_depth_type) { 18 | case 0: 19 | depth_path_ = dataset_path_ / "depth"; 20 | break; 21 | case 1: 22 | depth_path_ = dataset_path_ / "depth_filtered"; 23 | break; 24 | case 2: 25 | depth_path_ = dataset_path_ / "depth_with_noise"; 26 | break; 27 | default: 28 | throw std::runtime_error("Invalid depth type"); 29 | } 30 | 31 | focal_length_path_ = dataset_path_ / "focal.txt"; 32 | 33 | // gt_mesh_path_ = dataset_path_ / "gt_mesh.ply"; 34 | gt_mesh_path_ = dataset_path_ / "gt_mesh_culled.ply"; 35 | 36 | load_intrinsics(); 37 | } 38 | 39 | std::filesystem::path focal_length_path_; 40 | std::filesystem::path gt_mesh_path_; 41 | 42 | void load_data() override { 43 | depth_poses_ = load_poses(pose_path_, false, 0)[0]; 44 | TORCH_CHECK(depth_poses_.size(0) > 0); 45 | /* // https://github.com/PRBonn/semantic-kitti-api/issues/115 46 | // https://github.com/PRBonn/kiss-icp/issues/156#issuecomment-1551695713 47 | https://github.com/PRBonn/semantic-kitti-api/issues/133 48 | poses_ = T_S_B_.matmul(poses_.matmul(T_B_S_)); */ 49 | depth_poses_ = 50 | coords::change_world_system(depth_poses_, dataset_system_type_); 51 | depth_poses_ = 52 | coords::change_camera_system(depth_poses_, dataset_system_type_); 53 | color_poses_ = depth_poses_; 54 | 55 | load_colors(".png", "img", false, true); 56 | TORCH_CHECK(depth_poses_.size(0) == raw_color_filelists_.size()); 57 | load_depths(".png", "depth", false, true); 58 | TORCH_CHECK(raw_color_filelists_.size() == raw_depth_filelists_.size()); 59 | } 60 | 61 | virtual void load_focal_length() { 62 | if (!std::filesystem::exists(focal_length_path_)) { 63 | throw std::runtime_error("Focal Length file does not exist"); 64 | } 65 | std::ifstream import_file(focal_length_path_, std::ios::in); 66 | if (!import_file) { 67 | throw std::runtime_error("Could not open focal length file"); 68 | } 69 | std::string line; 70 | std::getline(import_file, line); 71 | sensor_.camera.focal_length_ = std::stof(line); 72 | import_file.close(); 73 | } 74 | 75 | void load_intrinsics() override { 76 | sensor_.camera.width = 640; 77 | sensor_.camera.height = 480; 78 | 79 | load_focal_length(); 80 | sensor_.camera.fx = sensor_.camera.focal_length_; 81 | sensor_.camera.fy = sensor_.camera.focal_length_; 82 | sensor_.camera.cx = 0.5 * (sensor_.camera.width - 1); 83 | sensor_.camera.cy = 0.5 * (sensor_.camera.height - 1); 84 | } 85 | 86 | std::string get_gt_mesh_path() override { return gt_mesh_path_.string(); } 87 | }; 88 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/oxford_spires_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "base_parser.h" 3 | #include "utils/coordinates.h" 4 | #include "utils/sensor_utils/cameras.hpp" 5 | #include 6 | 7 | namespace dataparser { 8 | struct Spires : DataParser { 9 | explicit Spires(const std::filesystem::path &_dataset_path, 10 | const torch::Device &_device = torch::kCPU, 11 | const bool &_preload = true, const float &_res_scale = 1.0, 12 | const sensor::Sensors &_sensor = sensor::Sensors()) 13 | : DataParser(_dataset_path, _device, _preload, _res_scale, 14 | coords::SystemType::OpenCV, _sensor) { 15 | // export undistorted images 16 | pose_path_ = dataset_path_ / "color_poses.txt"; 17 | depth_pose_path_ = dataset_path_ / "depth_poses.txt"; 18 | color_path_ = dataset_path_ / "undistorted_images"; 19 | depth_path_ = dataset_path_ / "depths"; // origin data 20 | // /media/chrisliu/T9/Datasets/Oxford_Spires_Dataset/bodleian_library/02 21 | // └── 02 22 | // ├── gt-tum.txt 23 | // ├── images 24 | // │   ├── cam0 25 | // │   ├── cam1 26 | // │   └── cam2 27 | // └── lidar-clouds 28 | dataset_name_ = dataset_path_.filename(); 29 | 30 | // pose_path_ = dataset_path_ / "gt-tum.txt"; 31 | // color_path_ = dataset_path_ / "images" / "cam0"; 32 | depth_type_ = DepthType::PCD; 33 | load_intrinsics(); 34 | load_data(); 35 | } 36 | 37 | torch::Tensor T_C0_L, T_C0_C2; 38 | 39 | std::filesystem::path depth_pose_path_; 40 | void load_data() override { 41 | if (!std::filesystem::exists(pose_path_) || 42 | !std::filesystem::exists(depth_pose_path_) || 43 | !std::filesystem::exists(color_path_) || 44 | !std::filesystem::exists(depth_path_)) { 45 | pose_path_ = dataset_path_ / "gt-tum.txt"; 46 | color_path_ = dataset_path_ / "images" / "cam0"; 47 | depth_path_ = dataset_path_ / "lidar-clouds"; 48 | 49 | auto pose_data = load_poses(pose_path_, false, 3); 50 | auto T_W_B = pose_data[0]; 51 | auto T_W_L = T_W_B.matmul(sensor_.T_B_L); 52 | time_stamps_ = pose_data[1]; 53 | TORCH_CHECK(T_W_L.size(0) > 0); 54 | // auto T_W_C0 = 55 | // coords::change_world_system(T_C0_C0, coords::SystemType::Kitti); 56 | color_poses_ = T_W_L.matmul(sensor_.T_C_L.inverse()); 57 | depth_poses_ = T_W_L; 58 | 59 | load_colors(".jpg", "", false, false); 60 | std::cout << "color_poses: " << color_poses_.size(0) << " color_files" 61 | << raw_color_filelists_.size() << std::endl; 62 | load_depths(".pcd", "", false, false); 63 | std::cout << "depth_poses: " << depth_poses_.size(0) 64 | << " depth_files: " << raw_depth_filelists_.size() << std::endl; 65 | 66 | // export undistorted images 67 | color_path_ = dataset_path_ / "undistorted_images"; 68 | std::filesystem::create_directories(color_path_); 69 | pose_path_ = dataset_path_ / "color_poses.txt"; 70 | std::ofstream color_pose_file(pose_path_); 71 | for (int i = 0; i < raw_color_filelists_.size(); i++) { 72 | auto color_image = get_image_cv_mat(i); 73 | auto undistorted_img = sensor_.camera.undistort(color_image); 74 | auto undistorted_img_path = 75 | color_path_ / raw_color_filelists_[i].filename(); 76 | cv::imwrite(undistorted_img_path, undistorted_img); 77 | 78 | auto T_W_C = color_poses_[i]; 79 | for (int i = 0; i < 4; i++) { 80 | for (int j = 0; j < 4; j++) { 81 | color_pose_file << T_W_C[i][j].item() << " "; 82 | } 83 | color_pose_file << "\n"; 84 | } 85 | } 86 | // export align pose depth 87 | depth_path_ = dataset_path_ / "depths"; 88 | std::filesystem::create_directories(depth_path_); 89 | std::ofstream depth_pose_file(depth_pose_path_); 90 | for (int i = 0; i < raw_depth_filelists_.size(); i++) { 91 | // copy depth file to undistorted_images 92 | std::filesystem::copy_file( 93 | raw_depth_filelists_[i], 94 | depth_path_ / raw_depth_filelists_[i].filename(), 95 | std::filesystem::copy_options::overwrite_existing); 96 | 97 | auto T_W_L = depth_poses_[i]; 98 | for (int i = 0; i < 4; i++) { 99 | for (int j = 0; j < 4; j++) { 100 | depth_pose_file << T_W_L[i][j].item() << " "; 101 | } 102 | depth_pose_file << "\n"; 103 | } 104 | } 105 | } 106 | 107 | time_stamps_ = torch::Tensor(); // reset time_stamps 108 | 109 | color_poses_ = load_poses(pose_path_, false, 0)[0]; 110 | TORCH_CHECK(color_poses_.size(0) > 0); 111 | depth_poses_ = load_poses(depth_pose_path_, false, 0)[0]; 112 | TORCH_CHECK(depth_poses_.size(0) > 0); 113 | 114 | load_colors(".jpg", "", false, true); 115 | TORCH_CHECK(color_poses_.size(0) == raw_color_filelists_.size()); 116 | load_depths(".pcd", "", false, true); 117 | TORCH_CHECK(depth_poses_.size(0) == raw_depth_filelists_.size()); 118 | } 119 | 120 | void load_intrinsics() override { 121 | // auto scale = 0.5f; // output image from Fastlivo is 640x512 122 | auto scale = 1.0f; // output image from Fastlivo is 1280x1024 123 | sensor_.camera.width = scale * sensor_.camera.width; 124 | sensor_.camera.height = scale * sensor_.camera.height; 125 | // HKU 126 | sensor_.camera.fx = scale * sensor_.camera.fx; 127 | sensor_.camera.fy = scale * sensor_.camera.fy; 128 | sensor_.camera.cx = scale * sensor_.camera.cx; 129 | sensor_.camera.cy = scale * sensor_.camera.cy; 130 | } 131 | 132 | std::vector get_distance_ndir_zdirn(const int &idx) override { 133 | /** 134 | * @description: 135 | * @return {distance, ndir, dir_norm}, where ndir.norm = 1; 136 | {[height width 1], [height width 3], [height width 1]} 137 | */ 138 | 139 | auto pointcloud = get_depth_image(idx); 140 | // [height width 1] 141 | auto distance = pointcloud.norm(2, -1, true); 142 | auto ndir = pointcloud / distance; 143 | return {distance, ndir, distance}; 144 | } 145 | }; 146 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/r3live_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "data_loader/data_parsers/rosbag_parser.hpp" 3 | #include "utils/coordinates.h" 4 | #include "utils/sensor_utils/cameras.hpp" 5 | #include 6 | 7 | namespace dataparser { 8 | struct R3live : Rosbag { 9 | explicit R3live(const std::filesystem::path &_bag_path, 10 | const torch::Device &_device = torch::kCPU, 11 | const bool &_preload = true, const float &_res_scale = 1.0, 12 | const int &_dataset_system_type = coords::SystemType::OpenCV) 13 | : Rosbag(_bag_path, _device, _preload, _res_scale, _dataset_system_type) { 14 | depth_type_ = DepthType::PLY; 15 | 16 | dataset_name_ = bag_path_.filename(); 17 | dataset_name_ = dataset_name_.replace_extension(); 18 | 19 | pose_topic = "/Odometry"; 20 | color_topic = "/camera/image_color/compressed"; 21 | depth_topic = "/cloud_registered_body"; 22 | 23 | // camera to imu 24 | sensor_.T_B_C = 25 | torch::tensor({{-0.00113207, -0.0158688, 0.999873, 0.050166}, 26 | {-0.9999999, -0.000486594, -0.00113994, 0.0474116}, 27 | {0.000504622, -0.999874, -0.0158682, -0.0312415}, 28 | {0.0, 0.0, 0.0, 1.0}}, 29 | torch::kFloat); 30 | // lidar to imu 31 | sensor_.T_B_L = torch::tensor({{1.0, 0.0, 0.0, 0.04165}, 32 | {0.0, 1.0, 0.0, 0.02326}, 33 | {0.0, 0.0, 1.0, -0.0284}, 34 | {0.0, 0.0, 0.0, 1.0}}, 35 | torch::kFloat); 36 | load_intrinsics(); 37 | } 38 | 39 | void load_intrinsics() override { 40 | // Replica/cam_params.json 41 | sensor_.camera.width = 1200; 42 | sensor_.camera.height = 1024; 43 | sensor_.camera.fx = 863.4241; 44 | sensor_.camera.fy = 863.4171; 45 | sensor_.camera.cx = 640.6808; 46 | sensor_.camera.cy = 518.3392; 47 | depth_scale_inv_ = 1.0; 48 | 49 | sensor_.camera.set_distortion(-0.1080, 0.1050, -1.2872e-04, 5.7923e-05, 50 | -0.0222); 51 | } 52 | }; 53 | } // namespace dataparser -------------------------------------------------------------------------------- /include/data_loader/data_parsers/replica_parser.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "base_parser.h" 3 | #include "utils/coordinates.h" 4 | #include "utils/sensor_utils/cameras.hpp" 5 | 6 | namespace dataparser { 7 | struct Replica : DataParser { 8 | explicit Replica(const std::string &_dataset_path, 9 | const torch::Device &_device = torch::kCPU, 10 | const bool &_preload = true, const float &_res_scale = 1.0) 11 | : DataParser(_dataset_path, _device, _preload, _res_scale, 12 | coords::SystemType::OpenCV) { 13 | pose_path_ = dataset_path_ / "traj.txt"; 14 | color_path_ = dataset_path_ / "results"; 15 | depth_path_ = color_path_; 16 | 17 | eval_pose_path_ = dataset_path_ / "eval" / "traj.txt"; 18 | eval_color_path_ = dataset_path_ / "eval" / "results"; 19 | eval_depth_path_ = eval_color_path_; 20 | 21 | gt_mesh_path_ = dataset_path_.parent_path() / "cull_replica_mesh" / 22 | (dataset_path_.filename().string() + "_culled.ply"); 23 | // gt_mesh_path_ = dataset_path_.parent_path() / "cull_replica_mesh" / 24 | // (dataset_path_.filename().string() + ".ply"); 25 | // gt_mesh_path_ = dataset_path_.parent_path() / 26 | // (dataset_path_.filename().string() + "_mesh.ply"); 27 | 28 | load_intrinsics(); 29 | if (std::filesystem::exists(pose_path_)) { 30 | load_data(); 31 | } else { 32 | std::cout << "pose_path_ does not exist: " << pose_path_ << std::endl; 33 | } 34 | } 35 | 36 | std::filesystem::path gt_mesh_path_; 37 | 38 | void load_data() override { 39 | depth_poses_ = load_poses(pose_path_, false, 1)[0]; 40 | color_poses_ = depth_poses_; 41 | TORCH_CHECK(depth_poses_.size(0) > 0); 42 | 43 | load_colors(".jpg", "frame", false, false); 44 | TORCH_CHECK(depth_poses_.size(0) == raw_color_filelists_.size()); 45 | load_depths(".png", "depth", false, false); 46 | TORCH_CHECK(raw_color_filelists_.size() == raw_depth_filelists_.size()); 47 | 48 | load_eval_data(); 49 | } 50 | 51 | void load_eval_data() override { 52 | eval_color_poses_ = load_poses(eval_pose_path_, false, 1)[0]; 53 | if (eval_color_poses_.size(0) <= 0) { 54 | std::cout << "eval_color_poses_.size(0) <= 0" << std::endl; 55 | return; 56 | } 57 | 58 | load_colors(".jpg", "frame", true); 59 | TORCH_CHECK(eval_color_poses_.size(0) == eval_color_filelists_.size()); 60 | load_depths(".png", "depth", true); 61 | TORCH_CHECK(eval_color_filelists_.size() == eval_depth_filelists_.size()); 62 | 63 | eval_depth_poses_ = eval_color_poses_; 64 | } 65 | 66 | void load_intrinsics() override { 67 | // Replica/cam_params.json 68 | sensor_.camera.width = 1200; 69 | sensor_.camera.height = 680; 70 | sensor_.camera.fx = 600.0; 71 | sensor_.camera.fy = 600.0; 72 | sensor_.camera.cx = 599.5; 73 | sensor_.camera.cy = 339.5; 74 | depth_scale_inv_ = 1.0 / 6553.5; 75 | } 76 | 77 | std::string get_gt_mesh_path() override { return gt_mesh_path_.string(); } 78 | }; 79 | } // namespace dataparser -------------------------------------------------------------------------------- /include/mesher/cumcubes/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Zhihao Liang. All rights reserved. 2 | from pathlib import Path 3 | from typing import List, Optional, Sequence, Tuple, Union, Callable 4 | 5 | import torch 6 | import numpy as np 7 | 8 | from . import src as _C 9 | from .utils import Timer, scale_to_bound 10 | 11 | 12 | def marching_cubes( 13 | density_grid: Union[torch.Tensor, np.ndarray], 14 | thresh: float, 15 | scale: Optional[Union[float, Sequence]]=None, 16 | verbose: bool=False, 17 | cpu: bool=False 18 | ) -> Tuple[torch.Tensor]: 19 | """python wrapper of marching cubes 20 | 21 | Args: 22 | density_grid (Union[torch.Tensor, np.ndarray]): 23 | input density grid to realize marching cube 24 | thresh (float): thresh of marching cubes 25 | scale (Optional[Union[float, Sequence]], optional): 26 | the scale of density grid. Defaults to None. 27 | verbose (bool, optional): 28 | print verbose informations or not. Defaults to False. 29 | cpu (bool, optional): 30 | cpu mode(wrapper of mcubes). Defaults to False. 31 | 32 | Returns: 33 | Tuple[torch.Tensor]: vertices and faces 34 | """ 35 | # get the bound according to the given scale 36 | lower: List[float] 37 | upper: List[float] 38 | # process scale as the bounding box 39 | if scale is None: 40 | lower = [0.0, 0.0, 0.0] 41 | upper = [density_grid.shape[0], density_grid.shape[1], density_grid.shape[2]] 42 | else: 43 | lower, upper = scale_to_bound(scale) 44 | 45 | if cpu or not torch.cuda.is_available(): 46 | try: 47 | import mcubes 48 | except: 49 | raise ImportError("the cpu mode cumcubes is the wrapper of `mcubes`, please install the mcubes") 50 | 51 | density_grid = density_grid.detach().cpu().numpy() 52 | vertices, faces = mcubes.marching_cubes(density_grid, thresh) 53 | offset = np.array(lower) 54 | scale = (np.array(upper) - np.array(lower)) / np.array(density_grid.shape) 55 | vertices = vertices / scale + offset 56 | 57 | vertices = torch.tensor(vertices) 58 | faces = torch.tensor(faces.astype(np.int64)) 59 | else: 60 | # process density_grid 61 | if isinstance(density_grid, np.ndarray): density_grid = torch.tensor(density_grid) 62 | density_grid = density_grid.cuda() 63 | density_grid = density_grid.to(torch.float32) 64 | 65 | if (density_grid.shape[0] < 2 or density_grid.shape[1] < 2 or density_grid.shape[2] < 2): 66 | raise ValueError() 67 | 68 | vertices, faces = _C.marching_cubes(density_grid, thresh, lower, upper) 69 | 70 | if verbose: 71 | print(f"#vertices={vertices.shape[0]}") 72 | print(f"#triangles={faces.shape[0]}") 73 | 74 | return vertices, faces 75 | 76 | 77 | def marching_cubes_func( 78 | scale: Optional[Union[float, Sequence]]=None, 79 | num_x: int=100, 80 | num_y: int=100, 81 | num_z: int=100, 82 | func: Callable=lambda x, y, z: x**2 + y**2 + z**2, 83 | thresh: float=16, 84 | verbose: bool=False, 85 | cpu: bool=False 86 | ) -> Tuple[torch.Tensor]: 87 | """python wrapper of marching cubes via the given function 88 | 89 | Args: 90 | scale (Optional[Union[float, Sequence]], optional): 91 | the scale of density grid. Defaults to None.. Defaults to None. 92 | num_x (int, optional): the sampling scale of x-axis. Defaults to 100. 93 | num_y (int, optional): the sampling scale of y-axis. Defaults to 100. 94 | num_z (int, optional): the sampling scale of z-axis. Defaults to 100. 95 | func (Callable, optional): the funcion for marching cudes. Defaults to lambdax. 96 | thresh (float, optional): thresh of marching cubes. Defaults to 16. 97 | verbose (bool, optional): 98 | print verbose informations or not. Defaults to False. 99 | cpu (bool, optional): 100 | cpu mode(wrapper of mcubes). Defaults to False. 101 | 102 | Returns: 103 | Tuple[torch.Tensor]: vertices and faces 104 | """ 105 | # convert the scale to the bound 106 | lower, upper = scale_to_bound(scale) 107 | 108 | if cpu or not torch.cuda.is_available(): 109 | try: 110 | import mcubes 111 | except: 112 | raise ImportError("the cpu mode cumcubes is the wrapper of `mcubes`, please install the mcubes") 113 | 114 | vertices, faces = mcubes.marching_cubes_func( 115 | tuple(lower), 116 | tuple(upper), 117 | num_x, 118 | num_y, 119 | num_z, 120 | func, 121 | thresh) 122 | 123 | vertices = torch.tensor(vertices) 124 | faces = torch.tensor(faces.astype(np.int64)) 125 | else: 126 | # get the sample grid 127 | x = torch.linspace(lower[0], upper[0], num_x) 128 | y = torch.linspace(lower[0], upper[0], num_y) 129 | z = torch.linspace(lower[0], upper[0], num_z) 130 | X, Y, Z = torch.meshgrid(x, y, z, indexing="ij") 131 | sample_points = torch.stack((X, Y, Z), dim=-1).reshape(num_x, num_y, num_z, 3) 132 | 133 | vertices, faces = _C.marching_cubes_func(sample_points, thresh, lower, upper, func) 134 | 135 | if verbose: 136 | print(f"#vertices={vertices.shape[0]}\n") 137 | print(f"#triangles={faces.shape[0]}\n") 138 | 139 | return vertices, faces 140 | 141 | 142 | def save_mesh( 143 | vertices: Union[torch.Tensor, np.ndarray], 144 | faces: Union[torch.Tensor, np.ndarray], 145 | colors: Optional[Union[torch.Tensor, np.ndarray]]=None, 146 | filename: Union[str, Path]="temp.ply", 147 | verbose: bool=False 148 | ) -> None: 149 | """save mesh into the given filename 150 | 151 | Args: 152 | vertices (Union[torch.Tensor, np.ndarray]): vertices of the mesh to save 153 | faces (Union[torch.Tensor, np.ndarray]): faces of the mesh to save 154 | colors (Optional[Union[torch.Tensor, np.ndarray]], optional): 155 | vertices of the mesh to save. Defaults to None. 156 | filename (Union[str, Path], optional): 157 | the save path. Defaults to "temp.ply". 158 | verbose (bool, optional): 159 | print verbose mention or not. Defaults to False. 160 | """ 161 | 162 | if isinstance(filename, Path): 163 | filename = str(filename) 164 | 165 | if isinstance(vertices, np.ndarray): vertices = torch.tensor(vertices) 166 | if isinstance(faces, np.ndarray): faces = torch.tensor(faces) 167 | 168 | # process colors 169 | if colors is None: 170 | colors = torch.ones_like(vertices) * 127 171 | elif isinstance(colors, np.ndarray): 172 | colors = torch.tensor(colors) 173 | colors = colors.to(torch.uint8) 174 | 175 | if filename.endswith(".ply"): 176 | _C.save_mesh_as_ply(filename, vertices, faces, colors) 177 | else: 178 | raise NotImplementedError() 179 | 180 | if verbose: 181 | print(f"save as {filename} successfully!") 182 | 183 | 184 | __all__ = ["Timer", "marching_cubes", "save_mesh"] 185 | -------------------------------------------------------------------------------- /include/mesher/cumcubes/include/cumcubes.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | using torch::Tensor; 6 | 7 | namespace mc 8 | { 9 | // defination 10 | std::vector marching_cubes( 11 | const Tensor&, 12 | const float, 13 | const std::vector, 14 | const std::vector); 15 | std::vector marching_cubes_wrapper( 16 | const Tensor&, 17 | const float, 18 | const float*, 19 | const float*); 20 | 21 | void save_mesh_as_ply(const std::string, Tensor, Tensor, Tensor); 22 | } // namespace mc(marching cube) 23 | 24 | // Utils 25 | #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") 26 | 27 | #define CHECK_CPU(x) TORCH_CHECK(!x.is_cuda(), #x " must be a CPU tensor") 28 | 29 | #define CHECK_CONTIGUOUS(x) \ 30 | TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 31 | 32 | #define CHECK_INPUT(x) \ 33 | CHECK_CUDA(x); \ 34 | CHECK_CONTIGUOUS(x) 35 | 36 | #define CHECK_CPU_INPUT(x) \ 37 | CHECK_CPU(x); \ 38 | CHECK_CONTIGUOUS(x) 39 | 40 | -------------------------------------------------------------------------------- /include/mesher/cumcubes/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | // Copyright 2021 Zhihao Liang 2 | #include "cumcubes.hpp" 3 | 4 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 5 | m.def("marching_cubes", &mc::marching_cubes); 6 | m.def("marching_cubes_func", &mc::marching_cubes_func); 7 | m.def("save_mesh_as_ply", &mc::save_mesh_as_ply); 8 | } 9 | 10 | -------------------------------------------------------------------------------- /include/mesher/cumcubes/src/cumcubes.cpp: -------------------------------------------------------------------------------- 1 | // Copyright 2021 Zhihao Liang 2 | #include 3 | #include 4 | #include 5 | 6 | #include "cumcubes.hpp" 7 | 8 | 9 | std::vector mc::marching_cubes( 10 | const torch::Tensor& density_grid, 11 | const float thresh, 12 | const std::vector lower, 13 | const std::vector upper 14 | ) { 15 | // check 16 | CHECK_INPUT(density_grid); 17 | TORCH_CHECK(density_grid.ndimension() == 3) 18 | 19 | assert(lower.size() == 3); 20 | assert(upper.size() == 3); 21 | 22 | const float l[3] = {lower[0], lower[1], lower[2]}; 23 | const float u[3] = {upper[0], upper[1], upper[2]}; 24 | 25 | std::vector results = mc::marching_cubes_wrapper(density_grid, thresh, l, u); 26 | 27 | return results; 28 | } 29 | 30 | void mc::save_mesh_as_ply( 31 | const std::string filename, 32 | torch::Tensor vertices, 33 | torch::Tensor faces, 34 | torch::Tensor colors 35 | ) { 36 | CHECK_CONTIGUOUS(vertices); 37 | CHECK_CONTIGUOUS(faces); 38 | CHECK_CONTIGUOUS(colors); 39 | assert(colors.dtype() == torch::kUInt8); 40 | 41 | if (vertices.is_cuda()) { vertices = vertices.to(torch::kCPU); } 42 | if (faces.is_cuda()) { faces = faces.to(torch::kCPU); } 43 | if (colors.is_cuda()) { colors = colors.to(torch::kCPU); } 44 | 45 | std::ofstream ply_file(filename, std::ios::out | std::ios::binary); 46 | ply_file << "ply\n"; 47 | ply_file << "format binary_little_endian 1.0\n"; 48 | ply_file << "element vertex " << vertices.size(0) << std::endl; 49 | ply_file << "property float x\n"; 50 | ply_file << "property float y\n"; 51 | ply_file << "property float z\n"; 52 | ply_file << "property uchar red\n"; 53 | ply_file << "property uchar green\n"; 54 | ply_file << "property uchar blue\n"; 55 | ply_file << "element face " << faces.size(0) << std::endl; 56 | ply_file << "property list int int vertex_index\n"; 57 | 58 | ply_file << "end_header\n"; 59 | 60 | const int32_t num_vertices = vertices.size(0), num_faces = faces.size(0); 61 | 62 | const float* vertices_ptr = vertices.data_ptr(); 63 | const uint8_t* colors_ptr = colors.data_ptr(); 64 | for (int32_t i = 0; i < num_vertices; ++i) { 65 | ply_file.write((char *)&(vertices_ptr[i * 3]), 3 * sizeof(float)); 66 | ply_file.write((char *)&(colors_ptr[i * 3]), 3 * sizeof(uint8_t)); 67 | } 68 | 69 | torch::Tensor faces_head = torch::ones({num_faces, 1}, 70 | torch::TensorOptions().dtype(torch::kInt).device(torch::kCPU)) * 3; 71 | 72 | torch::Tensor padded_faces = torch::cat({faces_head, faces}, 1); // [num_faces, 4] 73 | CHECK_CONTIGUOUS(padded_faces); 74 | 75 | const int32_t* faces_ptr = padded_faces.data_ptr(); 76 | ply_file.write((char *)&(faces_ptr[0]), num_faces * 4 * sizeof(int32_t)); 77 | 78 | ply_file.close(); 79 | } 80 | 81 | -------------------------------------------------------------------------------- /include/mesher/cumcubes/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import re 3 | from time import time 4 | from typing import List, Tuple, Union, Sequence 5 | 6 | import torch 7 | import numpy as np 8 | 9 | 10 | def scale_to_bound(scale: Union[float, Sequence]) -> Tuple[List[float]]: 11 | if isinstance(scale, float): 12 | lower = [0.0, 0.0, 0.0] 13 | upper = [scale, scale, scale] 14 | elif isinstance(scale, (list, tuple, np.ndarray, torch.Tensor)): 15 | if len(scale) == 3: 16 | lower = [0.0, 0.0, 0.0] 17 | upper = [i for i in scale] 18 | elif len(scale) == 2: 19 | if isinstance(scale[0], float): 20 | lower = [scale[0]] * 3 21 | upper = [scale[1]] * 3 22 | else: 23 | assert len(scale[0]) == len(scale[1]) == 3 24 | lower = [i for i in scale[0]] 25 | upper = [i for i in scale[1]] 26 | else: 27 | raise TypeError() 28 | else: 29 | raise TypeError() 30 | 31 | return lower, upper 32 | 33 | 34 | 35 | class TimerError(Exception): 36 | 37 | def __init__(self, message): 38 | self.message = message 39 | super(TimerError, self).__init__(message) 40 | 41 | 42 | class Timer: 43 | """A flexible Timer class. 44 | Examples: 45 | >>> import time 46 | >>> import mmcv 47 | >>> with mmcv.Timer(): 48 | >>> # simulate a code block that will run for 1s 49 | >>> time.sleep(1) 50 | 1.000 51 | >>> with mmcv.Timer(print_tmpl='it takes {:.1f} seconds'): 52 | >>> # simulate a code block that will run for 1s 53 | >>> time.sleep(1) 54 | it takes 1.0 seconds 55 | >>> timer = mmcv.Timer() 56 | >>> time.sleep(0.5) 57 | >>> print(timer.since_start()) 58 | 0.500 59 | >>> time.sleep(0.5) 60 | >>> print(timer.since_last_check()) 61 | 0.500 62 | >>> print(timer.since_start()) 63 | 1.000 64 | """ 65 | def __init__(self, print_tmpl=None, start=True): 66 | self._is_running = False 67 | if (print_tmpl 68 | is not None) and not re.findall(r"({:.*\df})", print_tmpl): 69 | print_tmpl += " {:.3f}" 70 | # raise ValueError("`print_tmpl` must has the `{:.nf}` to show time.") 71 | self.print_tmpl = print_tmpl if print_tmpl else "{:.3f}" 72 | if start: 73 | self.start() 74 | 75 | @property 76 | def is_running(self): 77 | """bool: indicate whether the timer is running""" 78 | return self._is_running 79 | 80 | def __enter__(self): 81 | self.start() 82 | return self 83 | 84 | def __exit__(self, type, value, traceback): 85 | print(self.print_tmpl.format(self.since_last_check())) 86 | self._is_running = False 87 | 88 | def start(self): 89 | """Start the timer.""" 90 | if not self._is_running: 91 | self._t_start = time() 92 | self._is_running = True 93 | self._t_last = time() 94 | 95 | def since_start(self): 96 | """Total time since the timer is started. 97 | Returns: 98 | float: Time in seconds. 99 | """ 100 | if not self._is_running: 101 | raise TimerError('timer is not running') 102 | self._t_last = time() 103 | return self._t_last - self._t_start 104 | 105 | def since_last_check(self): 106 | """Time since the last checking. 107 | Either :func:`since_start` or :func:`since_last_check` is a checking 108 | operation. 109 | Returns: 110 | float: Time in seconds. 111 | """ 112 | if not self._is_running: 113 | raise TimerError('timer is not running') 114 | dur = time() - self._t_last 115 | self._t_last = time() 116 | return dur 117 | 118 | -------------------------------------------------------------------------------- /include/mesher/cumcubes/version.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.0.3" 2 | -------------------------------------------------------------------------------- /include/mesher/mesher.cpp: -------------------------------------------------------------------------------- 1 | #include "mesher.h" 2 | #include "cumcubes.hpp" 3 | #include "utils/tqdm.hpp" 4 | 5 | using namespace std; 6 | 7 | namespace mesher { 8 | 9 | #ifdef ENABLE_ROS 10 | void tensor_to_mesh(mesh_msgs::MeshGeometry &mesh, 11 | mesh_msgs::MeshVertexColors &mesh_color, 12 | const torch::Tensor &vertice, const torch::Tensor &face, 13 | const torch::Tensor &color) { 14 | auto vertex_num = vertice.size(0); 15 | auto face_num = face.size(0); 16 | mesh.faces.resize(face_num); 17 | mesh.vertices.resize(vertex_num); 18 | mesh_color.vertex_colors.resize(vertex_num); 19 | 20 | auto vertice_float64 = vertice.to(torch::kFloat64).contiguous().cpu(); 21 | torch::Tensor color_cat; 22 | TORCH_CHECK(color.size(-1) == 3) 23 | color_cat = 24 | torch::cat({color, torch::ones({color.size(0), 1}, color.options())}, 1) 25 | .contiguous() 26 | .cpu(); 27 | 28 | memcpy(mesh.vertices.data(), vertice_float64.data_ptr(), 29 | vertice_float64.numel() * sizeof(double)); 30 | memcpy(mesh_color.vertex_colors.data(), color_cat.data_ptr(), 31 | color_cat.numel() * sizeof(float)); 32 | 33 | auto face_int32 = face.to(torch::kInt32).contiguous().cpu(); 34 | memcpy(mesh.faces.data(), face_int32.data_ptr(), 35 | face_int32.numel() * sizeof(uint32_t)); 36 | } 37 | #endif 38 | 39 | Mesher::Mesher() {}; 40 | 41 | void Mesher::save_mesh(const std::string &output_path, bool vis_attribute, 42 | const std::string &prefix, 43 | dataparser::DataParser::Ptr dataparser_ptr_) { 44 | if (vec_face_.empty()) { 45 | cout << "\n" 46 | << "No mesh to save! Please first meshing using 'v' first!" << "\n"; 47 | return; 48 | } 49 | 50 | auto vertice_attrs = torch::cat({vec_face_attr_}, 0).contiguous(); 51 | auto face = torch::cat({vec_face_}, 0).contiguous(); 52 | auto vertices = torch::cat({vec_vertice_}, 0).contiguous(); 53 | std::string filename = output_path + "/" + prefix + "mesh.ply"; 54 | if (vertice_attrs.numel() == 0) { 55 | vertice_attrs = torch::full_like(vertices, 127).to(torch::kUInt8); 56 | } 57 | mc::save_mesh_as_ply(filename, vertices, face, vertice_attrs); 58 | if (dataparser_ptr_) { 59 | std::string cull_filename = output_path + "/" + prefix + "mesh_culled.ply"; 60 | auto cull_results = 61 | cull_mesh(dataparser_ptr_, vertices, face, vertice_attrs); 62 | mc::save_mesh_as_ply(cull_filename, cull_results[0], cull_results[1], 63 | cull_results[2]); 64 | } 65 | vec_face_attr_.clear(); 66 | vec_face_.clear(); 67 | vec_vertice_.clear(); 68 | printf("\033[1;32mExport mesh saved to %s\n\033[0m", filename.c_str()); 69 | } 70 | 71 | std::vector 72 | Mesher::cull_mesh(dataparser::DataParser::Ptr dataparser_ptr_, 73 | torch::Tensor vertices, torch::Tensor faces, 74 | torch::Tensor vertice_attrs) { 75 | int n_imgs = dataparser_ptr_->raw_depth_filelists_.size(); 76 | auto whole_mask = 77 | torch::ones(vertices.size(0), vertices.options()).to(torch::kBool); 78 | auto device = vertices.device(); 79 | 80 | // Pre-compute constants and allocate reusable tensors 81 | static const auto K = torch::tensor({{dataparser_ptr_->sensor_.camera.fx, 0.f, 82 | dataparser_ptr_->sensor_.camera.cx}, 83 | {0.f, dataparser_ptr_->sensor_.camera.fy, 84 | dataparser_ptr_->sensor_.camera.cy}, 85 | {0.f, 0.f, 1.f}}, 86 | device); 87 | static const auto W = dataparser_ptr_->sensor_.camera.width; 88 | static const auto H = dataparser_ptr_->sensor_.camera.height; 89 | 90 | // Pre-compute homogeneous vertices (only done once) 91 | const auto ones = 92 | torch::ones({vertices.size(0), 1}, torch::kFloat32).to(device); 93 | const auto homo_points = torch::cat({vertices, ones}, 1).reshape({-1, 4, 1}); 94 | 95 | // Buffers for reuse in the loop 96 | auto cam_cord_homo = 97 | torch::empty({vertices.size(0), 4, 1}, vertices.options()); 98 | auto cam_cord = torch::empty({vertices.size(0), 3, 1}, vertices.options()); 99 | auto uv = torch::empty({vertices.size(0), 2}, vertices.options()); 100 | auto grid = torch::empty({vertices.size(0), 2}, vertices.options()); 101 | auto mask = 102 | torch::empty(vertices.size(0), vertices.options().dtype(torch::kBool)); 103 | 104 | auto iter_bar = tq::trange(n_imgs); 105 | iter_bar.set_prefix("Culling mesh"); 106 | 107 | for (const auto &i : iter_bar) { 108 | auto pose = 109 | dataparser_ptr_->get_pose(i, dataparser::DataType::RawDepth).to(device); 110 | auto depth = dataparser_ptr_->get_depth_image(i).to(device); 111 | auto w2c = torch::inverse(pose); 112 | 113 | // Transform points to camera space 114 | cam_cord_homo = w2c.matmul(homo_points); 115 | cam_cord = cam_cord_homo.slice(1, 0, 3); 116 | auto z = cam_cord.slice(1, -1); 117 | auto z_squeezed = z.squeeze(); 118 | 119 | // Project to image space 120 | uv = cam_cord / z.abs(); 121 | uv = K.matmul(uv); 122 | uv = uv.slice(1, 0, 2).squeeze(-1); 123 | 124 | // Create grid for sampling 125 | auto grid_x = uv.slice(1, 0, 1) / W; 126 | auto grid_y = uv.slice(1, 1) / H; 127 | grid = torch::cat({grid_x, grid_y}, 1); 128 | grid = 2 * grid - 1; // Convert to [-1,1] range 129 | 130 | // Sample depth from image 131 | auto input = depth.unsqueeze(0).unsqueeze(1).squeeze(-1); 132 | auto flow_field = grid.unsqueeze(0).unsqueeze(1); 133 | auto depth_samples = torch::nn::functional::grid_sample( 134 | input, flow_field, 135 | torch::nn::functional::GridSampleFuncOptions() 136 | .padding_mode(torch::kZeros) 137 | .align_corners(true)) 138 | .squeeze(); 139 | 140 | // Create visibility mask 141 | mask = (0 <= z_squeezed) & (uv.select(1, 0) < W) & (uv.select(1, 0) > 0) & 142 | (uv.select(1, 1) < H) & (uv.select(1, 1) > 0) & 143 | ((depth_samples + 0.02f) > z_squeezed); 144 | 145 | // Update the global mask 146 | whole_mask &= ~mask; 147 | } 148 | 149 | // Cull faces based on vertices 150 | auto face_mask = ~(whole_mask.index({faces}).all(1)); 151 | auto valid_face_idx = face_mask.nonzero().squeeze(); 152 | faces = faces.index({valid_face_idx}); 153 | 154 | return {vertices, faces, vertice_attrs}; 155 | } 156 | } // namespace mesher -------------------------------------------------------------------------------- /include/mesher/mesher.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | #include "data_loader/data_loader.h" 7 | 8 | #ifdef ENABLE_ROS 9 | #include 10 | #include 11 | #endif 12 | 13 | namespace mesher { 14 | #ifdef ENABLE_ROS 15 | void tensor_to_mesh(mesh_msgs::MeshGeometry &mesh, 16 | mesh_msgs::MeshVertexColors &mesh_color, 17 | const torch::Tensor &vertice, const torch::Tensor &face, 18 | const torch::Tensor &color); 19 | #endif 20 | 21 | class Mesher { 22 | public: 23 | Mesher(); 24 | 25 | std::vector vec_face_attr_, vec_face_, vec_vertice_; 26 | void save_mesh(const std::string &output_path, bool vis_attribute = false, 27 | const std::string &prefix = "", 28 | dataparser::DataParser::Ptr dataparser_ptr_ = nullptr); 29 | std::vector 30 | cull_mesh(dataparser::DataParser::Ptr dataparser_ptr_, torch::Tensor vertices, 31 | torch::Tensor faces, torch::Tensor vertice_attrs); 32 | }; 33 | } // namespace mesher -------------------------------------------------------------------------------- /include/neural_mapping/neural_mapping.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #ifdef ENABLE_ROS 6 | #include 7 | #include 8 | #include 9 | #endif 10 | 11 | // Add PCL headers for PCD file support 12 | #include 13 | #include 14 | 15 | #include "data_loader/data_loader.h" 16 | #include "neural_net/local_map.h" 17 | #include "rog_map_cuda/rog_map_class.cuh" 18 | 19 | class NeuralSLAM { 20 | public: 21 | typedef std::shared_ptr Ptr; 22 | NeuralSLAM(const int &mode, const std::filesystem::path &_config_path, 23 | const std::filesystem::path &_data_path = ""); 24 | 25 | #ifdef ENABLE_ROS 26 | NeuralSLAM(ros::NodeHandle &nh, const int &mode, 27 | const std::filesystem::path &_config_path, 28 | const std::filesystem::path &_data_path = ""); 29 | #endif 30 | 31 | private: 32 | rog_map::ROGMap::Ptr rog_map_ptr; 33 | 34 | std::vector vec_inrange_xyz; 35 | LocalMap::Ptr local_map_ptr; 36 | std::shared_ptr p_optimizer_; 37 | 38 | dataloader::DataLoader::Ptr data_loader_ptr; 39 | 40 | #ifdef ENABLE_ROS 41 | // ROS stuff 42 | ros::Subscriber rviz_pose_sub; 43 | ros::Publisher pose_pub, path_pub, odom_pub; 44 | ros::Publisher mesh_pub, mesh_color_pub, voxel_pub, vis_shift_map_pub, 45 | pointcloud_pub, occ_pub, unknown_pub, rgb_pub, depth_pub; 46 | 47 | nav_msgs::Path path_msg; 48 | 49 | torch::Tensor rviz_pose_ = torch::Tensor(); 50 | 51 | void register_subscriber(ros::NodeHandle &nh); 52 | void register_publisher(ros::NodeHandle &nh); 53 | 54 | void rviz_pose_callback(const geometry_msgs::PoseConstPtr &_rviz_pose_ptr); 55 | 56 | void visualization(const torch::Tensor &_xyz = {}, 57 | std_msgs::Header _header = std_msgs::Header()); 58 | void pub_pose(const torch::Tensor &_pose, 59 | std_msgs::Header _header = std_msgs::Header()); 60 | void pub_path(std_msgs::Header _header = std_msgs::Header()); 61 | void pub_voxel(const torch::Tensor &_xyz, 62 | std_msgs::Header _header = std_msgs::Header()); 63 | static void pub_pointcloud(const ros::Publisher &_pub, 64 | const torch::Tensor &_xyz, 65 | std_msgs::Header _header = std_msgs::Header()); 66 | static void pub_pointcloud(const ros::Publisher &_pub, 67 | const torch::Tensor &_xyz, 68 | const torch::Tensor &_rgb, 69 | std_msgs::Header _header = std_msgs::Header()); 70 | static void pub_image(const ros::Publisher &_pub, const torch::Tensor &_image, 71 | std_msgs::Header _header = std_msgs::Header()); 72 | 73 | void pub_render_image(std_msgs::Header _header); 74 | void pretrain_loop(); 75 | #endif 76 | 77 | // thread, buffer stuff 78 | std::mutex mapper_buf_mutex, train_mutex; 79 | std::condition_variable mapperer_buf_cond; 80 | std::queue mapper_pcl_buf, mapper_pose_buf; 81 | std::thread mapper_thread, keyboard_thread, misc_thread; 82 | 83 | bool save_image = false; 84 | 85 | void prefilter_data(const bool &export_img = false); 86 | 87 | bool build_occ_map(); 88 | 89 | DepthSamples sample(DepthSamples batch_ray_samples, const int &iter, 90 | const float &sample_std); 91 | 92 | torch::Tensor sdf_regularization(const torch::Tensor &xyz, 93 | const torch::Tensor &sdf = torch::Tensor(), 94 | const bool &curvate_enable = false, 95 | const std::string &name = ""); 96 | std::tuple sdf_train_batch_iter(const int &iter); 97 | torch::Tensor color_train_batch_iter(const int &iter); 98 | void train(int _opt_iter); 99 | 100 | void train_callback(const int &_opt_iter, const RaySamples &point_samples); 101 | 102 | void keyboard_loop(); 103 | void batch_train(); 104 | void misc_loop(); 105 | int misc_trigger = -1; 106 | 107 | void create_dir(const std::filesystem::path &base_path, 108 | std::filesystem::path &color_path, 109 | std::filesystem::path &depth_path, 110 | std::filesystem::path >_color_path, 111 | std::filesystem::path &render_color_path, 112 | std::filesystem::path >_depth_path, 113 | std::filesystem::path &render_depth_path, 114 | std::filesystem::path &acc_path); 115 | void render_image_colormap(int idx); 116 | 117 | std::vector render_image(const torch::Tensor &_pose, 118 | sensor::Cameras &camera, 119 | const float &_scale = 1.0, 120 | const bool &training = true); 121 | 122 | void render_path(bool eval, const int &fps = 30, const bool &save = true); 123 | void render_path(std::string pose_file, std::string camera_file = "", 124 | const int &fps = 30); 125 | 126 | float export_test_image(int idx = -1, const std::string &prefix = ""); 127 | void export_checkpoint(); 128 | 129 | void load_pretrained(const std::filesystem::path &_pretrained_path); 130 | void load_checkpoint(const std::filesystem::path &_checkpoint_path); 131 | 132 | void save_mesh(const bool &cull_mesh = false, const std::string &prefix = ""); 133 | void eval_mesh(); 134 | static void eval_render(); 135 | static void plot_log(const std::string &log_file); 136 | static void export_timing(bool print = false); 137 | 138 | bool end(); 139 | }; -------------------------------------------------------------------------------- /include/neural_net/encoding_map.cpp: -------------------------------------------------------------------------------- 1 | #include "encoding_map.h" 2 | #include "params/params.h" 3 | 4 | using namespace std; 5 | 6 | EncodingMap::EncodingMap(const torch::Tensor &_pos_W_M, std::string _name, 7 | float _x_min, float _x_max, float _y_min, float _y_max, 8 | float _z_min, float _z_max) 9 | : SubMap(_pos_W_M, _x_min, _x_max, _y_min, _y_max, _z_min, _z_max), 10 | name_(_name) { 11 | // int base_res = 16; 12 | // auto growth_factor = exp((log(8192) - log(base_res)) / (k_n_levels - 1)); 13 | 14 | // The number of encoded dimensions is n_levels * n_features_per_level 15 | nlohmann::json encoding_config = { 16 | {"otype", "Grid"}, 17 | {"type", "Hash"}, 18 | {"n_levels", k_n_levels}, 19 | {"n_features_per_level", k_n_features_per_level}, 20 | {"log2_hashmap_size", k_log2_hashmap_size}, 21 | {"base_resolution", 32}, 22 | {"per_level_scale", 2.0}, 23 | {"interpolation", "Linear"}}; 24 | 25 | p_encoder_tcnn_ = 26 | std::make_shared(3, encoding_config, "encoder_" + name_); 27 | 28 | p_color_encoder_tcnn_ = std::make_shared( 29 | 3, encoding_config, "color_encoder_" + name_); 30 | 31 | active_ = true; 32 | } 33 | 34 | torch::Tensor EncodingMap::encoding(const torch::Tensor &xyz, 35 | const int &encoding_type, 36 | const bool &normalized) { 37 | // Make sure the tcnn gets inputs between 0 and 1. 38 | torch::Tensor normalized_xyz; 39 | if (normalized) { 40 | normalized_xyz = xyz; 41 | } else { 42 | normalized_xyz = xyz_to_zp1_pts(xyz); 43 | } 44 | 45 | static bool scene_contraction = true; 46 | if (scene_contraction) { 47 | // static float zp1_leaf_size = k_leaf_size * k_map_size_inv; 48 | static float zp1_leaf_size = k_boundary_size * k_map_size_inv; 49 | static float inner_bound = 0.5f - zp1_leaf_size; 50 | auto tmp_xyz = normalized_xyz - 0.5f; 51 | auto tmp_xyz_norm = tmp_xyz.abs(); 52 | tmp_xyz = torch::where( 53 | tmp_xyz_norm < inner_bound, tmp_xyz, 54 | (inner_bound + zp1_leaf_size * (1 - inner_bound / tmp_xyz_norm)) * 55 | tmp_xyz / tmp_xyz_norm); 56 | normalized_xyz = tmp_xyz + 0.5f; 57 | } 58 | 59 | // [n, feat_dim] 60 | if (encoding_type == 0) 61 | return p_encoder_tcnn_->forward(normalized_xyz); 62 | else 63 | return p_color_encoder_tcnn_->forward(normalized_xyz); 64 | } 65 | 66 | void EncodingMap::activate() { 67 | state_mutex_.lock(); 68 | this->to(k_device); 69 | active_ = true; 70 | state_mutex_.unlock(); 71 | } 72 | 73 | void EncodingMap::freeze() { 74 | state_mutex_.lock(); 75 | this->to(torch::kCPU); 76 | active_ = false; 77 | state_mutex_.unlock(); 78 | } 79 | -------------------------------------------------------------------------------- /include/neural_net/encoding_map.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "neural_net/sub_map.h" 4 | #include "tcnn_binding/tcnn_binding.h" 5 | 6 | struct EncodingMap : SubMap { 7 | EncodingMap(const torch::Tensor &_pos_W_M, std::string _name, float _x_min, 8 | float _x_max, float _y_min, float _y_max, float _z_min, 9 | float _z_max); 10 | std::string name_; 11 | std::shared_ptr p_encoder_tcnn_, p_color_encoder_tcnn_; 12 | 13 | std::mutex state_mutex_; 14 | 15 | torch::Tensor encoding(const torch::Tensor &xyz, const int &encoding_type = 0, 16 | const bool &normalized = false); 17 | 18 | void activate(); 19 | 20 | void freeze(); 21 | }; -------------------------------------------------------------------------------- /include/neural_net/encodings/encodings.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hku-mars/M2Mapping/6075ad2e39a6f20fdb7db12003f9513a95b983cf/include/neural_net/encodings/encodings.cpp -------------------------------------------------------------------------------- /include/neural_net/encodings/encodings.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #include "tcnn_binding/tcnn_binding.h" 6 | struct SHEncoding : TCNNEncoding { 7 | /* Spherical harmonic encoding */ 8 | explicit SHEncoding(const int &_levels = 4) : levels_(_levels) { 9 | if (levels_ <= 0 || levels_ > 4) { 10 | throw std::invalid_argument("Spherical harmonic encoding only supports 1 " 11 | "to 4 levels, requested " + 12 | std::to_string(levels_)); 13 | } 14 | nlohmann::json encoding_config = { 15 | {"otype", "SphericalHarmonics"}, 16 | {"degree", levels_} // The SH degree up to which 17 | // to evaluate the encoding. 18 | // Produces degree^2 encoded 19 | // dimensions. 20 | }; 21 | 22 | init_encoding(3, encoding_config, "sh_encoding"); 23 | } 24 | int levels_; 25 | 26 | size_t get_out_dim() const override { return levels_ * levels_; } 27 | }; -------------------------------------------------------------------------------- /include/neural_net/local_map.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "mesher/mesher.h" 4 | #include "neural_net/encoding_map.h" 5 | #include "tcnn_binding/tcnn_binding.h" 6 | #include "utils/ray_utils/ray_utils.h" 7 | 8 | #ifdef ENABLE_ROS 9 | #include "ros/publisher.h" 10 | #endif 11 | struct LocalMap : SubMap { 12 | typedef std::shared_ptr Ptr; 13 | explicit LocalMap(const torch::Tensor &_pos_W_M, float _x_min, float _x_max, 14 | float _y_min, float _y_max, float _z_min, float _z_max); 15 | 16 | std::shared_ptr p_mesher_; 17 | 18 | std::shared_ptr p_encoding_map_; 19 | // tcnn 20 | std::shared_ptr p_decoder_tcnn_, p_color_decoder_tcnn_; 21 | std::shared_ptr p_dir_encoder_tcnn_; 22 | torch::Tensor dir_mask; 23 | 24 | torch::Tensor partition_pos_W_M_; // pos: Map(xyz) to World; 25 | torch::Tensor partition_map_index_; // index: Map to World; [3] 26 | 27 | EncodingMap *new_map(const torch::Tensor &_pos_W_M); 28 | 29 | void unfreeze_net(); 30 | 31 | void freeze_net(); 32 | 33 | torch::Tensor get_feat(const torch::Tensor &xyz, const int &encoding_type = 0, 34 | const bool &normalized = false); 35 | 36 | std::vector get_sdf(const torch::Tensor &xyz, 37 | const bool &normalized = false); 38 | std::vector get_gradient(const torch::Tensor &_xyz, 39 | const float &delta = 0.01, 40 | torch::Tensor _sdf = torch::Tensor(), 41 | bool _heissian = false); 42 | std::vector get_color(const torch::Tensor &xyz, 43 | const torch::Tensor &dir); 44 | 45 | #ifdef ENABLE_ROS 46 | static void pub_mesh(const ros::Publisher &_mesh_pub, 47 | const ros::Publisher &_mesh_color_pub, 48 | const torch::Tensor &vertice, const torch::Tensor &face, 49 | const torch::Tensor &color, 50 | const std_msgs::Header &_header, 51 | const std::string &_uuid); 52 | 53 | void meshing_(ros::Publisher &mesh_pub, ros::Publisher &mesh_color_pub, 54 | std_msgs::Header &header, float _res, bool _save = false, 55 | 56 | const std::string &uuid = "mesh_map"); 57 | #endif 58 | 59 | void meshing_(float _res, bool _save); 60 | 61 | DepthSamples sample(const DepthSamples &_samples, int voxel_sample_num = 1); 62 | 63 | DepthSamples filter_sample(DepthSamples &_samples); 64 | void freeze_decoder(); 65 | }; -------------------------------------------------------------------------------- /include/neural_net/sub_map.cpp: -------------------------------------------------------------------------------- 1 | #include "sub_map.h" 2 | #include "kaolin_wisp_cpp/spc_ops/spc_ops.h" 3 | #include "params/params.h" 4 | 5 | using namespace std; 6 | 7 | SubMap::SubMap(const torch::Tensor &_pos_W_M, float _x_min, float _x_max, 8 | float _y_min, float _y_max, float _z_min, float _z_max) 9 | : torch::nn::Module(), pos_W_M_(_pos_W_M.view({1, 3}).to(k_device)) { 10 | xyz_min_M_ = torch::tensor({{_x_min, _y_min, _z_min}}, pos_W_M_.options()); 11 | xyz_max_M_ = torch::tensor({{_x_max, _y_max, _z_max}}, pos_W_M_.options()); 12 | 13 | xyz_min_W_ = pos_W_M_ + xyz_min_M_; 14 | xyz_max_W_ = pos_W_M_ + xyz_max_M_; 15 | 16 | /// visualization init 17 | xyz_min_M_margin_ = xyz_min_M_ + 0.5 * k_leaf_size; 18 | xyz_max_M_margin_ = xyz_max_M_ - 0.5 * k_leaf_size; 19 | active_ = true; 20 | } 21 | 22 | void SubMap::update_octree_as(const torch::Tensor &_xyz, 23 | const bool &_is_prior) { 24 | // Make sure the kaolin gets inputs between -1 and 1. 25 | auto normalized_xyz = xyz_to_m1p1_pts(_xyz); 26 | auto qpts = spc_ops::quantize_points(normalized_xyz, k_octree_level); 27 | 28 | qpts = get<0>(torch::unique_dim(qpts.contiguous(), 0)); 29 | if (!_is_prior) { 30 | int res = std::pow(2, k_octree_level); 31 | qpts = spc_ops::points_to_neighbors(qpts).view({-1, 3}).clamp(0, res - 1); 32 | } 33 | p_acc_strcut_ = 34 | std::shared_ptr(from_quantized_points(qpts, k_octree_level)); 35 | } 36 | 37 | torch::Tensor SubMap::get_inrange_mask(const torch::Tensor &_xyz, 38 | float padding) const { 39 | torch::Tensor max_mask = 40 | _xyz < (xyz_max_W_ - padding - 1e-6).to(_xyz.device()); 41 | torch::Tensor min_mask = 42 | _xyz > (xyz_min_W_ + padding + 1e-6).to(_xyz.device()); 43 | torch::Tensor mask = max_mask & min_mask; 44 | return mask.all(1); 45 | } 46 | 47 | void SubMap::get_intersect_point(const torch::Tensor &_points, 48 | const torch::Tensor &_rays, 49 | torch::Tensor &z_nears, torch::Tensor &z_fars, 50 | torch::Tensor &mask_intersect, 51 | float padding) const { 52 | // Force the parallel rays to intersect with plane, and they will be removed 53 | // by sanity check, add small number to avoid judgement 54 | torch::Tensor mask_parallels = _rays == 0; 55 | torch::Tensor tmp_rays = 56 | _rays.index_put({mask_parallels}, _rays.index({mask_parallels}) + 1e-6); 57 | 58 | // [n,3] 59 | torch::Tensor tmp_z_nears = (xyz_min_W_ + padding - _points) / tmp_rays; 60 | torch::Tensor tmp_z_fars = (xyz_max_W_ - padding - _points) / tmp_rays; 61 | 62 | // Make sure near is closer than far 63 | torch::Tensor mask_exchange = tmp_z_nears > tmp_z_fars; 64 | z_nears = 65 | tmp_z_nears.index_put({mask_exchange}, tmp_z_fars.index({mask_exchange})); 66 | z_fars = 67 | tmp_z_fars.index_put({mask_exchange}, tmp_z_nears.index({mask_exchange})); 68 | 69 | z_nears = get<0>(torch::max(z_nears, 1)); 70 | z_fars = get<0>(torch::min(z_fars, 1)); 71 | 72 | // check if intersect 73 | mask_intersect = z_nears < z_fars; 74 | } 75 | 76 | torch::Tensor SubMap::get_valid_mask(const torch::Tensor &xyz, int level) { 77 | // Make sure the kaolin gets inputs between -1 and 1. 78 | auto normalized_xyz = xyz_to_m1p1_pts(xyz); 79 | return p_acc_strcut_->query(normalized_xyz, level).pidx > -1; 80 | } 81 | 82 | torch::Tensor SubMap::xyz_to_m1p1_pts(const torch::Tensor &_xyz) { 83 | return scale_to_m1p1(_xyz - pos_W_M_); 84 | } 85 | torch::Tensor SubMap::m1p1_pts_to_xyz(const torch::Tensor &_m1p1_pts) { 86 | return scale_from_m1p1(_m1p1_pts) + pos_W_M_; 87 | } 88 | torch::Tensor SubMap::scale_from_m1p1(const torch::Tensor &_m1p1_tensor) { 89 | return _m1p1_tensor * 0.5 * k_map_size; 90 | } 91 | torch::Tensor SubMap::scale_to_m1p1(const torch::Tensor &_ntensor) { 92 | return _ntensor * 2 * k_map_size_inv; 93 | } 94 | 95 | torch::Tensor SubMap::xyz_to_zp1_pts(const torch::Tensor &_xyz) { 96 | return 0.5f * xyz_to_m1p1_pts(_xyz) + 0.5f; 97 | } 98 | torch::Tensor SubMap::zp1_pts_to_xyz(const torch::Tensor &_zp1_pts) { 99 | return m1p1_pts_to_xyz(2.0f * _zp1_pts - 1.0f); 100 | } 101 | torch::Tensor SubMap::scale_from_zp1(const torch::Tensor &_zp1_tensor) { 102 | return _zp1_tensor * k_map_size; 103 | } 104 | torch::Tensor SubMap::scale_to_zp1(const torch::Tensor &_ntensor) { 105 | return _ntensor * k_map_size_inv; 106 | } 107 | float SubMap::scale_to_zp1(const float &_value) { 108 | return _value * k_map_size_inv; 109 | } -------------------------------------------------------------------------------- /include/neural_net/sub_map.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "kaolin_wisp_cpp/octree_as/octree_as.h" 4 | 5 | struct SubMap : torch::nn::Module { 6 | typedef std::shared_ptr Ptr; 7 | SubMap(const torch::Tensor &_pos_W_M, float _x_min, float _x_max, 8 | float _y_min, float _y_max, float _z_min, float _z_max); 9 | 10 | std::shared_ptr p_acc_strcut_; 11 | 12 | torch::Tensor pos_W_M_; // pos: Map(xyz) to World 13 | torch::Tensor xyz_max_W_, 14 | xyz_min_W_; // relative to World 15 | torch::Tensor xyz_max_M_, xyz_min_M_; // relative to Map 16 | torch::Tensor xyz_max_M_margin_, xyz_min_M_margin_; 17 | 18 | bool active_; 19 | 20 | void update_octree_as(const torch::Tensor &_xyz, 21 | const bool &_is_prior = false); 22 | 23 | torch::Tensor get_inrange_mask(const torch::Tensor &_xyz, 24 | float padding = 0.0) const; 25 | 26 | void get_intersect_point(const torch::Tensor &_points, 27 | const torch::Tensor &_rays, torch::Tensor &z_nears, 28 | torch::Tensor &z_fars, torch::Tensor &mask_intersect, 29 | float padding = 0.0) const; 30 | 31 | torch::Tensor get_valid_mask(const torch::Tensor &xyz, int level = -1); 32 | 33 | torch::Tensor xyz_to_m1p1_pts(const torch::Tensor &_xyz); 34 | torch::Tensor m1p1_pts_to_xyz(const torch::Tensor &_m1p1_pts); 35 | static torch::Tensor scale_from_m1p1(const torch::Tensor &_m1p1_tensor); 36 | static torch::Tensor scale_to_m1p1(const torch::Tensor &_tensor); 37 | 38 | torch::Tensor xyz_to_zp1_pts(const torch::Tensor &_xyz); 39 | torch::Tensor zp1_pts_to_xyz(const torch::Tensor &_zp1_pts); 40 | static torch::Tensor scale_from_zp1(const torch::Tensor &_zp1_tensor); 41 | static torch::Tensor scale_to_zp1(const torch::Tensor &_ntensor); 42 | float scale_to_zp1(const float &_value); 43 | }; -------------------------------------------------------------------------------- /include/optimizer/loss.cpp: -------------------------------------------------------------------------------- 1 | #include "loss.h" 2 | 3 | #include "loss_utils/loss_utils.h" 4 | 5 | namespace loss { 6 | 7 | torch::Tensor rgb_loss(const torch::Tensor &rgb, const torch::Tensor &rgb_gt, 8 | const std::string &name) { 9 | return torch::sqrt((rgb - rgb_gt).square() + 1e-4f).mean(); 10 | // return torch::abs(rgb - rgb_gt).mean(); 11 | } 12 | 13 | torch::Tensor distortion_loss(const torch::Tensor &render_dist, 14 | const std::string &name) { 15 | return render_dist.square().mean(); 16 | } 17 | 18 | torch::Tensor dssim_loss(const torch::Tensor &pred_image, 19 | const torch::Tensor >_image, 20 | const std::string &name) { 21 | auto ssim = loss_utils::ssim(pred_image.permute({2, 0, 1}).unsqueeze(0), 22 | gt_image.permute({2, 0, 1}).unsqueeze(0)); 23 | return 1.0f - ssim; 24 | } 25 | 26 | torch::Tensor sdf_loss(const torch::Tensor &pred_sdf, 27 | const torch::Tensor >_sdf, 28 | const torch::Tensor &pred_isigma) { 29 | auto isigma = pred_isigma.clamp_max(5e2f); 30 | 31 | if (isigma.isnan().any().item()) { 32 | // throw std::runtime_error("inv_bce_sigma.isnan().any()"); 33 | isigma = isigma.nan_to_num(5e2f); 34 | } 35 | 36 | // better avoid nan 37 | torch::Tensor sdf_loss = torch::binary_cross_entropy_with_logits( 38 | (-pred_sdf * isigma).clamp(-15.0f, 15.0f), 39 | torch::sigmoid(-gt_sdf * isigma).clamp(1e-7f, 1 - 1e-7f)); 40 | 41 | return sdf_loss; 42 | } 43 | 44 | torch::Tensor eikonal_loss(const torch::Tensor &grad, const std::string &name) { 45 | return (grad.norm(2, 1) - 1.0f).square().mean(); 46 | } 47 | 48 | torch::Tensor curvate_loss(const torch::Tensor &hessian, 49 | const std::string &name) { 50 | auto curvate_loss = hessian.sum(-1).abs().mean(); 51 | curvate_loss.nan_to_num_(0.0f, 0.0f, 0.0f); 52 | return curvate_loss; 53 | } 54 | 55 | } // namespace loss 56 | -------------------------------------------------------------------------------- /include/optimizer/loss.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace loss { 6 | torch::Tensor rgb_loss(const torch::Tensor &rgb, const torch::Tensor &rgb_gt, 7 | const std::string &name = ""); 8 | 9 | torch::Tensor distortion_loss(const torch::Tensor &render_dist, 10 | const std::string &name = ""); 11 | 12 | torch::Tensor dssim_loss(const torch::Tensor &pred_image, 13 | const torch::Tensor >_image, 14 | const std::string &name = ""); 15 | 16 | torch::Tensor sdf_loss(const torch::Tensor &pred_sdf, 17 | const torch::Tensor >_sdf, 18 | const torch::Tensor &pred_isigma); 19 | 20 | torch::Tensor eikonal_loss(const torch::Tensor &grad, 21 | const std::string &name = ""); 22 | 23 | torch::Tensor curvate_loss(const torch::Tensor &hessian, 24 | const std::string &name = ""); 25 | 26 | torch::Tensor entropy_loss(const torch::Tensor &alpha, 27 | const std::string &name = ""); 28 | 29 | } // namespace loss 30 | -------------------------------------------------------------------------------- /include/optimizer/loss_utils/loss_utils.cpp: -------------------------------------------------------------------------------- 1 | #include "loss_utils.h" 2 | #include "opencv2/opencv.hpp" 3 | 4 | namespace loss_utils { 5 | // 1D Gaussian kernel 6 | torch::Tensor gaussian(int window_size, float sigma) { 7 | torch::Tensor gauss = torch::empty(window_size); 8 | for (int x = 0; x < window_size; ++x) { 9 | gauss[x] = std::exp( 10 | -(std::pow(std::floor(static_cast(x - window_size) / 2.f), 2)) / 11 | (2.f * sigma * sigma)); 12 | } 13 | return gauss / gauss.sum(); 14 | } 15 | 16 | torch::Tensor create_window(int window_size, int channel) { 17 | auto _1D_window = gaussian(window_size, 1.5).unsqueeze(1); 18 | auto _2D_window = _1D_window.mm(_1D_window.t()).unsqueeze(0).unsqueeze(0); 19 | return _2D_window.expand({channel, 1, window_size, window_size}).contiguous(); 20 | } 21 | 22 | torch::Tensor _ssim(const torch::Tensor &img1, const torch::Tensor &img2, 23 | const torch::Tensor &window, int window_size, int channel, 24 | bool size_average) { 25 | auto mu1 = 26 | torch::nn::functional::conv2d(img1, window, 27 | torch::nn::functional::Conv2dFuncOptions() 28 | .padding(window_size / 2) 29 | .groups(channel)); 30 | auto mu1_sq = mu1.pow(2); 31 | auto sigma1_sq = 32 | torch::nn::functional::conv2d(img1 * img1, window, 33 | torch::nn::functional::Conv2dFuncOptions() 34 | .padding(window_size / 2) 35 | .groups(channel)) - 36 | mu1_sq; 37 | 38 | auto mu2 = 39 | torch::nn::functional::conv2d(img2, window, 40 | torch::nn::functional::Conv2dFuncOptions() 41 | .padding(window_size / 2) 42 | .groups(channel)); 43 | auto mu2_sq = mu2.pow(2); 44 | auto sigma2_sq = 45 | torch::nn::functional::conv2d(img2 * img2, window, 46 | torch::nn::functional::Conv2dFuncOptions() 47 | .padding(window_size / 2) 48 | .groups(channel)) - 49 | mu2_sq; 50 | 51 | auto mu1_mu2 = mu1 * mu2; 52 | auto sigma12 = 53 | torch::nn::functional::conv2d(img1 * img2, window, 54 | torch::nn::functional::Conv2dFuncOptions() 55 | .padding(window_size / 2) 56 | .groups(channel)) - 57 | mu1_mu2; 58 | 59 | static const float C1 = 0.01 * 0.01; 60 | static const float C2 = 0.03 * 0.03; 61 | auto ssim_map = ((2.f * mu1_mu2 + C1) * (2.f * sigma12 + C2)) / 62 | ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)); 63 | 64 | if (size_average) { 65 | return ssim_map.mean(); 66 | } else { 67 | return ssim_map.mean(1).mean(1).mean(1); 68 | } 69 | } 70 | 71 | torch::Tensor ssim(const torch::Tensor &img1, const torch::Tensor &img2, 72 | int window_size, int channel) { 73 | static const float C1 = 0.01 * 0.01; 74 | static const float C2 = 0.03 * 0.03; 75 | const auto window = create_window(window_size, channel) 76 | .to(torch::kFloat32) 77 | .to(torch::kCUDA, true); 78 | // [N, chaneel, H, W], data range: [0, 1] 79 | auto mu1 = 80 | torch::nn::functional::conv2d(img1, window, 81 | torch::nn::functional::Conv2dFuncOptions() 82 | .padding(window_size / 2) 83 | .groups(channel)); 84 | auto mu1_sq = mu1.pow(2); 85 | auto sigma1_sq = 86 | torch::nn::functional::conv2d(img1 * img1, window, 87 | torch::nn::functional::Conv2dFuncOptions() 88 | .padding(window_size / 2) 89 | .groups(channel)) - 90 | mu1_sq; 91 | 92 | auto mu2 = 93 | torch::nn::functional::conv2d(img2, window, 94 | torch::nn::functional::Conv2dFuncOptions() 95 | .padding(window_size / 2) 96 | .groups(channel)); 97 | auto mu2_sq = mu2.pow(2); 98 | auto sigma2_sq = 99 | torch::nn::functional::conv2d(img2 * img2, window, 100 | torch::nn::functional::Conv2dFuncOptions() 101 | .padding(window_size / 2) 102 | .groups(channel)) - 103 | mu2_sq; 104 | 105 | auto mu1_mu2 = mu1 * mu2; 106 | auto sigma12 = 107 | torch::nn::functional::conv2d(img1 * img2, window, 108 | torch::nn::functional::Conv2dFuncOptions() 109 | .padding(window_size / 2) 110 | .groups(channel)) - 111 | mu1_mu2; 112 | 113 | auto ssim_map = ((2.f * mu1_mu2 + C1) * (2.f * sigma12 + C2)) / 114 | ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)); 115 | 116 | return ssim_map.mean(); 117 | } 118 | 119 | float psnr(const torch::Tensor &rendered_img, const torch::Tensor >_img) { 120 | // Check both are in shape of [1, chaneel, H, W] 121 | TORCH_CHECK( 122 | rendered_img.dim() == 4 && gt_img.dim() == 4, 123 | "Both rendered_img and gt_img should be in shape of [1, chaneel, H, W]"); 124 | torch::Tensor squared_diff = 125 | (rendered_img - gt_img.to(rendered_img.device())).square(); 126 | torch::Tensor mse_val = 127 | squared_diff.reshape({rendered_img.size(0), -1}).mean(1, true); 128 | return (20.f * torch::log10(1.0f / mse_val.sqrt())).mean().item(); 129 | } 130 | 131 | } // namespace loss_utils -------------------------------------------------------------------------------- /include/optimizer/loss_utils/loss_utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace loss_utils { 6 | torch::Tensor gaussian(int window_size, float sigma); 7 | 8 | torch::Tensor create_window(int window_size, int channel); 9 | 10 | torch::Tensor _ssim(const torch::Tensor &img1, const torch::Tensor &img2, 11 | const torch::Tensor &window, int window_size, int channel, 12 | bool size_average = true); 13 | 14 | // torch::Tensor ssim(const torch::Tensor &img1, const torch::Tensor &img2, 15 | // int window_size = 11, bool size_average = true); 16 | torch::Tensor ssim(const torch::Tensor &img1, const torch::Tensor &img2, 17 | int window_size = 11, int channel = 3); 18 | float psnr(const torch::Tensor &rendered_img, const torch::Tensor >_img); 19 | 20 | } // namespace loss_utils -------------------------------------------------------------------------------- /include/params/params.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "utils/sensor_utils/sensors.hpp" 4 | #include 5 | #include 6 | 7 | extern bool k_debug; 8 | extern int k_dataset_type; 9 | extern bool k_preload; 10 | 11 | extern torch::Tensor k_map_origin; 12 | extern float k_prefilter; 13 | extern float k_max_time_diff_camera_and_pose, k_max_time_diff_lidar_and_pose; 14 | extern bool k_prob_map_en; 15 | 16 | extern std::filesystem::path k_dataset_path; 17 | 18 | extern int k_ds_pt_num, k_max_pt_num; 19 | 20 | extern std::filesystem::path k_output_path, k_package_path; 21 | 22 | extern torch::Device k_device; 23 | // parameter 24 | extern float k_x_max, k_x_min, k_y_max, k_y_min, k_z_max, k_z_min, k_min_range, 25 | k_max_range; 26 | extern float k_inner_map_size, k_map_size, k_map_size_inv, k_boundary_size; 27 | extern float k_leaf_size, k_leaf_size_inv; 28 | extern int k_octree_level, k_fill_level; 29 | extern int k_map_resolution; 30 | 31 | extern int k_iter_step, k_export_interval, k_export_ckp_interval, 32 | k_surface_sample_num, k_free_sample_num; 33 | extern float k_color_batch_pt_num; 34 | extern int k_batch_num; 35 | extern float k_sample_pts_per_ray; 36 | extern float k_sample_std; 37 | extern bool k_outlier_remove; 38 | extern double k_outlier_dist; 39 | extern int k_outlier_removal_interval; 40 | 41 | extern int k_hidden_dim, k_color_hidden_dim, k_geo_feat_dim; 42 | extern int k_geo_num_layer, k_color_num_layer, k_dir_embedding_degree; 43 | extern int k_n_levels, k_n_features_per_level, k_log2_hashmap_size; 44 | 45 | // abalation parmaeter 46 | extern float k_bce_sigma, k_bce_isigma, k_truncated_dis, k_sphere_trace_thr; 47 | extern float k_sdf_weight, k_eikonal_weight, k_curvate_weight, k_rgb_weight, 48 | k_dist_weight, k_rgb_weight_end, k_rgb_weight_init; 49 | extern float k_res_scale; 50 | 51 | extern int k_trace_iter; 52 | 53 | extern float k_t, k_lr, k_lr_end, k_weight_decay; 54 | 55 | extern sensor::Sensors k_sensor; 56 | extern torch::Tensor k_T_B_S; 57 | 58 | // visualization 59 | extern int k_vis_attribute; 60 | extern int k_vis_frame_step; 61 | extern int k_vis_batch_pt_num, k_batch_ray_num; 62 | extern float k_vis_res, k_export_res; 63 | extern int k_fps; 64 | 65 | extern int k_export_colmap_format, k_export_train_pcl; 66 | extern bool k_llff; 67 | extern bool k_cull_mesh; 68 | 69 | void read_params(const std::filesystem::path &_config_path, 70 | const std::filesystem::path &_data_path = "", 71 | const bool &_new_dir = true); 72 | void read_scene_params(const std::filesystem::path &_scene_config_path); 73 | void read_base_params(const std::filesystem::path &_base_config_path); 74 | void write_pt_params(); 75 | void read_pt_params(); -------------------------------------------------------------------------------- /include/tracer/sphere_trace/sphere_trace.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | namespace sphere_trace { 4 | void sphere_trace_cuda_impl(int64_t num_packs, int64_t num_nugs, 5 | at::Tensor ray_state, at::Tensor r, at::Tensor is, 6 | at::Tensor trans, at::Tensor m, at::Tensor t, 7 | at::Tensor T, at::Tensor z, at::Tensor R, 8 | at::Tensor IS, at::Tensor M, float surface_thr, 9 | at::Tensor start_idxes_in, at::Tensor curr_idxes_in, 10 | at::Tensor depth_io, at::Tensor step_mask); 11 | 12 | at::Tensor sphere_trace_cuda(at::Tensor ray_state, at::Tensor r, at::Tensor is, 13 | at::Tensor trans, at::Tensor m, at::Tensor t, 14 | at::Tensor T, at::Tensor z, at::Tensor R, 15 | at::Tensor IS, at::Tensor M, float surface_thr, 16 | at::Tensor start_idxes_in, 17 | at::Tensor curr_idxes_in, at::Tensor depth_io) { 18 | #ifdef WITH_CUDA 19 | TORCH_CHECK(ray_state.is_contiguous() && r.is_contiguous() && 20 | is.is_contiguous() && trans.is_contiguous() && 21 | m.is_contiguous() && t.is_contiguous() && T.is_contiguous() && 22 | z.is_contiguous() && R.is_contiguous() && IS.is_contiguous() && 23 | M.is_contiguous()); 24 | int64_t num_packs = ray_state.numel(); 25 | int64_t num_nugs = depth_io.size(0); 26 | auto step_mask = at::zeros_like(ray_state).to(at::kBool); 27 | sphere_trace_cuda_impl(num_packs, num_nugs, ray_state, r, is, trans, m, t, T, 28 | z, R, IS, M, surface_thr, start_idxes_in, 29 | curr_idxes_in, depth_io, step_mask); 30 | return step_mask; 31 | #else 32 | AT_ERROR(__func__); 33 | #endif // WITH_CUDA 34 | } 35 | } // namespace sphere_trace 36 | -------------------------------------------------------------------------------- /include/tracer/sphere_trace/sphere_trace.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace sphere_trace { 6 | 7 | at::Tensor sphere_trace_cuda(at::Tensor ray_state, at::Tensor r, at::Tensor is, 8 | at::Tensor trans, at::Tensor m, at::Tensor t, 9 | at::Tensor T, at::Tensor z, at::Tensor R, 10 | at::Tensor IS, at::Tensor M, float surface_thr, 11 | at::Tensor start_idxes_in, 12 | at::Tensor curr_idxes_in, at::Tensor depth_io); 13 | } // namespace sphere_trace 14 | -------------------------------------------------------------------------------- /include/tracer/sphere_trace/sphere_trace_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | namespace sphere_trace { 4 | typedef unsigned int uint; 5 | 6 | __global__ void sphere_trace_cuda_kernel( 7 | int64_t num_packs, int64_t num_nugs, int *ray_state, float *r, float *is, 8 | float *trans, float *m, float *t, float *T, float *z, float *R, float *IS, 9 | float *M, float surface_thr, const int *start_idxes_in, int *curr_idxes_in, 10 | const float2 *depth_io, bool *step_mask) { 11 | uint tidx = blockDim.x * blockIdx.x + threadIdx.x; 12 | if (tidx < num_packs) { 13 | if (ray_state[tidx] == 0) { 14 | return; 15 | } 16 | 17 | step_mask[tidx] = false; 18 | // as sdf is in training, we need more aggressive step 19 | auto tmp_z = T[tidx] - t[tidx]; 20 | // take the scale into acount to slack to avoid unconverged sdf make too 21 | // much revert step and resutls in noisy pixel. 22 | float slack = max(3.0f / is[tidx], surface_thr); 23 | bool do_back_step = tmp_z > (abs(R[tidx]) + abs(r[tidx]) + slack); 24 | if ((ray_state[tidx] == 3) || (m[tidx] == -1.0f)) { 25 | do_back_step = false; 26 | } 27 | if (do_back_step) { 28 | m[tidx] = -1.0f; 29 | } else { 30 | if (ray_state[tidx] == 2) { 31 | // inner object is Isotropic 32 | M[tidx] = -1.0f; 33 | } else if (ray_state[tidx] == 3) { 34 | // Skip gap, need skip a round 35 | m[tidx] = -1.0f; 36 | M[tidx] = 0.0f; 37 | ray_state[tidx] = 1; 38 | } else { 39 | float tmp_M = (R[tidx] - r[tidx]) / tmp_z; 40 | // clip M to [-1, 1] 41 | tmp_M = min(0.99f, max(-1.0f, tmp_M)); // never shrink it to [-1,0] 42 | float step_relaxation = 0.8f; 43 | m[tidx] = step_relaxation * m[tidx] + (1.0f - step_relaxation) * tmp_M; 44 | M[tidx] = m[tidx]; 45 | } 46 | if (M[tidx] < 0.0f) { 47 | float density = -M[tidx] * is[tidx] / (1 + exp(r[tidx] * is[tidx])); 48 | trans[tidx] *= exp(-density * tmp_z); 49 | step_mask[tidx] = true; 50 | z[tidx] = tmp_z; 51 | } 52 | 53 | is[tidx] = IS[tidx]; 54 | r[tidx] = R[tidx]; 55 | t[tidx] = T[tidx]; 56 | if ((ray_state[tidx] == 2) && (trans[tidx] < 1e-3f)) { 57 | ray_state[tidx] = 0; 58 | return; 59 | } 60 | if (r[tidx] <= surface_thr) { 61 | ray_state[tidx] = 2; 62 | } 63 | } 64 | float omega = 2.0f / (1.0f - m[tidx]); 65 | // avoid stuck in 0 66 | float tmp_step = max(abs(r[tidx]) * omega, surface_thr); 67 | T[tidx] = t[tidx] + tmp_step; 68 | 69 | // find next depth bound 70 | uint max_iidx = 71 | (tidx == num_packs - 1) ? num_nugs : start_idxes_in[tidx + 1]; 72 | 73 | int iidx = curr_idxes_in[tidx]; 74 | float query = t[tidx]; 75 | float entry = depth_io[iidx].x; 76 | float exit = depth_io[iidx].y; 77 | while (true) { 78 | if ((query >= entry && query <= exit) || (query < entry)) { 79 | break; 80 | } 81 | iidx++; 82 | if (iidx >= max_iidx) { 83 | // it matter: helps improve strcuture 84 | if (query > exit) { 85 | ray_state[tidx] = 0; 86 | } 87 | return; 88 | } 89 | 90 | entry = depth_io[iidx].x; 91 | if ((entry - exit) > 1e-6) { 92 | T[tidx] = entry + 1e-6f; 93 | ray_state[tidx] = 3; 94 | } 95 | exit = depth_io[iidx].y; 96 | } 97 | curr_idxes_in[tidx] = iidx; 98 | } 99 | } 100 | void sphere_trace_cuda_impl(int64_t num_packs, int64_t num_nugs, 101 | at::Tensor ray_state, at::Tensor r, at::Tensor is, 102 | at::Tensor trans, at::Tensor m, at::Tensor t, 103 | at::Tensor T, at::Tensor z, at::Tensor R, 104 | at::Tensor IS, at::Tensor M, float surface_thr, 105 | at::Tensor start_idxes_in, at::Tensor curr_idxes_in, 106 | at::Tensor depth_io, at::Tensor step_mask) { 107 | sphere_trace_cuda_kernel<<<(num_packs + 1023) / 1024, 1024>>>( 108 | num_packs, num_nugs, ray_state.data_ptr(), r.data_ptr(), 109 | is.data_ptr(), trans.data_ptr(), m.data_ptr(), 110 | t.data_ptr(), T.data_ptr(), z.data_ptr(), 111 | R.data_ptr(), IS.data_ptr(), M.data_ptr(), 112 | surface_thr, start_idxes_in.data_ptr(), 113 | curr_idxes_in.data_ptr(), 114 | reinterpret_cast(depth_io.data_ptr()), 115 | step_mask.data_ptr()); 116 | } 117 | } // namespace sphere_trace 118 | -------------------------------------------------------------------------------- /include/tracer/tracer.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "neural_net/local_map.h" 4 | namespace tracer { 5 | 6 | std::vector 7 | render_ray(const LocalMap::Ptr &_local_map_ptr, const torch::Tensor &ray_o, 8 | const torch::Tensor &ray_d, const int &render_type = 0, 9 | const int &iter_num = 1, const bool training = true); 10 | 11 | std::vector 12 | render_from_pts(const LocalMap::Ptr &_local_map_ptr, torch::Tensor pts, 13 | torch::Tensor ray_d, torch::Tensor ridx, torch::Tensor depth, 14 | torch::Tensor delta, const int &num_rays, 15 | const int &render_type = 0, 16 | torch::Tensor slope = torch::Tensor()); 17 | 18 | void plot_ray_analysis(const LocalMap::Ptr &_local_map_ptr, torch::Tensor pts, 19 | torch::Tensor sdf, torch::Tensor depth, 20 | torch::Tensor delta, torch::Tensor alpha, 21 | std::string save_path = ""); 22 | 23 | std::vector sphere_trace_adaptive_sampling( 24 | const LocalMap::Ptr &_local_map_ptr, torch::Tensor ray_o, 25 | torch::Tensor ray_d, const torch::Tensor &depth_io, torch::Tensor ridx, 26 | const int &iter_num = 100, const float &surface_thr = 0.001, 27 | bool training = true); 28 | } // namespace tracer -------------------------------------------------------------------------------- /include/utils/bin_utils/endian.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2023, ETH Zurich and UNC Chapel Hill. 2 | // All rights reserved. 3 | // 4 | // Redistribution and use in source and binary forms, with or without 5 | // modification, are permitted provided that the following conditions are met: 6 | // 7 | // * Redistributions of source code must retain the above copyright 8 | // notice, this list of conditions and the following disclaimer. 9 | // 10 | // * Redistributions in binary form must reproduce the above copyright 11 | // notice, this list of conditions and the following disclaimer in the 12 | // documentation and/or other materials provided with the distribution. 13 | // 14 | // * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of 15 | // its contributors may be used to endorse or promote products derived 16 | // from this software without specific prior written permission. 17 | // 18 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 21 | // ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE 22 | // LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 23 | // CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 24 | // SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 25 | // INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 26 | // CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 27 | // ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 28 | // POSSIBILITY OF SUCH DAMAGE. 29 | 30 | #pragma once 31 | 32 | #include 33 | #include 34 | #include 35 | 36 | namespace colmap { 37 | 38 | // Reverse the order of each byte. 39 | template 40 | T ReverseBytes(const T& data); 41 | 42 | // Check the order in which bytes are stored in computer memory. 43 | bool IsLittleEndian(); 44 | bool IsBigEndian(); 45 | 46 | // Convert data between endianness and the native format. Note that, for float 47 | // and double types, these functions are only valid if the format is IEEE-754. 48 | // This is the case for pretty much most processors. 49 | template 50 | T LittleEndianToNative(T x); 51 | template 52 | T BigEndianToNative(T x); 53 | template 54 | T NativeToLittleEndian(T x); 55 | template 56 | T NativeToBigEndian(T x); 57 | 58 | // Read data in little endian format for cross-platform support. 59 | template 60 | T ReadBinaryLittleEndian(std::istream* stream); 61 | template 62 | void ReadBinaryLittleEndian(std::istream* stream, std::vector* data); 63 | 64 | // Write data in little endian format for cross-platform support. 65 | template 66 | void WriteBinaryLittleEndian(std::ostream* stream, const T& data); 67 | template 68 | void WriteBinaryLittleEndian(std::ostream* stream, const std::vector& data); 69 | 70 | //////////////////////////////////////////////////////////////////////////////// 71 | // Implementation 72 | //////////////////////////////////////////////////////////////////////////////// 73 | 74 | template 75 | T ReverseBytes(const T& data) { 76 | T data_reversed = data; 77 | std::reverse(reinterpret_cast(&data_reversed), 78 | reinterpret_cast(&data_reversed) + sizeof(T)); 79 | return data_reversed; 80 | } 81 | 82 | inline bool IsLittleEndian() { 83 | #ifdef BOOST_BIG_ENDIAN 84 | return false; 85 | #else 86 | return true; 87 | #endif 88 | } 89 | 90 | inline bool IsBigEndian() { 91 | #ifdef BOOST_BIG_ENDIAN 92 | return true; 93 | #else 94 | return false; 95 | #endif 96 | } 97 | 98 | template 99 | T LittleEndianToNative(const T x) { 100 | if (IsLittleEndian()) { 101 | return x; 102 | } else { 103 | return ReverseBytes(x); 104 | } 105 | } 106 | 107 | template 108 | T BigEndianToNative(const T x) { 109 | if (IsBigEndian()) { 110 | return x; 111 | } else { 112 | return ReverseBytes(x); 113 | } 114 | } 115 | 116 | template 117 | T NativeToLittleEndian(const T x) { 118 | if (IsLittleEndian()) { 119 | return x; 120 | } else { 121 | return ReverseBytes(x); 122 | } 123 | } 124 | 125 | template 126 | T NativeToBigEndian(const T x) { 127 | if (IsBigEndian()) { 128 | return x; 129 | } else { 130 | return ReverseBytes(x); 131 | } 132 | } 133 | 134 | template 135 | T ReadBinaryLittleEndian(std::istream* stream) { 136 | T data_little_endian; 137 | stream->read(reinterpret_cast(&data_little_endian), sizeof(T)); 138 | return LittleEndianToNative(data_little_endian); 139 | } 140 | 141 | template 142 | void ReadBinaryLittleEndian(std::istream* stream, std::vector* data) { 143 | for (size_t i = 0; i < data->size(); ++i) { 144 | (*data)[i] = ReadBinaryLittleEndian(stream); 145 | } 146 | } 147 | 148 | template 149 | void WriteBinaryLittleEndian(std::ostream* stream, const T& data) { 150 | const T data_little_endian = NativeToLittleEndian(data); 151 | stream->write(reinterpret_cast(&data_little_endian), sizeof(T)); 152 | } 153 | 154 | template 155 | void WriteBinaryLittleEndian(std::ostream* stream, const std::vector& data) { 156 | for (const auto& elem : data) { 157 | WriteBinaryLittleEndian(stream, elem); 158 | } 159 | } 160 | 161 | } // namespace colmap 162 | -------------------------------------------------------------------------------- /include/utils/coordinates.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace coords { 6 | /* clang-format off 7 | https://docs.nerf.studio/quickstart/data_conventions.html#coordinate-conventions 8 | 9 | # Camera/view coordinate systems 10 | We use the COLMAP/OpenCV coordinate convention for cameras, where the Y and Z 11 | axes are flipped from ours but the +X axis remains the same. Other codebases may 12 | use the OpenGL/Blender (and original NeRF) convention, where the +X is right, +Y 13 | is up, and +Z is pointing back and away from the camera. -Z is the look-at 14 | direction. 15 | 16 | # World coordinate systems 17 | Our world space is oriented such that the up vector is +Z. The XY plane is 18 | parallel to the ground plane. In the viewer, you’ll notice that red, green, and 19 | blue vectors correspond to X, Y, and Z respectively. 20 | 21 | 22 | - | Camera | World 23 | - OpenCV: right handed, y-down, z-lookat | right handed, z-up 24 | - COLMAP, Replica 25 | - Blender: right handed, y-up, -z-lookat | right handed, z-up 26 | - NerfStudio 27 | - OpenGL: right handed, y-up, -z-lookat | right handed, y-up 28 | - Kitti: right handed, y-down, z-lookat | right handed, -y-up 29 | 30 | Default coordinate systems: OpenCV 31 | clang-format on */ 32 | 33 | enum SystemType { OpenCV = 0, Blender = 1, OpenGL = 2, Kitti = 3 }; 34 | 35 | torch::Tensor opencv_camera_coords(); 36 | 37 | torch::Tensor opencv_camera_transform(); 38 | 39 | torch::Tensor blender_camera_coords(); 40 | 41 | torch::Tensor blender_camera_transform(); 42 | 43 | torch::Tensor opencv_to_opengl_world_transform(); 44 | 45 | torch::Tensor opengl_to_opencv_world_rotation(); 46 | torch::Tensor opencv_to_opengl_camera_rotation(); 47 | 48 | torch::Tensor opengl_to_opencv_world_transform(); 49 | 50 | torch::Tensor rotpos_to_transform34( 51 | const torch::Tensor &rotation, 52 | const torch::Tensor &translation = torch::zeros({3, 1}, torch::kFloat32)); 53 | 54 | torch::Tensor rotpos_to_transform44( 55 | const torch::Tensor &rotation, 56 | const torch::Tensor &translation = torch::zeros({3, 1}, torch::kFloat32)); 57 | 58 | torch::Tensor change_camera_system(const torch::Tensor &pose, 59 | const int &pose_system_type); 60 | 61 | torch::Tensor change_world_system(const torch::Tensor &pose, 62 | const int &_system_type); 63 | 64 | torch::Tensor reset_world_system(const torch::Tensor &pose, 65 | const int &_system_type); 66 | 67 | } // namespace coords -------------------------------------------------------------------------------- /include/utils/loss_utils.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace loss_utils { 6 | torch::Tensor gaussian(int window_size, float sigma) { 7 | torch::Tensor gauss = torch::empty(window_size); 8 | for (int x = 0; x < window_size; x++) { 9 | gauss[x] = exp(-pow(x - window_size / 2, 2) / (2 * pow(sigma, 2))); 10 | } 11 | return gauss / gauss.sum(); 12 | } 13 | 14 | torch::Tensor create_window(int window_size, int channel) { 15 | torch::Tensor _1D_window = gaussian(window_size, 1.5).unsqueeze(1); 16 | torch::Tensor _2D_window = _1D_window.mm(_1D_window.t()) 17 | .to(torch::kFloat32) 18 | .unsqueeze(0) 19 | .unsqueeze(0); 20 | torch::Tensor window = torch::autograd::Variable( 21 | _2D_window.expand({channel, 1, window_size, window_size}).contiguous()); 22 | return window; 23 | } 24 | 25 | torch::Tensor _ssim(const torch::Tensor &img1, const torch::Tensor &img2, 26 | const torch::Tensor &window, int window_size, int channel, 27 | bool size_average = true) { 28 | torch::Tensor mu1 = 29 | torch::conv2d(img1, window, {}, 1, window_size / 2, 1, channel); 30 | torch::Tensor mu2 = 31 | torch::conv2d(img2, window, {}, 1, window_size / 2, 1, channel); 32 | 33 | torch::Tensor mu1_sq = mu1.pow(2); 34 | torch::Tensor mu2_sq = mu2.pow(2); 35 | torch::Tensor mu1_mu2 = mu1 * mu2; 36 | 37 | torch::Tensor sigma1_sq = 38 | torch::conv2d(img1 * img1, window, {}, 1, window_size / 2, 1, channel) - 39 | mu1_sq; 40 | torch::Tensor sigma2_sq = 41 | torch::conv2d(img2 * img2, window, {}, 1, window_size / 2, 1, channel) - 42 | mu2_sq; 43 | torch::Tensor sigma12 = 44 | torch::conv2d(img1 * img2, window, {}, 1, window_size / 2, 1, channel) - 45 | mu1_mu2; 46 | 47 | static float C1 = 0.01 * 0.01; 48 | static float C2 = 0.03 * 0.03; 49 | 50 | torch::Tensor ssim_map = 51 | ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / 52 | ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)); 53 | 54 | if (size_average) { 55 | return ssim_map.mean(); 56 | } else { 57 | return ssim_map.mean(1).mean(1).mean(1); 58 | } 59 | } 60 | 61 | torch::Tensor ssim(const torch::Tensor &img1, const torch::Tensor &img2, 62 | int window_size = 11, bool size_average = true) { 63 | // [N, chaneel, H, W], data range: [0, 1] 64 | int channel = img1.size(-3); 65 | torch::Tensor window = create_window(window_size, channel); 66 | 67 | if (img1.is_cuda()) { 68 | window = window.cuda(); 69 | } 70 | window = window.to(img1.dtype()); 71 | 72 | return _ssim(img1, img2, window, window_size, channel, size_average); 73 | } 74 | 75 | } // namespace loss_utils -------------------------------------------------------------------------------- /include/utils/ray_utils/ray_utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | typedef pcl::PointXYZ PointT; 8 | typedef pcl::PointCloud PointCloudT; 9 | 10 | struct RaySamples { 11 | torch::Tensor origin; // [N, 3] 12 | torch::Tensor direction; // [N, 3] 13 | torch::Tensor zdir_norm; // [N, 1] 14 | 15 | torch::Tensor pred_xyz; // [N,3] 16 | torch::Tensor pred_sdf; // [N,1] 17 | torch::Tensor pred_isigma; // [N,1] 18 | 19 | torch::Tensor ridx; // [N] 20 | 21 | int64_t size(const int &dim = 0) const; 22 | 23 | torch::Device device() const; 24 | 25 | RaySamples 26 | index(const torch::ArrayRef &index) const; 27 | 28 | RaySamples index_select(const int64_t dim, const torch::Tensor &index) const; 29 | 30 | RaySamples slice(int64_t dim = 0, c10::optional start = c10::nullopt, 31 | c10::optional end = c10::nullopt, 32 | int64_t step = 1) const; 33 | 34 | RaySamples cat(const RaySamples &other) const; 35 | 36 | RaySamples to(const torch::TensorOptions &option) const; 37 | RaySamples contiguous() const; 38 | RaySamples clone() const; 39 | }; 40 | 41 | struct DepthSamples : RaySamples { 42 | torch::Tensor xyz; // [N, 3] 43 | torch::Tensor depth; // [N,1] 44 | 45 | torch::Tensor ray_sdf; // [N,1] 46 | 47 | DepthSamples 48 | index(const torch::ArrayRef &index) const; 49 | 50 | DepthSamples index_select(const int64_t dim, 51 | const torch::Tensor &index) const; 52 | 53 | DepthSamples slice(int64_t dim = 0, 54 | c10::optional start = c10::nullopt, 55 | c10::optional end = c10::nullopt, 56 | int64_t step = 1) const; 57 | 58 | DepthSamples cat(const DepthSamples &other) const; 59 | 60 | DepthSamples to(const torch::TensorOptions &option) const; 61 | 62 | DepthSamples contiguous() const; 63 | 64 | DepthSamples clone() const; 65 | }; 66 | 67 | struct ColorSamples : RaySamples { 68 | 69 | torch::Tensor rgb; // [N,3] 70 | 71 | ColorSamples 72 | index(const torch::ArrayRef &index) const; 73 | 74 | ColorSamples index_select(const int64_t dim, 75 | const torch::Tensor &index) const; 76 | 77 | ColorSamples slice(int64_t dim = 0, 78 | c10::optional start = c10::nullopt, 79 | c10::optional end = c10::nullopt, 80 | int64_t step = 1) const; 81 | 82 | ColorSamples cat(const ColorSamples &other) const; 83 | 84 | ColorSamples to(const torch::TensorOptions &option) const; 85 | 86 | ColorSamples contiguous() const; 87 | 88 | ColorSamples clone() const; 89 | }; 90 | 91 | struct RaySamplesVec { 92 | 93 | std::vector origin; // [N, 3] 94 | std::vector direction; // [N, 3] 95 | std::vector zdir_norm; // [N, 1] 96 | 97 | std::vector pred_xyz; // [N,3] 98 | std::vector pred_sdf; // [N,1] 99 | std::vector pred_isigma; // [N,1] 100 | 101 | int64_t size(const int &dim = 0) const; 102 | 103 | void emplace_back(const RaySamples &ray); 104 | 105 | void clear(); 106 | void shrink_to_fit(); 107 | 108 | void release(); 109 | 110 | RaySamples cat() const; 111 | }; 112 | 113 | struct DepthSamplesVec : RaySamplesVec { 114 | std::vector xyz; // [N, 3] 115 | std::vector depth; // [N,1] 116 | 117 | std::vector ray_sdf; // [N,1] 118 | 119 | void emplace_back(const DepthSamples &input); 120 | 121 | void clear(); 122 | 123 | void shrink_to_fit(); 124 | 125 | void release(); 126 | 127 | DepthSamples cat() const; 128 | }; 129 | 130 | struct ColorSamplesVec : RaySamplesVec { 131 | std::vector rgb; // [N,3] 132 | 133 | void emplace_back(const ColorSamples &input); 134 | 135 | void clear(); 136 | 137 | void shrink_to_fit(); 138 | 139 | void release(); 140 | 141 | ColorSamples cat() const; 142 | }; -------------------------------------------------------------------------------- /include/utils/sensor_utils/cameras.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | namespace sensor { 6 | 7 | static torch::Tensor get_image_coords(const int &height, const int &width, 8 | const float &pixel_offset = 0.5) { 9 | // [height width 2], stored as (y, x) coordinates 10 | auto image_coords = 11 | torch::meshgrid({torch::arange(height), torch::arange(width)}, "ij"); 12 | return torch::stack(image_coords, -1) + pixel_offset; 13 | } 14 | 15 | static torch::Tensor get_image_coords_zdir(const int &height, const int &width, 16 | const float &fx, const float &fy, 17 | const float &cx, const float &cy, 18 | torch::Device device = torch::kCPU, 19 | const float &pixel_offset = 0.5) { 20 | // [height width 2] 21 | auto vu = get_image_coords(height, width).to(device); 22 | // [height width] 23 | auto pt_x = (vu.select(2, 1) - cx) / fx; 24 | auto pt_y = (vu.select(2, 0) - cy) / fy; 25 | auto pt_z = torch::ones_like(pt_x); 26 | // [height width 3] 27 | return torch::stack({pt_x, pt_y, pt_z}, -1); 28 | } 29 | 30 | static std::vector 31 | get_image_coords_ndir(const int &height, const int &width, const float &fx, 32 | const float &fy, const float &cx, const float &cy, 33 | torch::Device device = torch::kCPU, 34 | const float &pixel_offset = 0.5) { 35 | // [height width 3] 36 | auto zdir = get_image_coords_zdir(height, width, fx, fy, cx, cy, device); 37 | // [height width 1] 38 | auto zdir_norm = zdir.norm(2, -1, true); 39 | auto ndir = zdir / zdir_norm; 40 | return {ndir, zdir_norm}; 41 | } 42 | 43 | struct Cameras { 44 | float fx; 45 | float fy; 46 | float cx; 47 | float cy; 48 | int width; 49 | int height; 50 | 51 | // please set distortion parameters by set_distortion 52 | int model = 0; // 0: pinhole; 1:equidistant; 53 | float d0 = 0; 54 | float d1 = 0; 55 | float d2 = 0; 56 | float d3 = 0; 57 | float d4 = 0; 58 | 59 | float focal_length_; 60 | 61 | void set_distortion(const float &_d0, const float &_d1 = 0, 62 | const float &_d2 = 0, const float &_d3 = 0, 63 | const float &_d4 = 0) { 64 | d0 = _d0; 65 | d1 = _d1; 66 | d2 = _d2; 67 | d3 = _d3; 68 | d4 = _d4; 69 | 70 | cv::Mat K = 71 | (cv::Mat_(3, 3) << fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0); 72 | 73 | cv::Mat new_K; 74 | 75 | if (model == 0) { 76 | cv::Mat D = (cv::Mat_(1, 5) << d0, d1, d2, d3, d4); 77 | 78 | new_K = cv::getOptimalNewCameraMatrix(K, D, cv::Size(width, height), 0, 79 | cv::Size(width, height), 0, true); 80 | std::cout << "origin K:\n" << K << "\n"; 81 | std::cout << "undistort K:\n" << new_K << "\n"; 82 | cv::initUndistortRectifyMap( 83 | K, // Intrinsics (distorted image) 84 | D, // Distortion (distorted image) 85 | cv::Mat_::eye(3, 3), // Rectification (distorted image) 86 | new_K, // New camera matrix (undistorted image) 87 | cv::Size(width, height), // Image resolution (distorted image) 88 | CV_16SC2, // Map type 89 | undist_map_x_, // Undistortion map for X axis 90 | undist_map_y_ // Undistortion map for Y axis 91 | ); 92 | } else if (model == 1) { 93 | cv::Mat D = (cv::Mat_(1, 4) << d0, d1, d2, d3); 94 | 95 | new_K = cv::getOptimalNewCameraMatrix(K, D, cv::Size(width, height), 0, 96 | cv::Size(width, height), 0, true); 97 | std::cout << "origin K:\n" << K << "\n"; 98 | std::cout << "undistort K:\n" << new_K << "\n"; 99 | cv::fisheye::initUndistortRectifyMap( 100 | K, // Intrinsics (distorted image) 101 | D, // Distortion (distorted image) 102 | cv::Mat_::eye(3, 3), // Rectification (distorted image) 103 | new_K, // New camera matrix (undistorted image) 104 | cv::Size(width, height), // Image resolution (distorted image) 105 | CV_16SC2, // Map type 106 | undist_map_x_, // Undistortion map for X axis 107 | undist_map_y_ // Undistortion map for Y axis 108 | ); 109 | } else { 110 | throw std::runtime_error("Invalid camera model"); 111 | } 112 | distortion_ = true; 113 | 114 | fx = new_K.at(0, 0); 115 | fy = new_K.at(1, 1); 116 | cx = new_K.at(0, 2); 117 | cy = new_K.at(1, 2); 118 | d0 = d1 = d2 = d3 = d4 = 0; 119 | } 120 | 121 | cv::Mat undistort(const cv::Mat &_img) const { 122 | cv::Mat undist_img; 123 | if (distortion_) 124 | cv::remap(_img, undist_img, undist_map_x_, undist_map_y_, 125 | cv::INTER_LINEAR); 126 | else 127 | undist_img = _img.clone(); 128 | return undist_img; 129 | } 130 | 131 | std::vector 132 | generate_rays_from_coords(const torch::Tensor &_coords, 133 | const float &_scale = 1.0) const { 134 | TORCH_CHECK(_coords.dtype() == torch::kFloat32, 135 | "Input coordinates must be float32"); 136 | auto y = _coords.select(-1, 0); 137 | auto x = _coords.select(-1, 1); 138 | // [height width 3] 139 | auto zrays = 140 | torch::stack({(x - _scale * cx) / (_scale * fx), 141 | (y - _scale * cy) / (_scale * fy), torch::ones_like(x)}, 142 | -1); 143 | auto rays_norm = torch::norm(zrays, 2, -1, true); 144 | auto nrays = zrays / rays_norm; 145 | return {nrays, rays_norm}; 146 | } 147 | 148 | std::vector generate_rays(const torch::Tensor &_pose, 149 | const float &_scale = 1.0) { 150 | if (_scale != res_scale_) { 151 | camera_ray_results_ = generate_rays_from_coords( 152 | get_image_coords(_scale * height, _scale * width), _scale); 153 | camera_rays_ = camera_ray_results_[0].view({-1, 3}); 154 | camera_ray_norms_ = camera_ray_results_[1].view({-1, 1}); 155 | res_scale_ = _scale; 156 | } 157 | 158 | auto pos = _pose.slice(1, 3, 4).view({1, 3}); 159 | auto ray_o = pos.expand({camera_rays_.size(0), 3}); 160 | auto rot = _pose.slice(1, 0, 3); 161 | auto ray_d = camera_rays_.to(_pose.device()).matmul(rot.t()); 162 | 163 | return {ray_o, ray_d, camera_ray_norms_}; 164 | } 165 | 166 | private: 167 | float res_scale_ = 0.0; 168 | std::vector camera_ray_results_; 169 | torch::Tensor camera_rays_, camera_ray_norms_; 170 | 171 | bool distortion_ = false; 172 | cv::Mat undist_map_x_, undist_map_y_; 173 | }; 174 | } // namespace sensor -------------------------------------------------------------------------------- /include/utils/sensor_utils/sensors.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "cameras.hpp" 3 | 4 | namespace sensor { 5 | 6 | struct Sensors { 7 | sensor::Cameras camera; 8 | 9 | // extrinsics 10 | torch::Tensor 11 | T_C_L; // [4, 4]; extrinsic param, transformation from lidar to camera 12 | torch::Tensor 13 | T_B_L; // [4, 4]; extrinsic param, transformation from lidar to body 14 | torch::Tensor 15 | T_B_C; // [4, 4]; extrinsic param, transformation from camera to body 16 | }; 17 | } // namespace sensor 18 | -------------------------------------------------------------------------------- /include/utils/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | #include 11 | 12 | #include "ray_utils/ray_utils.h" 13 | 14 | #ifdef ENABLE_ROS 15 | #include 16 | 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include 22 | #endif 23 | 24 | namespace utils { 25 | 26 | std::vector unique(const torch::Tensor &x, const int &dim = 0); 27 | 28 | // get cpu data memory usage 29 | double get_cpu_mem_usage(); 30 | 31 | double get_gpu_mem_usage(); 32 | 33 | torch::Tensor downsample_point(const torch::Tensor &points, int ds_pt_num); 34 | 35 | template 36 | RaySamplesT downsample_sample(const RaySamplesT &samples, int ds_pt_num); 37 | 38 | DepthSamples downsample(const DepthSamples &samples, float pt_size, 39 | float dir_size); 40 | ColorSamples downsample(const ColorSamples &samples, float pt_size, 41 | float dir_size); 42 | 43 | void extract_ray_depth_from_pcl(const torch::Tensor &pointcloud, 44 | torch::Tensor &rays_d, torch::Tensor &depths, 45 | float min_range = 0.01); 46 | 47 | torch::Tensor pointcloud_to_tensor(PointCloudT &pointcloud); 48 | 49 | torch::Tensor vec_eigen_to_tensor( 50 | std::vector> 51 | &vec_eigen); 52 | 53 | std::vector tensor_to_vector(const torch::Tensor &_tensor); 54 | 55 | pcl::PointCloud 56 | tensor_to_pointcloud(const torch::Tensor &_tensor); 57 | pcl::PointCloud 58 | tensor_to_pointcloud(const torch::Tensor &_xyz, const torch::Tensor &_rgb); 59 | 60 | cv::Mat 61 | apply_colormap_to_depth(const cv::Mat &depth, double near = 0.0, 62 | double far = 0.0, 63 | cv::ColormapTypes colormap = cv::COLORMAP_TURBO); 64 | cv::Mat tensor_to_cv_mat(const torch::Tensor &_image, 65 | const float &_depth_scale = 1000.0); 66 | 67 | void pointcloud_to_raydepth(PointCloudT &_pointcloud, int _ds_pt_num, 68 | torch::Tensor &_rays_d, torch::Tensor &_depths, 69 | float min_range = 0.01, 70 | torch::Device device = torch::kCPU); 71 | 72 | template 73 | RaySamplesT sample_batch_pts(const RaySamplesT &_samples, int sdf_batch_ray_num = -1, 74 | int batch_type = 1, int iter = 0); 75 | 76 | void sample_surface_pts(const DepthSamples &_samples, 77 | DepthSamples &surface_samples, 78 | int surface_sample_num = 4, float std = 0.1); 79 | 80 | void sample_strat_pts(const DepthSamples &_samples, torch::Tensor &sample_pts, 81 | torch::Tensor &sample_dirs, torch::Tensor &sample_gts, 82 | int strat_sample_num = 4, float std = 0.1, 83 | float strat_near_ratio = 0.2, 84 | float strat_far_ratio = 0.8); 85 | 86 | torch::Tensor cal_nn_dist(const torch::Tensor &_input_pts, 87 | const torch::Tensor &_target_pts); 88 | 89 | torch::Tensor get_verteices(const torch::Tensor &xyz_index); 90 | 91 | torch::Tensor get_width_verteices(const torch::Tensor &xyz_index, 92 | float width = 1.0); 93 | 94 | torch::Tensor cal_tri_inter_coef(const torch::Tensor &xyz_weight); 95 | 96 | torch::Tensor cal_inter_pair_coef(const torch::Tensor &vertex_sdf, 97 | const torch::Tensor &face_edge_pair_index, 98 | float iso_value = 0.0); 99 | 100 | torch::Tensor normalized_quat_to_rotmat(const torch::Tensor &quat); 101 | 102 | torch::Tensor quat_to_rot(const torch::Tensor &quat, const bool &xyzw = false); 103 | 104 | torch::Tensor rot_to_quat(const torch::Tensor &rotation); 105 | 106 | torch::Tensor positional_encode(const torch::Tensor &xyz); 107 | 108 | torch::Tensor meshgrid_3d(float x_min, float x_max, float y_min, float y_max, 109 | float z_min, float z_max, float resolution, 110 | torch::Device &device); 111 | 112 | #ifdef ENABLE_ROS 113 | sensor_msgs::PointCloud2 tensor_to_pointcloud_msg(const torch::Tensor &_xyz, 114 | const torch::Tensor &_rgb); 115 | 116 | sensor_msgs::Image tensor_to_img_msg(const torch::Tensor &_image); 117 | 118 | visualization_msgs::Marker get_vix_voxel_map(const torch::Tensor &_xyz, 119 | float voxel_size, float r = 1.0, 120 | float g = 1.0, float b = 1.0); 121 | 122 | visualization_msgs::Marker get_vis_shift_map(torch::Tensor _pos_W_M, 123 | float _x_min, float _x_max, 124 | float _y_min, float _y_max, 125 | float _z_min, float _z_max); 126 | #endif 127 | } // namespace utils -------------------------------------------------------------------------------- /launch/rviz.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | neural_mapping 4 | 0.0.0 5 | The Neural Mapping package 6 | 7 | 8 | 9 | 10 | liujianheng 11 | 12 | 13 | 14 | 15 | 16 | TODO 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | catkin 52 | geometry_msgs 53 | sensor_msgs 54 | geometry_msgs 55 | sensor_msgs 56 | 57 | message_generation 58 | message_runtime 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | -------------------------------------------------------------------------------- /pics/pipeline_h.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hku-mars/M2Mapping/6075ad2e39a6f20fdb7db12003f9513a95b983cf/pics/pipeline_h.jpg -------------------------------------------------------------------------------- /src/neural_mapping_node.cpp: -------------------------------------------------------------------------------- 1 | #include "neural_mapping/neural_mapping.h" 2 | #ifdef ENABLE_ROS 3 | #include 4 | #else 5 | #include 6 | #endif 7 | 8 | #define BACKWARD_HAS_DW 1 9 | #include "backward.hpp" 10 | namespace backward { 11 | backward::SignalHandling sh; 12 | } 13 | 14 | /** 15 | * @brief Main entry point for the neural_mapping node 16 | * 17 | * Supports two operating modes: 18 | * - "view": Visualizes a pretrained model 19 | * - "train": Trains a new neural mapping model 20 | */ 21 | int main(int argc, char **argv) { 22 | // Set random seeds for reproducibility 23 | torch::manual_seed(0); 24 | torch::cuda::manual_seed_all(0); 25 | 26 | #ifdef ENABLE_ROS 27 | // Initialize ROS node 28 | ros::init(argc, argv, "neural_mapping"); 29 | ros::NodeHandle nh("neural_mapping"); 30 | ros::console::set_logger_level(ROSCONSOLE_DEFAULT_NAME, 31 | ros::console::levels::Info); 32 | #endif 33 | 34 | // Check for minimum required arguments 35 | if (argc < 2) { 36 | std::cerr << "Usage: neural_mapping_node [args...]\n" 37 | << " Modes: view | train " 38 | "\n"; 39 | return 1; 40 | } 41 | 42 | std::string mode = std::string(argv[1]); 43 | NeuralSLAM::Ptr neural_mapping_ptr; 44 | 45 | try { 46 | if (mode == "view") { 47 | // View mode - load pretrained model 48 | if (argc != 3) { 49 | std::cerr << "Usage: neural_mapping_node view \n"; 50 | return 1; 51 | } 52 | 53 | auto pretrained_path = std::filesystem::path(argv[2]); 54 | auto config_path = pretrained_path / "config/scene/config.yaml"; 55 | 56 | #ifdef ENABLE_ROS 57 | neural_mapping_ptr = std::make_shared(nh, 0, config_path); 58 | #else 59 | neural_mapping_ptr = std::make_shared(0, config_path); 60 | #endif 61 | std::cout << "View mode initialized with model: " << pretrained_path 62 | << std::endl; 63 | } else if (mode == "train") { 64 | // Train mode - create new model 65 | if (argc != 4) { 66 | std::cerr 67 | << "Usage: neural_mapping_node train \n"; 68 | return 1; 69 | } 70 | 71 | std::string config_path = std::string(argv[2]); 72 | std::string data_path = std::string(argv[3]); 73 | 74 | #ifdef ENABLE_ROS 75 | neural_mapping_ptr = 76 | std::make_shared(nh, 1, config_path, data_path); 77 | #else 78 | neural_mapping_ptr = 79 | std::make_shared(1, config_path, data_path); 80 | #endif 81 | std::cout << "Training mode initialized with config: " << config_path 82 | << std::endl; 83 | } else { 84 | std::cerr << "Invalid mode: " << mode << "\n"; 85 | return 1; 86 | } 87 | 88 | #ifdef ENABLE_ROS 89 | ros::spin(); 90 | #else 91 | // Wait indefinitely in non-ROS mode 92 | std::condition_variable cv; 93 | std::mutex cv_m; 94 | std::unique_lock lock(cv_m); 95 | cv.wait(lock); 96 | #endif 97 | } catch (const std::exception &e) { 98 | std::cerr << "Error: " << e.what() << std::endl; 99 | return 1; 100 | } 101 | 102 | return 0; 103 | } -------------------------------------------------------------------------------- /submodules/.gitignore: -------------------------------------------------------------------------------- 1 | libtorch/ 2 | !.gitignore -------------------------------------------------------------------------------- /supplement.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hku-mars/M2Mapping/6075ad2e39a6f20fdb7db12003f9513a95b983cf/supplement.pdf --------------------------------------------------------------------------------