From be182e654f89b0db61bcfbc6c7230b4d3c20febf Mon Sep 17 00:00:00 2001 From: DEAR18 Date: Mon, 16 Dec 2024 23:23:58 +0100 Subject: [PATCH] Add visualization --- CMakeLists.txt | 19 ++++++- opensplat.cpp | 18 ++++++ visualizer.cpp | 150 +++++++++++++++++++++++++++++++++++++++++++++++++ visualizer.hpp | 59 +++++++++++++++++++ 4 files changed, 245 insertions(+), 1 deletion(-) create mode 100644 visualizer.cpp create mode 100644 visualizer.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 57f1473..5c6b672 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -163,6 +163,7 @@ endif() find_package(Torch REQUIRED) find_package(OpenCV HINTS "${OPENCV_DIR}" REQUIRED) +find_package(Pangolin QUIET) if (NOT WIN32 AND NOT APPLE) set(CMAKE_CUDA_COMPILER "${CUDA_TOOLKIT_ROOT_DIR}/bin/nvcc") @@ -216,7 +217,20 @@ endif() add_library(gsplat_cpu rasterizer/gsplat-cpu/gsplat_cpu.cpp) target_include_directories(gsplat_cpu PRIVATE ${TORCH_INCLUDE_DIRS}) -add_executable(opensplat opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp opensfm.cpp input_data.cpp tensor_math.cpp) +set(OPENSPLAT_SRC_FILES opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp +kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp +rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp opensfm.cpp input_data.cpp +tensor_math.cpp) + +if (Pangolin_FOUND) + message(STATUS "Found Pangolin. Including Pangolin-related files.") + list(APPEND OPENSPLAT_SRC_FILES visualizer.cpp) + add_definitions(-DUSE_VISUALIZATION) +else() + message(WARNING "Pangolin not found. Skipping Pangolin-related files.") +endif() + +add_executable(opensplat ${OPENSPLAT_SRC_FILES}) install(TARGETS opensplat DESTINATION bin) set_property(TARGET opensplat PROPERTY CXX_STANDARD 17) target_include_directories(opensplat PRIVATE @@ -224,6 +238,9 @@ target_include_directories(opensplat PRIVATE ${GPU_INCLUDE_DIRS} ) target_link_libraries(opensplat PUBLIC ${STDPPFS_LIBRARY} ${GPU_LIBRARIES} ${GSPLAT_LIBS} ${TORCH_LIBRARIES} ${OpenCV_LIBS}) +if (Pangolin_FOUND) + target_link_libraries(opensplat PUBLIC ${Pangolin_LIBRARIES}) +endif() target_link_libraries(opensplat PRIVATE nlohmann_json::nlohmann_json cxxopts::cxxopts diff --git a/opensplat.cpp b/opensplat.cpp index ee37e6d..4e3fdf7 100644 --- a/opensplat.cpp +++ b/opensplat.cpp @@ -7,6 +7,10 @@ #include "constants.hpp" #include +#ifdef USE_VISUALIZATION +#include "visualizer.hpp" +#endif + namespace fs = std::filesystem; using namespace torch::indexing; @@ -99,6 +103,11 @@ int main(int argc, char *argv[]){ displayStep = 1; } +#ifdef USE_VISUALIZATION + Visualizer visualizer; + visualizer.Initialize(numIters); +#endif + try{ InputData inputData = inputDataFromX(projectRoot); @@ -152,6 +161,15 @@ int main(int argc, char *argv[]){ cv::cvtColor(image, image, cv::COLOR_RGB2BGR); cv::imwrite((fs::path(valRender) / (std::to_string(step) + ".png")).string(), image); } + +#ifdef USE_VISUALIZATION + visualizer.SetInitialGaussianNum(inputData.points.xyz.size(0)); + visualizer.SetLoss(step, mainLoss.item()); + visualizer.SetGaussians(model.means, model.scales, model.featuresDc, + model.opacities); + visualizer.SetImage(rgb, gt); + visualizer.Draw(); +#endif } inputData.saveCameras((fs::path(outputScene).parent_path() / "cameras.json").string(), keepCrs); diff --git a/visualizer.cpp b/visualizer.cpp new file mode 100644 index 0000000..2cd0227 --- /dev/null +++ b/visualizer.cpp @@ -0,0 +1,150 @@ +#include "visualizer.hpp" + +#include +#include +#include + +#include + +bool Visualizer::Initialize(int iter_num) { + pangolin::CreateWindowAndBind("OpenSplat", 1200, 1000); + glEnable(GL_DEPTH_TEST); + + cam_state_ = std::make_unique( + pangolin::ProjectionMatrix(1200, 1000, 420, 420, 600, 500, 0.1f, 1000), + pangolin::ModelViewLookAt(-1, 1, -1, 0, 0, 0, pangolin::AxisNegY)); + + point_cloud_viewer_ = std::make_unique(); + point_cloud_viewer_->SetBounds(1 / 4.0f, 1.0f, 0.0f, 1 / 2.0f, true); + point_cloud_viewer_->SetHandler(new pangolin::Handler3D(*cam_state_)); + pangolin::DisplayBase().AddDisplay(*point_cloud_viewer_); + + render_viewer_ = std::make_unique(); + render_viewer_->SetBounds(1 / 4.0f, 1.0f, 1 / 2.0f, 1.0f, true); + pangolin::DisplayBase().AddDisplay(*render_viewer_); + + loss_log_.SetLabels({"loss"}); + float plotter_range_x = iter_num > 0 ? iter_num : 2000.0f; + float plotter_range_y = 0.3; + loss_viewer_ = std::make_unique( + &loss_log_, 0.0f, plotter_range_x, 0.0f, plotter_range_y, 1.f, 0.01f); + loss_viewer_->SetBounds(0.0f, 1 / 4.0f, 0.0f, 2 / 3.0f, true); + loss_viewer_->Track("$i"); + pangolin::DisplayBase().AddDisplay(*loss_viewer_); + + panel_viewer_ = std::make_unique("panel"); + panel_viewer_->SetBounds(0.0f, 1 / 4.0f, 2 / 3.0f, 1.0f, true); + pangolin::DisplayBase().AddDisplay(*panel_viewer_); + + step_ = std::make_unique>("panel.step", 0); + init_gaussian_num_ = + std::make_unique>("panel.init gaussian num", 19190); + gaussian_num_ = std::make_unique>("panel.gaussian num", 0); + loss_ = std::make_unique>("panel.loss", 0.0f); + pause_button_ = + std::make_unique>("panel.Start/Pause", false, false); + + return true; +} + +void Visualizer::SetLoss(int step, float loss) { + loss_log_.Log(loss); + + if (loss_) { + *loss_ = loss; + } + if (step_) { + *step_ = step; + } +} + +void Visualizer::SetInitialGaussianNum(int num) { + if (init_gaussian_num_) { + *init_gaussian_num_ = num; + } +} + +void Visualizer::SetGaussians(const torch::Tensor& means, + const torch::Tensor& covariances, + const torch::Tensor& colors, + const torch::Tensor& opacities) { + means_ = means.cpu(); + covariances_ = covariances.cpu(); + colors_ = colors.cpu(); + opacities_ = opacities.cpu(); + + if (gaussian_num_) { + *gaussian_num_ = means_.size(0); + } +} + +void Visualizer::SetImage(const torch::Tensor& rendered_img, + const torch::Tensor& gt_img) { + rendered_img_ = (rendered_img.cpu() * 255).to(torch::kUInt8); + gt_img_ = (gt_img.cpu() * 255).to(torch::kUInt8); +} + +void Visualizer::Draw() { + glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT); + + DrawGaussians(); + DrawImage(); + + pangolin::FinishFrame(); + + while (*pause_button_) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + pangolin::WindowInterface* window = pangolin::GetBoundWindow(); + if (window) { + window->ProcessEvents(); + } else { + break; + } + } +} + +bool Visualizer::DrawGaussians() { + if (!point_cloud_viewer_) return false; + + static const double c0 = 0.28209479177387814; + auto sh2rgb = [](float sh) { + return static_cast(std::max(std::min(sh * c0 + 0.5, 1.0), 0.0)); + }; + + point_cloud_viewer_->Activate(*cam_state_); + glColor3f(1.0, 1.0, 1.0); + + int gaussian_num = means_.size(0); + auto mean_accessor = means_.accessor(); + auto color_accessor = colors_.accessor(); + + glBegin(GL_POINTS); + for (int i = 0; i < gaussian_num; ++i) { + glColor3f(sh2rgb(color_accessor[i][0]), sh2rgb(color_accessor[i][1]), + sh2rgb(color_accessor[i][2])); + glVertex3f(mean_accessor[i][0], mean_accessor[i][1], mean_accessor[i][2]); + } + glEnd(); + + return true; +} + +bool Visualizer::DrawImage() { + if (!render_viewer_) return false; + + torch::Tensor concatenated_img; + concatenated_img = torch::cat({rendered_img_, gt_img_}, 0); + + const int width = concatenated_img.size(1); + const int height = concatenated_img.size(0); + pangolin::GlTexture imageTexture(width, height, GL_RGB, false, 0, GL_RGB, + GL_UNSIGNED_BYTE); + unsigned char* data = concatenated_img.data_ptr(); + imageTexture.Upload(data, GL_RGB, GL_UNSIGNED_BYTE); + + render_viewer_->Activate(); + glColor3f(1.0, 1.0, 1.0); + imageTexture.RenderToViewport(true); + + return true; +} diff --git a/visualizer.hpp b/visualizer.hpp new file mode 100644 index 0000000..0214ea4 --- /dev/null +++ b/visualizer.hpp @@ -0,0 +1,59 @@ +#ifndef VISUALIZER_H +#define VISUALIZER_H + +#include + +#include +#include +#include +#include +#include + +class Visualizer { + public: + Visualizer() = default; + ~Visualizer() = default; + + bool Initialize(int iter_num = -1); + + void SetInitialGaussianNum(int num); + + void SetLoss(int step, float loss); + + void SetGaussians(const torch::Tensor& means, + const torch::Tensor& covariances, + const torch::Tensor& colors, + const torch::Tensor& opacities); + + void SetImage(const torch::Tensor& rendered_img, const torch::Tensor& gt_img); + + void Draw(); + + private: + bool DrawGaussians(); + + bool DrawImage(); + + private: + std::unique_ptr cam_state_; + std::unique_ptr point_cloud_viewer_; + std::unique_ptr render_viewer_; + std::unique_ptr loss_viewer_; + pangolin::DataLog loss_log_; + std::unique_ptr panel_viewer_; + std::unique_ptr> step_; + std::unique_ptr> init_gaussian_num_; + std::unique_ptr> gaussian_num_; + std::unique_ptr> loss_; + std::unique_ptr> pause_button_; + + torch::Tensor means_; + torch::Tensor covariances_; + torch::Tensor colors_; + torch::Tensor opacities_; + + torch::Tensor rendered_img_; + torch::Tensor gt_img_; +}; + +#endif