Add a visualizer based on Pangolin

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,14 +217,30 @@ 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
     ${PROJECT_SOURCE_DIR}/rasterizer
     ${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

opensplat.cpp

@@ -7,6 +7,10 @@
 #include "constants.hpp"
 #include <cxxopts.hpp>
 
+#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<float>());
+            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);

visualizer.cpp

@@ -0,0 +1,150 @@
+#include "visualizer.hpp"
+
+#include <algorithm>
+#include <chrono>
+#include <thread>
+
+#include <pangolin/display/display.h>
+
+bool Visualizer::Initialize(int iter_num) {
+  pangolin::CreateWindowAndBind("OpenSplat", 1200, 1000);
+  glEnable(GL_DEPTH_TEST);
+
+  cam_state_ = std::make_unique<pangolin::OpenGlRenderState>(
+      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<pangolin::View>();
+  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<pangolin::View>();
+  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<pangolin::Plotter>(
+      &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<pangolin::Panel>("panel");
+  panel_viewer_->SetBounds(0.0f, 1 / 4.0f, 2 / 3.0f, 1.0f, true);
+  pangolin::DisplayBase().AddDisplay(*panel_viewer_);
+
+  step_ = std::make_unique<pangolin::Var<int>>("panel.step", 0);
+  init_gaussian_num_ =
+      std::make_unique<pangolin::Var<int>>("panel.init gaussian num", 19190);
+  gaussian_num_ = std::make_unique<pangolin::Var<int>>("panel.gaussian num", 0);
+  loss_ = std::make_unique<pangolin::Var<float>>("panel.loss", 0.0f);
+  pause_button_ =
+      std::make_unique<pangolin::Var<bool>>("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<float>(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<float, 2>();
+  auto color_accessor = colors_.accessor<float, 2>();
+
+  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<unsigned char>();
+  imageTexture.Upload(data, GL_RGB, GL_UNSIGNED_BYTE);
+
+  render_viewer_->Activate();
+  glColor3f(1.0, 1.0, 1.0);
+  imageTexture.RenderToViewport(true);
+
+  return true;
+}

visualizer.hpp

@@ -0,0 +1,59 @@
+#ifndef VISUALIZER_H
+#define VISUALIZER_H
+
+#include <memory>
+
+#include <pangolin/display/view.h>
+#include <pangolin/display/widgets.h>
+#include <pangolin/plot/plotter.h>
+#include <pangolin/var/var.h>
+#include <torch/torch.h>
+
+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<pangolin::OpenGlRenderState> cam_state_;
+  std::unique_ptr<pangolin::View> point_cloud_viewer_;
+  std::unique_ptr<pangolin::View> render_viewer_;
+  std::unique_ptr<pangolin::Plotter> loss_viewer_;
+  pangolin::DataLog loss_log_;
+  std::unique_ptr<pangolin::Panel> panel_viewer_;
+  std::unique_ptr<pangolin::Var<int>> step_;
+  std::unique_ptr<pangolin::Var<int>> init_gaussian_num_;
+  std::unique_ptr<pangolin::Var<int>> gaussian_num_;
+  std::unique_ptr<pangolin::Var<float>> loss_;
+  std::unique_ptr<pangolin::Var<bool>> pause_button_;
+
+  torch::Tensor means_;
+  torch::Tensor covariances_;
+  torch::Tensor colors_;
+  torch::Tensor opacities_;
+
+  torch::Tensor rendered_img_;
+  torch::Tensor gt_img_;
+};
+
+#endif