├── .gitmodules
├── AvgMerger.cpp
├── AvgMerger.h
├── CMakeLists.txt
├── ClusterMerger.cpp
├── ClusterMerger.h
├── FlatGenerator.cpp
├── FlatGenerator.h
├── PointbasedKdTreeGenerator.cpp
├── PointbasedKdTreeGenerator.h
├── appearance_filter.cpp
├── appearance_filter.h
├── common.h
├── dependencies
    └── json.hpp
├── ext.cpp
├── half.hpp
├── hierarchy_explicit_loader.cpp
├── hierarchy_explicit_loader.h
├── hierarchy_loader.cpp
├── hierarchy_loader.h
├── hierarchy_writer.cpp
├── hierarchy_writer.h
├── loader.cpp
├── loader.h
├── mainHierarchyCreator.cpp
├── mainHierarchyMerger.cpp
├── rotation_aligner.cpp
├── rotation_aligner.h
├── runtime_maintenance.cu
├── runtime_maintenance.h
├── runtime_switching.cu
├── runtime_switching.h
├── setup.py
├── torch
    ├── torch_interface.cpp
    └── torch_interface.h
├── traversal.cpp
├── traversal.h
├── types.h
├── writer.cpp
└── writer.h


/.gitmodules:
--------------------------------------------------------------------------------
1 | [submodule "dependencies/eigen"]
2 | 	path = dependencies/eigen
3 | 	url = https://gitlab.com/libeigen/eigen.git
4 | 


--------------------------------------------------------------------------------
/AvgMerger.cpp:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #include "AvgMerger.h"
13 | 
14 | void AvgMerger::mergeRec(ExplicitTreeNode* node, const std::vector<Gaussian>& leaf_gaussians)
15 | {
16 | 	Gaussian g;
17 | 	g.position = Eigen::Vector3f::Zero();
18 | 	g.rotation = Eigen::Vector4f::Zero();
19 | 	g.opacity = 0;
20 | 	g.scale = Eigen::Vector3f::Zero();
21 | 	g.shs = SHs::Zero();
22 | 
23 | 	float div = node->children.size() + node->leaf_indices.size();
24 | 
25 | 	auto avgmerge = [&div](Gaussian& g, const Gaussian& x) {
26 | 		g.position += x.position / div;
27 | 		g.opacity += x.opacity / div;
28 | 		g.scale += x.scale;
29 | 		g.rotation += x.rotation / div;
30 | 		g.shs += x.shs / div;
31 | 	};
32 | 
33 | 	for (auto& child : node->children)
34 | 	{
35 | 		mergeRec(child, leaf_gaussians);
36 | 		avgmerge(g, child->merged[0]);
37 | 
38 | 		for (auto& child_leaf : child->leaf_indices)
39 | 			avgmerge(g, leaf_gaussians[child_leaf]);
40 | 	}
41 | 
42 | 	g.rotation = g.rotation.normalized();
43 | 
44 | 	node->merged.push_back(g);
45 | }
46 | 
47 | void AvgMerger::merge(ExplicitTreeNode* root, const std::vector<Gaussian>& leaf_gaussians)
48 | {
49 | 	mergeRec(root, leaf_gaussians);
50 | }


--------------------------------------------------------------------------------
/AvgMerger.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "common.h"
15 | 
16 | class AvgMerger
17 | {
18 | private:
19 | 	void mergeRec(ExplicitTreeNode* node, const std::vector<Gaussian>& leaf_gaussians);
20 | public:
21 | 	void merge(ExplicitTreeNode* root, const std::vector<Gaussian>& gaussians);
22 | };


--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | cmake_minimum_required (VERSION 3.0)
 2 | project (GaussianHierarchy LANGUAGES CUDA CXX)
 3 | 
 4 | add_library (GaussianHierarchy
 5 |  FlatGenerator.h
 6 |  FlatGenerator.cpp
 7 |  PointbasedKdTreeGenerator.h
 8 |  PointbasedKdTreeGenerator.cpp
 9 |  ClusterMerger.h
10 |  ClusterMerger.cpp
11 |  appearance_filter.h
12 |  appearance_filter.cpp
13 |  AvgMerger.h
14 |  AvgMerger.cpp
15 |  writer.h
16 |  writer.cpp
17 |  common.h 
18 |  loader.h 
19 |  loader.cpp
20 |  hierarchy_loader.h 
21 |  hierarchy_loader.cpp
22 |  hierarchy_explicit_loader.h
23 |  hierarchy_explicit_loader.cpp
24 |  hierarchy_writer.h
25 |  hierarchy_writer.cpp
26 |  traversal.h
27 |  traversal.cpp
28 |  runtime_maintenance.h
29 |  runtime_maintenance.cu
30 |  runtime_switching.h
31 |  runtime_switching.cu
32 |  rotation_aligner.h
33 |  rotation_aligner.cpp
34 |  half.hpp
35 |  types.h)
36 | target_include_directories(GaussianHierarchy PRIVATE dependencies/eigen)
37 | 
38 | set_property(TARGET GaussianHierarchy PROPERTY CXX_STANDARD 17)
39 | set_target_properties(GaussianHierarchy PROPERTIES CUDA_ARCHITECTURES "70;75;86")
40 | set(CMAKE_CUDA_STANDARD 17)
41 | 
42 | target_include_directories(GaussianHierarchy
43 |                            PUBLIC
44 |                             $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
45 |                             $<INSTALL_INTERFACE:include>
46 |                            )
47 | 						   
48 | install(TARGETS
49 |         GaussianHierarchy
50 |     EXPORT GaussianHierarchyTargets
51 |     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
52 | )
53 | install(FILES runtime_maintenance.h runtime_switching.h hierarchy_loader.h types.h DESTINATION include)
54 | install(EXPORT GaussianHierarchyTargets
55 |   FILE GaussianHierarchyConfig.cmake
56 |   DESTINATION ${CMAKE_INSTALL_PREFIX}/cmake
57 | )
58 | 
59 | add_executable (GaussianHierarchyCreator
60 |  mainHierarchyCreator.cpp
61 | )
62 | 
63 | add_executable (GaussianHierarchyMerger
64 |  mainHierarchyMerger.cpp
65 | )
66 | 
67 | target_include_directories(GaussianHierarchyCreator PRIVATE dependencies/eigen)
68 | set_property(TARGET GaussianHierarchyCreator PROPERTY CXX_STANDARD 17)
69 | target_link_libraries(GaussianHierarchyCreator PUBLIC GaussianHierarchy)
70 | 
71 | target_include_directories(GaussianHierarchyMerger PRIVATE dependencies/eigen)
72 | set_property(TARGET GaussianHierarchyMerger PROPERTY CXX_STANDARD 17)
73 | target_link_libraries(GaussianHierarchyMerger PUBLIC GaussianHierarchy)


--------------------------------------------------------------------------------
/ClusterMerger.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #include "ClusterMerger.h"
 13 | #include <Eigen/Dense>
 14 | #include <Eigen/Eigenvalues>
 15 | 
 16 | float ellipseSurface(Eigen::Vector3f scale)
 17 | {
 18 | 	return scale[0] * scale[1] +
 19 | 		scale[0] * scale[2] +
 20 | 		scale[1] * scale[2];
 21 | }
 22 | 
 23 | void ClusterMerger::mergeRec(ExplicitTreeNode* node, const std::vector<Gaussian>& leaf_gaussians)
 24 | {
 25 | 	Gaussian clustered;
 26 | 	clustered.position = Eigen::Vector3f::Zero();
 27 | 	clustered.rotation = Eigen::Vector4f::Zero();
 28 | 	clustered.opacity = 0;
 29 | 	clustered.scale = Eigen::Vector3f::Zero();
 30 | 	clustered.shs = SHs::Zero();
 31 | 	clustered.covariance = Cov::Zero();
 32 | 
 33 | 	std::vector<const Gaussian*> toMerge;
 34 | 	for (auto& child : node->children)
 35 | 	{
 36 | 		mergeRec(child, leaf_gaussians);
 37 | 		if(child->merged.size())
 38 | 			toMerge.push_back(&child->merged[0]);
 39 | 
 40 | 		for (auto& child_leaf : child->leaf_indices)
 41 | 			toMerge.push_back(&leaf_gaussians[child_leaf]);
 42 | 	}
 43 | 
 44 | 	if (node->depth == 0)
 45 | 		return;
 46 | 
 47 | 	float weight_sum = 0;
 48 | 	std::vector<float> weights;
 49 | 	for (const Gaussian* g : toMerge)
 50 | 	{
 51 | 		float w = g->opacity * ellipseSurface(g->scale);
 52 | 		weights.push_back(w);
 53 | 		weight_sum += w;
 54 | 	}
 55 | 	for (int i = 0; i < weights.size(); i++)
 56 | 		weights[i] = weights[i] / weight_sum;
 57 | 
 58 | 	for (int i = 0; i < toMerge.size(); i++)
 59 | 	{
 60 | 		const Gaussian* g = toMerge[i];
 61 | 		float a = weights[i];
 62 | 
 63 | 		clustered.position += a * g->position;
 64 | 		clustered.shs += a * g->shs;
 65 | 	}
 66 | 
 67 | 	for (int i = 0; i < toMerge.size(); i++)
 68 | 	{
 69 | 		const Gaussian* g = toMerge[i];
 70 | 		float a = weights[i];
 71 | 
 72 | 		Eigen::Vector3f diff = g->position - clustered.position;
 73 | 
 74 | 		clustered.covariance[0] += a * (g->covariance[0] + diff.x() * diff.x());
 75 | 		clustered.covariance[1] += a * (g->covariance[1] + diff.y() * diff.x());
 76 | 		clustered.covariance[2] += a * (g->covariance[2] + diff.z() * diff.x());
 77 | 		clustered.covariance[3] += a * (g->covariance[3] + diff.y() * diff.y());
 78 | 		clustered.covariance[4] += a * (g->covariance[4] + diff.z() * diff.y());
 79 | 		clustered.covariance[5] += a * (g->covariance[5] + diff.z() * diff.z());
 80 | 	}
 81 | 
 82 | 	Eigen::Matrix3f matrix;
 83 | 	matrix <<
 84 | 		clustered.covariance[0], clustered.covariance[1], clustered.covariance[2],
 85 | 		clustered.covariance[1], clustered.covariance[3], clustered.covariance[4],
 86 | 		clustered.covariance[2], clustered.covariance[4], clustered.covariance[5];
 87 | 
 88 | 	Eigen::SelfAdjointEigenSolver<Eigen::Matrix3f> eigensolver(matrix);
 89 | 	auto eigenvalues = eigensolver.eigenvalues();
 90 | 	auto eigenvectors = eigensolver.eigenvectors();
 91 | 
 92 | 	for (int i = 0; eigenvalues.hasNaN() && i < 5; i++)
 93 | 	{
 94 | 		throw std::runtime_error("Found Nans!");
 95 | 	}
 96 | 
 97 | 	int loops = 0;
 98 | 	while(eigenvalues[0] == 0 || eigenvalues[1] == 0 || eigenvalues[2] == 0)
 99 | 	{
100 | 		matrix(0, 0) += std::max(matrix(0, 0) * 0.0001f, std::numeric_limits<float>::epsilon());
101 | 		matrix(1, 1) += std::max(matrix(1, 1) * 0.0001f, std::numeric_limits<float>::epsilon());
102 | 		matrix(2, 2) += std::max(matrix(2, 2) * 0.0001f, std::numeric_limits<float>::epsilon());
103 | 
104 | 		Eigen::SelfAdjointEigenSolver<Eigen::Matrix3f> eigensolver2(matrix);
105 | 		eigenvalues = eigensolver2.eigenvalues();
106 | 		eigenvectors = eigensolver2.eigenvectors();
107 | 
108 | 		loops++;
109 | 		if (loops % 10 == 0)
110 | 			std::cout << "Working hard..." << std::endl;
111 | 	}
112 | 
113 | 	auto v1 = eigenvectors.col(0);
114 | 	auto v2 = eigenvectors.col(1);
115 | 	auto v3 = eigenvectors.col(2);
116 | 
117 | 	auto test = v1.cross(v2);
118 | 	if (test.dot(v3) < 0)
119 | 		eigenvectors.col(2) *= -1;
120 | 
121 | 	float a = sqrt(abs(eigenvalues.x()));
122 | 	float b = sqrt(abs(eigenvalues.y()));
123 | 	float c = sqrt(abs(eigenvalues.z()));
124 | 
125 | 	auto q = quatFromMatrix(eigenvectors);
126 | 
127 | 	clustered.scale = { a, b, c };
128 | 	clustered.rotation = { q.w(), -q.x(), -q.y(), -q.z() };
129 | 
130 | 	auto q2 = Eigen::Quaternionf(eigenvectors);
131 | 	clustered.rotation = { q2.w(), q2.x(), q2.y(), q2.z() };
132 | 
133 | 	clustered.opacity = weight_sum / (ellipseSurface(clustered.scale));
134 | 
135 | 	node->merged.push_back(clustered);
136 | 
137 | 	Gaussian g;
138 | 	if (node->depth == 0)
139 | 	{
140 | 		g = leaf_gaussians[node->leaf_indices[0]];
141 | 	}
142 | 	else
143 | 	{
144 | 		g = node->merged[0];
145 | 	}
146 | 	float minextent = g.scale.array().minCoeff();
147 | 	float maxextent = g.scale.array().maxCoeff();
148 | 
149 | 	if (node->depth != 0)
150 | 	{
151 | 		for (int i = 0; i < 2; i++)
152 | 		{
153 | 			minextent = std::max(minextent, node->children[i]->bounds.minn.w());
154 | 			maxextent = std::max(maxextent, node->children[i]->bounds.maxx.w());
155 | 		}
156 | 	}
157 | 
158 | 	auto diff = node->bounds.maxx - node->bounds.minn;
159 | 	node->bounds.minn.w() = std::max(std::max(diff.x(), diff.y()), diff.z());
160 | 	node->bounds.maxx.w() = std::max(std::max(diff.x(), diff.y()), diff.z());
161 | }
162 | 
163 | void ClusterMerger::merge(ExplicitTreeNode* root, const std::vector<Gaussian>& leaf_gaussians)
164 | {
165 | 	mergeRec(root, leaf_gaussians);
166 | }


--------------------------------------------------------------------------------
/ClusterMerger.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "common.h"
15 | 
16 | class ClusterMerger
17 | {
18 | private:
19 | 	void mergeRec(ExplicitTreeNode* node, const std::vector<Gaussian>& leaf_gaussians);
20 | public:
21 | 	void merge(ExplicitTreeNode* root, const std::vector<Gaussian>& gaussians);
22 | };


--------------------------------------------------------------------------------
/FlatGenerator.cpp:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #include "FlatGenerator.h"
13 | 
14 | ExplicitTreeNode* FlatGenerator::generate(const std::vector<Gaussian>& gaussians)
15 | {
16 | 	auto node = new ExplicitTreeNode();
17 | 
18 | 	Point minn = { FLT_MAX, FLT_MAX, FLT_MAX };
19 | 	Point maxx = { -FLT_MAX, -FLT_MAX, -FLT_MAX };
20 | 	for (int i = 0; i < gaussians.size(); i++)
21 | 	{
22 | 		const Gaussian& g = gaussians[i];
23 | 		minn = minn.cwiseMin(g.position);
24 | 		maxx = maxx.cwiseMax(g.position);
25 | 		node->leaf_indices.push_back(i);
26 | 	}
27 | 	node->bounds = { minn, maxx };
28 | 	node->depth = 0;
29 | 
30 | 	return std::move(node);
31 | }


--------------------------------------------------------------------------------
/FlatGenerator.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "common.h"
15 | 
16 | class FlatGenerator
17 | {
18 | public:
19 | 	ExplicitTreeNode* generate(const std::vector<Gaussian>& gaussians);
20 | };
21 | 


--------------------------------------------------------------------------------
/PointbasedKdTreeGenerator.cpp:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | 
13 | #include "PointbasedKdTreeGenerator.h"
14 | #include <numeric>
15 | 
16 | ExplicitTreeNode* recKdTree(const std::vector<Gaussian>& gaussians, int* g_indices, int start, int num)
17 | {
18 | 	auto node = new ExplicitTreeNode;
19 | 
20 | 	Point minn = { FLT_MAX, FLT_MAX, FLT_MAX };
21 | 	Point maxx = { -FLT_MAX, -FLT_MAX, -FLT_MAX };
22 | 	for (int i = 0; i < num; i++)
23 | 	{
24 | 		const Gaussian& g = gaussians[g_indices[start + i]];
25 | 		float r = 3.0f * g.scale.maxCoeff();
26 | 		auto gmin = g.position;
27 | 		gmin.array() -= r;
28 | 		auto gmax = g.position;
29 | 		gmax.array() += r;
30 | 		minn = minn.cwiseMin(gmin);
31 | 		maxx = maxx.cwiseMax(gmax);
32 | 	}
33 | 	node->bounds = { minn, maxx };
34 | 	
35 | 	if (num == 1)
36 | 	{
37 | 		node->depth = 0;
38 | 		node->leaf_indices.push_back(g_indices[start]);
39 | 	}
40 | 	else
41 | 	{
42 | 		int axis = 0;
43 | 		float greatest_dist = 0;
44 | 		for (int i = 0; i < 3; i++)
45 | 		{
46 | 			float dist = maxx[i] - minn[i];
47 | 			if (dist > greatest_dist)
48 | 			{
49 | 				greatest_dist = dist;
50 | 				axis = i;
51 | 			}
52 | 		}
53 | 
54 | 		int* range = g_indices + start;
55 | 		int pivot = num / 2 - 1;
56 | 		std::nth_element(range, range + pivot, range + num,
57 | 			[&](const int a, const int b) { return gaussians[a].position[axis] < gaussians[b].position[axis]; }
58 | 		);
59 | 
60 | 		node->children.push_back(recKdTree(gaussians, g_indices, start, pivot + 1));
61 | 		node->children.push_back(recKdTree(gaussians, g_indices, start + pivot + 1, num - (pivot + 1)));
62 | 		node->depth = std::max(node->children[0]->depth, node->children[1]->depth) + 1;
63 | 	}
64 | 
65 | 	return node;
66 | }
67 | 
68 | ExplicitTreeNode* PointbasedKdTreeGenerator::generate(const std::vector<Gaussian>& gaussians)
69 | {
70 | 	std::vector<int> indices(gaussians.size());
71 | 	std::iota(indices.begin(), indices.end(), 0);
72 | 	return recKdTree(gaussians, indices.data(), 0, gaussians.size());
73 | }
74 | 


--------------------------------------------------------------------------------
/PointbasedKdTreeGenerator.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "common.h"
15 | 
16 | class PointbasedKdTreeGenerator
17 | {
18 | public:
19 | 	ExplicitTreeNode* generate(const std::vector<Gaussian>& gaussians);
20 | };
21 | 


--------------------------------------------------------------------------------
/appearance_filter.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #include "appearance_filter.h"
 13 | #include <vector>
 14 | #include <Eigen/Dense>
 15 | #include "common.h"
 16 | #include <iostream>
 17 | #include <fstream>
 18 | #include <map>
 19 | #include "writer.h"
 20 | #include <numeric>
 21 | #include "runtime_switching.h"
 22 | 
 23 | template <typename T>
 24 | T ReadBinaryLittleEndian(std::ifstream* infile)
 25 | {
 26 | 	T val;
 27 | 	infile->read((char*)&val, sizeof(T));
 28 | 	return val;
 29 | }
 30 | 
 31 | template <typename T>
 32 | void ReadBinaryLittleEndian(std::ifstream* infile, std::vector<T>* vals)
 33 | {
 34 | 	infile->read((char*)vals->data(), sizeof(T) * vals->size());
 35 | }
 36 | 
 37 | void AppearanceFilter::init(const char* colmappath)
 38 | {
 39 | 	const std::string basePathName = std::string(colmappath) + "/sparse/0/";
 40 | 	const std::string camerasListing = basePathName + "/cameras.bin";
 41 | 	const std::string imagesListing = basePathName + "/images.bin";
 42 | 
 43 | 	std::ifstream camerasFile(camerasListing, std::ios::binary);
 44 | 	std::ifstream imagesFile(imagesListing, std::ios::binary);
 45 | 
 46 | 	if (camerasFile.good())
 47 | 	{
 48 | 		std::string line;
 49 | 
 50 | 		std::map<size_t, CameraParametersColmap> cameraParameters;
 51 | 		const size_t num_cameras = ReadBinaryLittleEndian<uint64_t>(&camerasFile);
 52 | 
 53 | 		for (size_t i = 0; i < num_cameras; ++i)
 54 | 		{
 55 | 			CameraParametersColmap params;
 56 | 
 57 | 			params.id = ReadBinaryLittleEndian<uint32_t>(&camerasFile);
 58 | 			int model_id = ReadBinaryLittleEndian<int>(&camerasFile);
 59 | 			params.width = ReadBinaryLittleEndian<uint64_t>(&camerasFile);
 60 | 			params.height = ReadBinaryLittleEndian<uint64_t>(&camerasFile);
 61 | 			std::vector<double> Params(4);
 62 | 
 63 | 			ReadBinaryLittleEndian<double>(&camerasFile, &Params);
 64 | 			params.fx = float(Params[0]);
 65 | 			params.fy = float(Params[1]);
 66 | 			params.dx = float(Params[2]);
 67 | 			params.dy = float(Params[3]);
 68 | 			cameraParameters[params.id] = params;
 69 | 		}
 70 | 
 71 | 		const size_t num_reg_images = ReadBinaryLittleEndian<uint64_t>(&imagesFile);
 72 | 		for (size_t i = 0; i < num_reg_images; ++i)
 73 | 		{
 74 | 			unsigned int	cId = ReadBinaryLittleEndian<unsigned int>(&imagesFile);
 75 | 			float			qw = float(ReadBinaryLittleEndian<double>(&imagesFile));
 76 | 			float			qx = float(ReadBinaryLittleEndian<double>(&imagesFile));
 77 | 			float			qy = float(ReadBinaryLittleEndian<double>(&imagesFile));
 78 | 			float			qz = float(ReadBinaryLittleEndian<double>(&imagesFile));
 79 | 			float			tx = float(ReadBinaryLittleEndian<double>(&imagesFile));
 80 | 			float			ty = float(ReadBinaryLittleEndian<double>(&imagesFile));
 81 | 			float			tz = float(ReadBinaryLittleEndian<double>(&imagesFile));
 82 | 			size_t			id = ReadBinaryLittleEndian<unsigned int>(&imagesFile);
 83 | 
 84 | 			char name_char;
 85 | 			do {
 86 | 				imagesFile.read(&name_char, 1);
 87 | 			} while (name_char != '\0');
 88 | 
 89 | 			// ignore the 2d points
 90 | 			const size_t num_points2D = ReadBinaryLittleEndian<uint64_t>(&imagesFile);
 91 | 			for (size_t j = 0; j < num_points2D; ++j) {
 92 | 				ReadBinaryLittleEndian<double>(&imagesFile);
 93 | 				ReadBinaryLittleEndian<double>(&imagesFile);
 94 | 				ReadBinaryLittleEndian<uint64_t>(&imagesFile);
 95 | 			}
 96 | 
 97 | 			if (cameraParameters.find(id) == cameraParameters.end())
 98 | 			{
 99 | 				id = 1;
100 | 			}
101 | 
102 | 			const Eigen::Quaternionf quat(qw, qx, qy, qz);
103 | 			const Eigen::Matrix3f orientation = quat.toRotationMatrix().transpose();
104 | 			Eigen::Vector3f translation(tx, ty, tz);
105 | 
106 | 			Eigen::Vector3f position = -(orientation * translation);
107 | 
108 | 			cameras.push_back({ cameraParameters[id], position });
109 | 		}
110 | 	}
111 | }
112 | 
113 | bool verify_rec(const ExplicitTreeNode* node, const std::map<const ExplicitTreeNode*, int>& tree2base, const std::vector<int>& seen, int parent_seen)
114 | {
115 | 	auto entry = tree2base.find(node);
116 | 	if (entry == tree2base.end())
117 | 		throw std::runtime_error("Looking for an entry that does not exist in tree!");
118 | 
119 | 	int id = entry->second;
120 | 	
121 | 	if (seen[id])
122 | 	{
123 | 		if (parent_seen != -1)
124 | 		{
125 | 			return false;
126 | 		}
127 | 		parent_seen = id;
128 | 	}
129 | 
130 | 	for (auto child : node->children)
131 | 	{
132 | 		if (!verify_rec(child, tree2base, seen, parent_seen))
133 | 			return false;
134 | 	}
135 | 	return true;
136 | }
137 | 
138 | bool bottomRec(ExplicitTreeNode* node, const std::map<const ExplicitTreeNode*, int>& tree2base, const std::vector<int>& seen, std::vector<ExplicitTreeNode*>& bottom)
139 | {
140 | 	auto entry = tree2base.find(node);
141 | 	if (entry == tree2base.end())
142 | 		throw std::runtime_error("Looking for an entry that does not exist in tree!");
143 | 
144 | 	int id = entry->second;
145 | 
146 | 	if (node->children.size() != 0)
147 | 	{
148 | 		bool some = false;
149 | 		bool all = true;
150 | 		for (auto child : node->children)
151 | 		{
152 | 			bool result = bottomRec(child, tree2base, seen, bottom);
153 | 			all &= result;
154 | 			some |= result;
155 | 		}
156 | 
157 | 		if (some && !all)
158 | 			throw std::runtime_error("Incomplete cut!");
159 | 
160 | 		if (all)
161 | 			return true;
162 | 	}
163 | 
164 | 	if (seen[id])
165 | 	{
166 | 		bottom.push_back(node);
167 | 		return true;
168 | 	}
169 | 	else
170 | 		return false;
171 | }
172 | 
173 | 
174 | void andBelowRec(ExplicitTreeNode* node, const std::map<const ExplicitTreeNode*, int>& tree2base, const std::vector<int>& marked, std::vector<ExplicitTreeNode*>& bottomandbelow, bool bebelow = false)
175 | {
176 | 	auto entry = tree2base.find(node);
177 | 	if (entry == tree2base.end())
178 | 		throw std::runtime_error("Looking for an entry that does not exist in tree!");
179 | 
180 | 	int id = entry->second;
181 | 
182 | 	if (marked[id])
183 | 		bebelow = true;
184 | 
185 | 	if (bebelow)
186 | 		bottomandbelow.push_back(node);
187 | 
188 | 	if (node->children.size() != 0)
189 | 	{
190 | 		for (auto child : node->children)
191 | 		{
192 | 			andBelowRec(child, tree2base, marked, bottomandbelow, bebelow);
193 | 		}
194 | 	}
195 | }
196 | 
197 | void recCollapse(
198 | 	ExplicitTreeNode* topnode,
199 | 	ExplicitTreeNode* node,
200 | 	std::map<const ExplicitTreeNode*, int>& tree2base,
201 | 	std::vector<int>& marked)
202 | {
203 | 	auto it = tree2base.find(node);
204 | 	if (it == tree2base.end())
205 | 		throw std::runtime_error("Looking for an entry that does not exist in tree!");
206 | 
207 | 	if (marked[it->second] || node->depth == 0)
208 | 	{		
209 | 		topnode->children.push_back(node);
210 | 	}
211 | 	else
212 | 	{
213 | 		for (auto& child : node->children)
214 | 			recCollapse(topnode, child, tree2base, marked);
215 | 	}
216 | }
217 | 
218 | void collapseUnused(
219 | 	std::vector<ExplicitTreeNode*>& bottom,
220 | 	std::map<const ExplicitTreeNode*, int>& tree2base,
221 | 	std::vector<int>& marked
222 | 	)
223 | {
224 | 	for (auto& node : bottom)
225 | 	{
226 | 		auto it = tree2base.find(node);
227 | 		if (it == tree2base.end())
228 | 			throw std::runtime_error("Looking for an entry that does not exist in tree!");
229 | 		if (node->depth == 0 || marked[it->second])
230 | 			continue;
231 | 
232 | 		auto backup = node->children;
233 | 		node->children.clear();
234 | 		for (auto& child : backup)
235 | 			recCollapse(node, child, tree2base, marked);
236 | 	}
237 | }
238 | 
239 | void recVisitAndCount(ExplicitTreeNode* node, int& nodes, int& leaves)
240 | {
241 | 	nodes++;
242 | 	if (node->depth == 0)
243 | 		leaves++;
244 | 	for (auto& child : node->children)
245 | 	{
246 | 		recVisitAndCount(child, nodes, leaves);
247 | 	}
248 | }
249 | 
250 | void recRelSizes(ExplicitTreeNode* node, std::vector<float>& sizes, float parentsize)
251 | {
252 | 	auto extent = node->bounds.maxx - node->bounds.minn;
253 | 	Eigen::Vector3f extent3 = { extent.x(), extent.y(), extent.z() };
254 | 	float mysize = extent3.norm();
255 | 
256 | 	if (parentsize != -1)
257 | 	{
258 | 		sizes.push_back(mysize / parentsize);
259 | 	}
260 | 
261 | 	for (auto child : node->children)
262 | 		recRelSizes(child, sizes, mysize);
263 | }
264 | 
265 | void AppearanceFilter::filter(ExplicitTreeNode* root, const std::vector<Gaussian>& gaussians, float orig_limit, float layermultiplier)
266 | {
267 | 	{
268 | 		std::vector<Eigen::Vector3f> positions;
269 | 		std::vector<Eigen::Vector4f> rotations;
270 | 		std::vector<Eigen::Vector3f> log_scales;
271 | 		std::vector<float> opacities;
272 | 		std::vector<SHs> shs;
273 | 		std::vector<Node> basenodes;
274 | 		std::vector<Box> boxes;
275 | 
276 | 		std::map<int, const ExplicitTreeNode*> base2tree;
277 | 
278 | 		Writer::makeHierarchy(
279 | 			gaussians,
280 | 			root,
281 | 			positions,
282 | 			rotations,
283 | 			log_scales,
284 | 			opacities,
285 | 			shs,
286 | 			basenodes,
287 | 			boxes,
288 | 			&base2tree);
289 | 
290 | 		std::map<const ExplicitTreeNode*, int> tree2base;
291 | 		for (auto entry : base2tree)
292 | 		{
293 | 			tree2base.insert(std::make_pair(entry.second, entry.first));
294 | 		}
295 | 
296 | 		std::vector<Point> campositions;
297 | 		for (int i = 0; i < cameras.size(); i++)
298 | 		{
299 | 			campositions.push_back(cameras[i].position);
300 | 		}
301 | 
302 | 		std::vector<int> seen(basenodes.size());
303 | 		std::vector<int> marked(basenodes.size(), 0);
304 | 
305 | 		int nodes_count = 0, leaves_count = 0;
306 | 		recVisitAndCount(root, nodes_count, leaves_count);
307 | 
308 | 		int last_bottom_size = 0;
309 | 
310 | 		float limit = orig_limit;
311 | 
312 | 		while (true)
313 | 		{
314 | 			std::vector<float> v;
315 | 			recRelSizes(root, v, -1);
316 | 
317 | 			double sum = std::accumulate(std::begin(v), std::end(v), 0.0);
318 | 			double mu = sum / v.size();
319 | 			double accum = 0.0;
320 | 			std::for_each(std::begin(v), std::end(v), [&](const double d) {
321 | 				accum += (d - mu) * (d - mu);
322 | 				});
323 | 			double stdev = sqrt(accum / (v.size() - 1));
324 | 
325 | 			std::cout << "Child-to-parent size relation (mean, std):" << mu << " " << stdev << std::endl;
326 | 
327 | 			Switching::markVisibleForAllViewpoints(limit,
328 | 				(int*)basenodes.data(),
329 | 				basenodes.size(),
330 | 				(float*)boxes.data(),
331 | 				(float*)campositions.data(),
332 | 				campositions.size(),
333 | 				seen.data(),
334 | 				0, 0, 0
335 | 			);
336 | 
337 | 			std::vector<ExplicitTreeNode*> bottom;
338 | 			bottomRec(root, tree2base, seen, bottom);
339 | 
340 | 			if (limit > 1)
341 | 				break;
342 | 
343 | 			last_bottom_size = bottom.size();
344 | 
345 | 			collapseUnused(bottom, tree2base, marked);
346 | 
347 | 			for (int i = 0; i < bottom.size(); i++)
348 | 				marked[tree2base[bottom[i]]] = 1;
349 | 
350 | 			nodes_count = 0;
351 | 			leaves_count = 0;
352 | 			recVisitAndCount(root, nodes_count, leaves_count);
353 | 			std::cout << "After collapse: " << nodes_count << " nodes, " << leaves_count << " leaves reachable" << std::endl;
354 | 
355 | 			limit *= layermultiplier;
356 | 		}
357 | 	}
358 | }
359 | 
360 | 
361 | void AppearanceFilter::writeAnchors(const char* filename, ExplicitTreeNode * root, const std::vector<Gaussian>&gaussians, float limit)
362 | {
363 | 	std::cout << "Identifying and writing anchor points" << std::endl;
364 | 
365 | 	std::vector<Eigen::Vector3f> positions;
366 | 	std::vector<Eigen::Vector4f> rotations;
367 | 	std::vector<Eigen::Vector3f> log_scales;
368 | 	std::vector<float> opacities;
369 | 	std::vector<SHs> shs;
370 | 	std::vector<Node> basenodes;
371 | 	std::vector<Box> boxes;
372 | 
373 | 	std::map<int, const ExplicitTreeNode*> base2tree;
374 | 
375 | 	Writer::makeHierarchy(
376 | 		gaussians,
377 | 		root,
378 | 		positions,
379 | 		rotations,
380 | 		log_scales,
381 | 		opacities,
382 | 		shs,
383 | 		basenodes,
384 | 		boxes,
385 | 		&base2tree);
386 | 
387 | 	std::map<const ExplicitTreeNode*, int> tree2base;
388 | 	for (auto entry : base2tree)
389 | 		tree2base.insert(std::make_pair(entry.second, entry.first));
390 | 
391 | 	std::vector<Point> campositions;
392 | 	for (int i = 0; i < cameras.size(); i++)
393 | 		campositions.push_back(cameras[i].position);
394 | 
395 | 	std::vector<int> seen(basenodes.size());
396 | 	std::vector<int> marked(basenodes.size(), 0);
397 | 
398 | 	Switching::markVisibleForAllViewpoints(limit,
399 | 		(int*)basenodes.data(),
400 | 		basenodes.size(),
401 | 		(float*)boxes.data(),
402 | 		(float*)campositions.data(),
403 | 		campositions.size(),
404 | 		seen.data(),
405 | 		0, 0, 0
406 | 	);
407 | 
408 | 	std::vector<ExplicitTreeNode*> bottom;
409 | 	bottomRec(root, tree2base, seen, bottom);
410 | 
411 | 	for (int i = 0; i < bottom.size(); i++)
412 | 		marked[tree2base[bottom[i]]] = 1;
413 | 
414 | 	std::vector<ExplicitTreeNode*> bottomandbelow;
415 | 	andBelowRec(root, tree2base, marked, bottomandbelow);
416 | 
417 | 	std::ofstream anchors(filename, std::ios_base::binary);
418 | 
419 | 	int size = 0;
420 | 	for (int i = 0; i < bottomandbelow.size(); i++)
421 | 	{
422 | 		int index = tree2base[bottomandbelow[i]];
423 | 		size += basenodes[index].count_leafs + basenodes[index].count_merged;
424 | 	}
425 | 	anchors.write((char*)&size, sizeof(int));
426 | 
427 | 	for (int i = 0; i < bottomandbelow.size(); i++)
428 | 	{
429 | 		int index = tree2base[bottomandbelow[i]];
430 | 
431 | 		for (int j = basenodes[index].start; j < basenodes[index].start + basenodes[index].count_leafs + basenodes[index].count_merged; j++)
432 | 		{
433 | 			anchors.write((char*)&j, sizeof(int));
434 | 		}
435 | 	}
436 | }


--------------------------------------------------------------------------------
/appearance_filter.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include <vector>
15 | #include <Eigen/Dense>
16 | #include "common.h"
17 | #include <iostream>
18 | #include <fstream>
19 | 
20 | class AppearanceFilter
21 | {
22 | private:
23 | 
24 | 	struct CameraParametersColmap {
25 | 		size_t id;
26 | 		size_t width;
27 | 		size_t height;
28 | 		float  fx;
29 | 		float  fy;
30 | 		float  dx;
31 | 		float  dy;
32 | 	};
33 | 
34 | 	struct Camera
35 | 	{
36 | 		CameraParametersColmap& params;
37 | 		Eigen::Vector3f position;
38 | 	};
39 | 
40 | 	std::vector<CameraParametersColmap> camera_params;
41 | 	std::vector<Camera> cameras;
42 | 
43 | public:
44 | 
45 | 	void init(const char* colmappath);
46 | 
47 | 	void filter(ExplicitTreeNode* root, const std::vector<Gaussian>& gaussians, float orig_limit, float layermultiplier);
48 | 
49 | 	void writeAnchors(const char* filename, ExplicitTreeNode* root, const std::vector<Gaussian>& gaussians, float limit);
50 | 
51 | };


--------------------------------------------------------------------------------
/common.h:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #pragma once
 13 | 
 14 | #include <Eigen/Dense>
 15 | #include <iostream>
 16 | #include "types.h"
 17 | #include <float.h>
 18 | #include <memory>
 19 | 
 20 | static float sigmoid(const float m1)
 21 | {
 22 | 	return 1.0f / (1.0f + exp(-m1));
 23 | }
 24 | 
 25 | 
 26 | typedef Eigen::Matrix<float, 6, 1> Cov;
 27 | typedef Eigen::Vector3f Point;
 28 | 
 29 | struct RichPoint
 30 | {
 31 | 	Eigen::Vector3f position;
 32 | 	Eigen::Vector3f normal;
 33 | 	float shs[48];
 34 | 	float opacity;
 35 | 	Eigen::Vector3f scale;
 36 | 	float rotation[4];
 37 | };
 38 | 
 39 | struct LessRichPoint
 40 | {
 41 | 	Eigen::Vector3f position;
 42 | 	Eigen::Vector3f normal;
 43 | 	float shs[12];
 44 | 	float opacity;
 45 | 	Eigen::Vector3f scale;
 46 | 	float rotation[4];
 47 | };
 48 | 
 49 | struct Gaussian
 50 | {
 51 | 	Eigen::Vector3f position;
 52 | 	SHs shs;
 53 | 	float opacity;
 54 | 	Eigen::Vector3f scale;
 55 | 	Eigen::Vector4f rotation;
 56 | 	Cov covariance;
 57 | 
 58 | 	Box bounds99(Eigen::Vector3f& minn, Eigen::Vector3f& maxx) const
 59 | 	{
 60 | 		Eigen::Vector3f off(3.f * sqrt(covariance[0]), 3.f * sqrt(covariance[3]), 3.f * sqrt(covariance[5]));
 61 | 		return Box(position - off, position + off);
 62 | 	}
 63 | };
 64 | 
 65 | struct ExplicitTreeNode
 66 | {
 67 | 	int depth = -1;
 68 | 	Box bounds = Box({FLT_MAX, FLT_MAX, FLT_MAX}, {-FLT_MAX, -FLT_MAX, -FLT_MAX});
 69 | 	std::vector<ExplicitTreeNode*> children;
 70 | 	std::vector<int> leaf_indices;
 71 | 	std::vector<Gaussian> merged;
 72 | };
 73 | 
 74 | static Box getBounds(const std::vector<Gaussian>& gaussians)
 75 | {
 76 | 	Eigen::Vector3f minn(FLT_MAX, FLT_MAX, FLT_MAX);
 77 | 	Eigen::Vector3f maxx = -minn;
 78 | 
 79 | 	Eigen::Vector3f gMax, gMin;
 80 | 	for (int i = 0; i < gaussians.size(); i++)
 81 | 	{
 82 | 		gaussians[i].bounds99(gMin, gMax);
 83 | 		maxx = maxx.cwiseMax(gMax);
 84 | 		minn = minn.cwiseMin(gMin);
 85 | 	}
 86 | 
 87 | 	return Box(minn, maxx);
 88 | }
 89 | 
 90 | static Eigen::Matrix3f matrixFromQuat(Eigen::Vector4f rot)
 91 | {
 92 | 	float s = rot.x();
 93 | 	float x = rot.y();
 94 | 	float y = rot.z();
 95 | 	float z = rot.w();
 96 | 
 97 | 	Eigen::Matrix3f R;
 98 | 	R <<
 99 | 		1.f - 2.f * (y * y + z * z), 2.f * (x * y - s * z), 2.f * (x * z + s * y),
100 | 		2.f * (x * y + s * z), 1.f - 2.f * (x * x + z * z), 2.f * (y * z - s * x),
101 | 		2.f * (x * z - s * y), 2.f * (y * z + s * x), 1.f - 2.f * (x * x + y * y);
102 | 
103 | 	return R;
104 | }
105 | 
106 | static void computeCovariance(const Eigen::Vector3f& scale, const Eigen::Vector4f& rot, Cov& covariance, bool debug = false)
107 | {
108 | 	Eigen::Matrix3f L = Eigen::Matrix3f::Identity();
109 | 
110 | 	L(0, 0) = scale.x();
111 | 	L(1, 1) = scale.y();
112 | 	L(2, 2) = scale.z();
113 | 
114 | 	auto R = matrixFromQuat(rot);
115 | 
116 | 	Eigen::Matrix3f T = R * L;
117 | 	Eigen::Matrix3f Tf;
118 | 	Tf <<
119 | 		T(0, 0), T(1, 0), T(2, 0),
120 | 		T(0, 1), T(1, 1), T(2, 1),
121 | 		T(0, 2), T(1, 2), T(2, 2);
122 | 	Eigen::Matrix3f T2 = T * Tf;
123 | 
124 | 	covariance[0] = T2(0, 0);
125 | 	covariance[1] = T2(0, 1);
126 | 	covariance[2] = T2(0, 2);
127 | 	covariance[3] = T2(1, 1);
128 | 	covariance[4] = T2(1, 2);
129 | 	covariance[5] = T2(2, 2);
130 | 
131 | 	if (debug)
132 | 		std::cout << covariance << std::endl;
133 | }
134 | 
135 | template <typename T, int Options>
136 | Eigen::Quaternion<T, 0> quatFromMatrix(const Eigen::Matrix<T, 3, 3, Options, 3, 3>& m) 
137 | {
138 | 	Eigen::Quaternion<T, 0> q;
139 | 	float trace = m(0, 0) + m(1, 1) + m(2, 2) + 1.f;
140 | 	if (trace > 0)
141 | 	{
142 | 		float s = 0.5f / sqrtf(trace);
143 | 		q.x() = (m(1, 2) - m(2, 1)) * s;
144 | 		q.y() = (m(2, 0) - m(0, 2)) * s;
145 | 		q.z() = (m(0, 1) - m(1, 0)) * s;
146 | 		q.w() = 0.25f / s;
147 | 	}
148 | 	else
149 | 	{
150 | 		if ((m(0, 0) > m(1, 1)) && (m(0, 0) > m(2, 2)))
151 | 		{
152 | 			float s = sqrtf(1.f + m(0, 0) - m(1, 1) - m(2, 2)) * 2.f;
153 | 			q.x() = 0.5f / s;
154 | 			q.y() = (m(1, 0) + m(0, 1)) / s;
155 | 			q.z() = (m(2, 0) + m(0, 2)) / s;
156 | 			q.w() = (m(2, 1) + m(1, 2)) / s;
157 | 		}
158 | 		else if (m(1, 1) > m(2, 2))
159 | 		{
160 | 			float s = sqrtf(1.f - m(0, 0) + m(1, 1) - m(2, 2)) * 2.f;
161 | 			q.x() = (m(1, 0) + m(0, 1)) / s;
162 | 			q.y() = 0.5f / s;
163 | 			q.z() = (m(2, 1) + m(1, 2)) / s;
164 | 			q.w() = (m(2, 0) + m(0, 2)) / s;
165 | 		}
166 | 		else
167 | 		{
168 | 			float s = sqrtf(1.f - m(0, 0) - m(1, 1) + m(2, 2)) * 2.f;
169 | 			q.x() = (m(2, 0) + m(0, 2)) / s;
170 | 			q.y() = (m(2, 1) + m(1, 2)) / s;
171 | 			q.z() = 0.5f / s;
172 | 			q.w() = (m(1, 0) + m(0, 1)) / s;
173 | 		}
174 | 	}
175 | 	return q;
176 | }
177 | 


--------------------------------------------------------------------------------
/ext.cpp:
--------------------------------------------------------------------------------
 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 | #include <torch/extension.h>
13 | #include "torch/torch_interface.h"
14 | 
15 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
16 |   m.def("load_hierarchy", &LoadHierarchy);
17 |   m.def("write_hierarchy", &WriteHierarchy);
18 |   m.def("expand_to_target", &ExpandToTarget);
19 |   m.def("expand_to_size", &ExpandToSize);
20 |   m.def("get_interpolation_weights", &GetTsIndexed);
21 | }


--------------------------------------------------------------------------------
/hierarchy_explicit_loader.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | 
 13 | #include "hierarchy_loader.h"
 14 | #include "hierarchy_explicit_loader.h"
 15 | #include <vector>
 16 | #include <Eigen/Dense>
 17 | #include "common.h"
 18 | #include <iostream>
 19 | #include <fstream>
 20 | #include "half.hpp"
 21 | 
 22 | float getWeight(Eigen::Vector3f& pos, int chunk_id,
 23 | 	std::vector<Eigen::Vector3f>& chunk_centers)
 24 | {
 25 | 	float dist_to_current_center = (pos - chunk_centers[chunk_id]).norm();
 26 | 	float min_dist_to_other_center = 1e12f;
 27 | 	for (int other_chunk_id(0); other_chunk_id < chunk_centers.size(); other_chunk_id++)
 28 | 	{
 29 | 		if (other_chunk_id != chunk_id)
 30 | 		{
 31 | 			float dist_to_other_center = (pos - chunk_centers[other_chunk_id]).norm();
 32 | 			if (min_dist_to_other_center > dist_to_other_center)
 33 | 			{
 34 | 				min_dist_to_other_center = dist_to_other_center;
 35 | 			}
 36 | 		}
 37 | 	}
 38 | 	float falloff = 0.05f;
 39 | 	if (dist_to_current_center <= (1.f - falloff) * min_dist_to_other_center)
 40 | 	{
 41 | 		return 1.f;
 42 | 	}
 43 | 	else if (dist_to_current_center > (1.f + falloff) * min_dist_to_other_center)
 44 | 	{
 45 | 		return 0.f;
 46 | 	}
 47 | 	else
 48 | 	{
 49 | 		float a = -1.f / (2 * falloff * min_dist_to_other_center);
 50 | 		float b = (1.f + falloff) / (2 * falloff);
 51 | 		return a * dist_to_current_center + b;
 52 | 	}
 53 | }
 54 | 
 55 | std::vector<ExplicitTreeNode*> buildTreeRec(ExplicitTreeNode* expliciteNode,
 56 | 	std::vector<Gaussian>& gaussians,
 57 | 	int chunk_id,
 58 | 	std::vector<Eigen::Vector3f>& chunk_centers,
 59 | 	Node& node,
 60 | 	int node_id,
 61 | 	std::vector<Eigen::Vector3f>& pos,
 62 | 	std::vector<SHs>& shs,
 63 | 	std::vector<float>& alphas,
 64 | 	std::vector<Eigen::Vector3f>& scales,
 65 | 	std::vector<Eigen::Vector4f>& rot,
 66 | 	std::vector<Node>& nodes,
 67 | 	std::vector<Box>& boxes)
 68 | {
 69 | 	expliciteNode->depth = node.depth;
 70 | 	expliciteNode->bounds = boxes[node_id];
 71 | 	int n_valid_gaussians = 0;
 72 | 	if (node.depth > 0)
 73 | 	{
 74 | 		for (int n(0); n < node.count_merged; n++)
 75 | 		{
 76 | 			float weigth = getWeight(pos[node.start + n], chunk_id, chunk_centers);
 77 | 			if (weigth > 0.f)
 78 | 			{
 79 | 				n_valid_gaussians++;
 80 | 				Gaussian g;
 81 | 				g.position = pos[node.start + n];
 82 | 				g.rotation = rot[node.start + n];
 83 | 				g.opacity = alphas[node.start + n] * weigth;
 84 | 				g.scale = scales[node.start + n].array().exp();
 85 | 				g.shs = shs[node.start + n];
 86 | 
 87 | 				expliciteNode->merged.push_back(g);
 88 | 			}
 89 | 		}
 90 | 	}
 91 | 	else
 92 | 	{
 93 | 		for (int n(0); n < node.count_leafs; n++)
 94 | 		{
 95 | 			float weigth = getWeight(pos[node.start + n], chunk_id, chunk_centers);
 96 | 			if (weigth > 0.f)
 97 | 			{
 98 | 				n_valid_gaussians++;
 99 | 				Gaussian g;
100 | 				g.position = pos[node.start + n];
101 | 				g.rotation = rot[node.start + n];
102 | 				g.opacity = alphas[node.start + n] * weigth;
103 | 				g.scale = scales[node.start + n].array().exp();
104 | 				g.shs = shs[node.start + n];
105 | 
106 | 				expliciteNode->leaf_indices.push_back(gaussians.size());
107 | 				gaussians.push_back(g);
108 | 			}
109 | 		}
110 | 	}
111 | 
112 | 	std::vector<ExplicitTreeNode*> children;
113 | 	for (int i = 0; i < node.count_children; i++)
114 | 	{
115 | 		ExplicitTreeNode* newNode = new ExplicitTreeNode;
116 | 		std::vector<ExplicitTreeNode*> newChildren = buildTreeRec(newNode, gaussians, chunk_id, chunk_centers,
117 | 			nodes[node.start_children + i], node.start_children + i,
118 | 			pos, shs, alphas, scales, rot, nodes, boxes
119 | 		);
120 | 		children.insert(children.end(), newChildren.begin(), newChildren.end());
121 | 	}
122 | 
123 | 	if (n_valid_gaussians > 0)
124 | 	{
125 | 		expliciteNode->children = children;
126 | 		return std::vector<ExplicitTreeNode*>(1, expliciteNode);
127 | 	}
128 | 	else
129 | 	{
130 | 		return children;
131 | 	}
132 | }
133 | 
134 | void HierarchyExplicitLoader::loadExplicit(
135 | 	const char* filename, std::vector<Gaussian>& gaussians, ExplicitTreeNode* root,
136 | 	int chunk_id, std::vector<Eigen::Vector3f>& chunk_centers)
137 | {
138 | 	std::vector<Eigen::Vector3f> pos;
139 | 	std::vector<SHs> shs;
140 | 	std::vector<float> alphas;
141 | 	std::vector<Eigen::Vector3f> scales;
142 | 	std::vector<Eigen::Vector4f> rot;
143 | 	std::vector<Node> nodes;
144 | 	std::vector<Box> boxes;
145 | 	HierarchyLoader::load(filename, pos, shs, alphas, scales, rot, nodes, boxes);
146 | 
147 | 	pos[0] = chunk_centers[chunk_id];
148 | 	buildTreeRec(root, gaussians, chunk_id, chunk_centers, nodes[0], 0, pos, shs, alphas, scales, rot, nodes, boxes);
149 | }
150 | 


--------------------------------------------------------------------------------
/hierarchy_explicit_loader.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include <vector>
15 | #include <Eigen/Dense>
16 | #include "common.h"
17 | #include <iostream>
18 | #include <fstream>
19 | 
20 | class HierarchyExplicitLoader
21 | {
22 | public:
23 | 
24 | 	static void loadExplicit(const char* filename,
25 | 		std::vector<Gaussian>& gaussian, ExplicitTreeNode* root,
26 | 		int chunk_id, std::vector<Eigen::Vector3f>& chunk_centers);
27 | };


--------------------------------------------------------------------------------
/hierarchy_loader.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #include "hierarchy_loader.h"
 13 | #include <vector>
 14 | #include <Eigen/Dense>
 15 | #include "common.h"
 16 | #include <iostream>
 17 | #include <fstream>
 18 | #include "half.hpp"
 19 | 
 20 | struct HalfBox2
 21 | {
 22 | 	half_float::half minn[4];
 23 | 	half_float::half maxx[4];
 24 | };
 25 | 
 26 | void HierarchyLoader::load(const char* filename,
 27 | 	std::vector<Eigen::Vector3f>& pos,
 28 | 	std::vector<SHs>& shs,
 29 | 	std::vector<float>& alphas,
 30 | 	std::vector<Eigen::Vector3f>& scales,
 31 | 	std::vector<Eigen::Vector4f>& rot,
 32 | 	std::vector<Node>& nodes,
 33 | 	std::vector<Box>& boxes)
 34 | {
 35 | 	std::ifstream infile(filename, std::ios_base::binary);
 36 | 
 37 | 	if (!infile.good())
 38 | 		throw std::runtime_error("File not found!");
 39 | 
 40 | 	int P;
 41 | 	infile.read((char*)&P, sizeof(int));
 42 | 
 43 | 	if (P >= 0)
 44 | 	{
 45 | 		pos.resize(P);
 46 | 		shs.resize(P);
 47 | 		alphas.resize(P);
 48 | 		scales.resize(P);
 49 | 		rot.resize(P);
 50 | 
 51 | 		infile.read((char*)pos.data(), P * sizeof(Eigen::Vector3f));
 52 | 		infile.read((char*)rot.data(), P * sizeof(Eigen::Vector4f));
 53 | 		infile.read((char*)scales.data(), P * sizeof(Eigen::Vector3f));
 54 | 		infile.read((char*)alphas.data(), P * sizeof(float));
 55 | 		infile.read((char*)shs.data(), P * sizeof(SHs));
 56 | 
 57 | 		int N;
 58 | 		infile.read((char*)&N, sizeof(int));
 59 | 
 60 | 		nodes.resize(N);
 61 | 		boxes.resize(N);
 62 | 
 63 | 		infile.read((char*)nodes.data(), N * sizeof(Node));
 64 | 		infile.read((char*)boxes.data(), N * sizeof(Box));
 65 | 	}
 66 | 	else
 67 | 	{
 68 | 		size_t allP = -P;
 69 | 
 70 | 		pos.resize(allP);
 71 | 		shs.resize(allP);
 72 | 		alphas.resize(allP);
 73 | 		scales.resize(allP);
 74 | 		rot.resize(allP);
 75 | 
 76 | 		std::vector<half_float::half> half_rotations(allP * 4);
 77 | 		std::vector<half_float::half> half_scales(allP * 3);
 78 | 		std::vector<half_float::half> half_opacities(allP);
 79 | 		std::vector<half_float::half> half_shs(allP * 48);
 80 | 
 81 | 		infile.read((char*)pos.data(), allP * sizeof(Eigen::Vector3f));
 82 | 		infile.read((char*)half_rotations.data(), allP * 4 * sizeof(half_float::half));
 83 | 		infile.read((char*)half_scales.data(), allP * 3 * sizeof(half_float::half));
 84 | 		infile.read((char*)half_opacities.data(), allP * sizeof(half_float::half));
 85 | 		infile.read((char*)half_shs.data(), allP * 48 * sizeof(half_float::half));
 86 | 
 87 | 		for (size_t i = 0; i < allP; i++)
 88 | 		{
 89 | 			for (size_t j = 0; j < 4; j++)
 90 | 				rot[i][j] = half_rotations[i * 4 + j];
 91 | 			for (size_t j = 0; j < 3; j++)
 92 | 				scales[i][j] = half_scales[i * 3 + j];
 93 | 			alphas[i] = half_opacities[i];
 94 | 			for (size_t j = 0; j < 48; j++)
 95 | 				shs[i][j] = half_shs[i * 48 + j];
 96 | 		}
 97 | 
 98 | 		int N;
 99 | 		infile.read((char*)&N, sizeof(int));
100 | 		size_t allN = N;
101 | 
102 | 		nodes.resize(allN);
103 | 		boxes.resize(allN);
104 | 
105 | 		std::vector<HalfNode> half_nodes(allN);
106 | 		std::vector<HalfBox2> half_boxes(allN);
107 | 
108 | 		infile.read((char*)half_nodes.data(), allN * sizeof(HalfNode));
109 | 		infile.read((char*)half_boxes.data(), allN * sizeof(HalfBox2));
110 | 
111 | 		for (int i = 0; i < allN; i++)
112 | 		{
113 | 			nodes[i].parent = half_nodes[i].parent;
114 | 			nodes[i].start = half_nodes[i].start;
115 | 			nodes[i].start_children = half_nodes[i].start_children;
116 | 			nodes[i].depth = half_nodes[i].dccc[0];
117 | 			nodes[i].count_children = half_nodes[i].dccc[1];
118 | 			nodes[i].count_leafs = half_nodes[i].dccc[2];
119 | 			nodes[i].count_merged = half_nodes[i].dccc[3];
120 | 
121 | 			for (int j = 0; j < 4; j++)
122 | 			{
123 | 				boxes[i].minn[j] = half_boxes[i].minn[j];
124 | 				boxes[i].maxx[j] = half_boxes[i].maxx[j];
125 | 			}
126 | 		}
127 | 	}
128 | }
129 | 


--------------------------------------------------------------------------------
/hierarchy_loader.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "types.h"
15 | 
16 | #include <vector>
17 | #include <Eigen/Dense>
18 | #include <iostream>
19 | #include <fstream>
20 | 
21 | class HierarchyLoader
22 | {
23 | public:
24 | 	static void load(const char* filename,
25 | 		std::vector<Eigen::Vector3f>& pos,
26 | 		std::vector<SHs>& shs,
27 | 		std::vector<float>& alphas,
28 | 		std::vector<Eigen::Vector3f>& scales,
29 | 		std::vector<Eigen::Vector4f>& rot,
30 | 		std::vector<Node>& nodes,
31 | 		std::vector<Box>& boxes);
32 | 	
33 | };


--------------------------------------------------------------------------------
/hierarchy_writer.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | 
 13 | #include "hierarchy_writer.h"
 14 | #include <vector>
 15 | #include <Eigen/Dense>
 16 | #include "common.h"
 17 | #include <iostream>
 18 | #include <fstream>
 19 | #include "half.hpp"
 20 | 
 21 | struct HalfBox
 22 | {
 23 | 	half_float::half minn[4];
 24 | 	half_float::half maxx[4];
 25 | };
 26 | 
 27 | void HierarchyWriter::write(const char* filename,
 28 | 	int allG, int allNB,
 29 | 	Eigen::Vector3f* positions,
 30 | 	SHs* shs,
 31 | 	float* opacities,
 32 | 	Eigen::Vector3f* log_scales,
 33 | 	Eigen::Vector4f* rotations,
 34 | 	Node* nodes,
 35 | 	Box* boxes,
 36 | 	bool compressed)
 37 | {
 38 | 	std::ofstream outfile(filename, std::ios_base::binary);
 39 | 
 40 | 	if (!outfile.good())
 41 | 		throw std::runtime_error("File not created!");
 42 | 
 43 | 	size_t allP = allG;
 44 | 	size_t allN = allNB;
 45 | 
 46 | 	if (!compressed)
 47 | 	{
 48 | 		outfile.write((char*)(&allG), sizeof(int));
 49 | 		outfile.write((char*)positions, allP * sizeof(Eigen::Vector3f));
 50 | 		outfile.write((char*)rotations, allP * sizeof(Eigen::Vector4f));
 51 | 		outfile.write((char*)log_scales, allP * sizeof(Eigen::Vector3f));
 52 | 		outfile.write((char*)opacities, allP * sizeof(float));
 53 | 		outfile.write((char*)shs, allP * sizeof(SHs));
 54 | 
 55 | 		outfile.write((char*)(&allNB), sizeof(int));
 56 | 		outfile.write((char*)nodes, allN * sizeof(Node));
 57 | 		outfile.write((char*)boxes, allN * sizeof(Box));
 58 | 	}
 59 | 	else
 60 | 	{
 61 | 		int indi = -allG;
 62 | 		outfile.write((char*)(&indi), sizeof(int));
 63 | 
 64 | 		std::vector<half_float::half> half_rotations(allP * 4);
 65 | 		std::vector<half_float::half> half_scales(allP * 3);
 66 | 		std::vector<half_float::half> half_opacities(allP);
 67 | 		std::vector<half_float::half> half_shs(allP * 48);
 68 | 
 69 | 		int checksum = 0;
 70 | 
 71 | 		for (size_t i = 0; i < allP; i++)
 72 | 		{
 73 | 			for (size_t j = 0; j < 4; j++)
 74 | 				half_rotations[i * 4 + j] = rotations[i][j];
 75 | 			for (size_t j = 0; j < 3; j++)
 76 | 				half_scales[i * 3 + j] = log_scales[i][j];
 77 | 			half_opacities[i] = opacities[i];
 78 | 			for (size_t j = 0; j < 48; j++)
 79 | 				half_shs[i * 48 + j] = shs[i][j];
 80 | 		}
 81 | 		outfile.write((char*)positions, allP * sizeof(Eigen::Vector3f));
 82 | 		outfile.write((char*)half_rotations.data(), allP * 4 * sizeof(half_float::half));
 83 | 		outfile.write((char*)half_scales.data(), allP * 3 * sizeof(half_float::half));
 84 | 		outfile.write((char*)half_opacities.data(), allP * sizeof(half_float::half));
 85 | 		outfile.write((char*)half_shs.data(), allP * 48 * sizeof(half_float::half));
 86 | 
 87 | 		checksum += allP * 4 * 3 + allP * 8 + allP * 6 + allP * 2 + allP * 96;
 88 | 		std::cout << checksum / (1024 * 1024) << " " << checksum / (1000000) << std::endl;
 89 | 
 90 | 		std::vector<HalfNode> half_nodes(allN);
 91 | 		std::vector<HalfBox> half_boxes(allN);
 92 | 
 93 | 		for (int i = 0; i < allN; i++)
 94 | 		{
 95 | 			half_nodes[i].parent = nodes[i].parent;
 96 | 			half_nodes[i].start = nodes[i].start;
 97 | 			half_nodes[i].start_children = nodes[i].start_children;
 98 | 			if (nodes[i].depth > 32000 || nodes[i].count_children > 32000 || nodes[i].count_leafs > 32000 || nodes[i].count_merged > 32000)
 99 | 				throw std::runtime_error("Would lose information!");
100 | 			half_nodes[i].dccc[0] = (short)nodes[i].depth;
101 | 			half_nodes[i].dccc[1] = (short)nodes[i].count_children;
102 | 			half_nodes[i].dccc[2] = (short)nodes[i].count_leafs;
103 | 			half_nodes[i].dccc[3] = (short)nodes[i].count_merged;
104 | 
105 | 			for (int j = 0; j < 4; j++)
106 | 			{
107 | 				half_boxes[i].minn[j] = boxes[i].minn[j];
108 | 				half_boxes[i].maxx[j] = boxes[i].maxx[j];
109 | 			}
110 | 		}
111 | 
112 | 		outfile.write((char*)(&allNB), sizeof(int));
113 | 		outfile.write((char*)half_nodes.data(), allN * sizeof(HalfNode));
114 | 		outfile.write((char*)half_boxes.data(), allN * sizeof(HalfBox));
115 | 
116 | 		checksum += allN * sizeof(HalfNode) + allN * sizeof(HalfBox);
117 | 		std::cout << checksum / (1024 * 1024) << " " << checksum / (1000000) << std::endl;
118 | 	}
119 | }


--------------------------------------------------------------------------------
/hierarchy_writer.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include <vector>
15 | #include <Eigen/Dense>
16 | #include "common.h"
17 | #include <iostream>
18 | #include <fstream>
19 | 
20 | class HierarchyWriter
21 | {
22 | public:
23 | 	void write(const char* filename,
24 | 		int allP, int allN,
25 | 		Eigen::Vector3f* positions,
26 | 		SHs* shs,
27 | 		float* opacities,
28 | 		Eigen::Vector3f* log_scales,
29 | 		Eigen::Vector4f* rotations,
30 | 		Node* nodes,
31 | 		Box* boxes,
32 | 		bool compressed = true);
33 | };


--------------------------------------------------------------------------------
/loader.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | 
 13 | #include "loader.h"
 14 | #include <iostream>
 15 | #include <fstream>
 16 | #include <string>
 17 | 
 18 | void Loader::loadPlyDir(const char* filename, std::vector<Gaussian>& gaussians)
 19 | {
 20 | 	int num_skip;
 21 | 	std::ifstream cfgfile(std::string(filename) + "/pc_info.txt");
 22 | 	cfgfile >> num_skip;
 23 | 	std::cout << "Skipping " << num_skip << std::endl;
 24 | 
 25 | 	std::ifstream infile(std::string(filename) + "/point_cloud.ply", std::ios_base::binary);
 26 | 
 27 | 	if (!infile.good())
 28 | 		throw std::runtime_error("File not found!");
 29 | 
 30 | 	std::string buff;
 31 | 	std::getline(infile, buff);
 32 | 	std::getline(infile, buff);
 33 | 
 34 | 	std::string dummy;
 35 | 	std::getline(infile, buff);
 36 | 	std::stringstream ss(buff);
 37 | 	int count;
 38 | 	ss >> dummy >> dummy >> count;
 39 | 
 40 | 	while (std::getline(infile, buff))
 41 | 		if (buff.compare("end_header") == 0)
 42 | 			break;
 43 | 
 44 | 	gaussians.resize(count);
 45 | 
 46 | 	std::vector<LessRichPoint> points(count);
 47 | 	infile.read((char*)points.data(), count * sizeof(LessRichPoint));
 48 | 
 49 | 	for (int i = 0; i < gaussians.size() - num_skip; i++)
 50 | 	{
 51 | 		Gaussian& g = gaussians[i];
 52 | 		LessRichPoint& p = points[i + num_skip];
 53 | 
 54 | 		g.opacity = sigmoid(p.opacity);
 55 | 		g.position = p.position;
 56 | 		g.rotation = Eigen::Vector4f(p.rotation[0], p.rotation[1], p.rotation[2], p.rotation[3]).normalized();
 57 | 		g.scale = p.scale.array().exp();
 58 | 		for (int j = 0; j < 3; j++)
 59 | 			g.shs[j] = p.shs[j];
 60 | 		for (int j = 1; j < 4; j++)
 61 | 		{
 62 | 			g.shs[(j - 1) + 3]  = p.shs[j * 3 + 0];
 63 | 			g.shs[(j - 1) + 18] = p.shs[j * 3 + 1];
 64 | 			g.shs[(j - 1) + 33] = p.shs[j * 3 + 2];
 65 | 		}
 66 | 		for (int j = 4; j < 16; j++)
 67 | 		{
 68 | 			g.shs[(j - 1) + 3]  = 0;
 69 | 			g.shs[(j - 1) + 18] = 0;
 70 | 			g.shs[(j - 1) + 33] = 0;
 71 | 		}
 72 | 		computeCovariance(g.scale, g.rotation, g.covariance);
 73 | 	}
 74 | }
 75 | 
 76 | void Loader::loadPly(const char* filename, std::vector<Gaussian>& gaussians, int skyboxpoints)
 77 | {
 78 | 	std::ifstream infile(filename, std::ios_base::binary);
 79 | 	std::cout << filename << std::endl;
 80 | 	if (!infile.good())
 81 | 		throw std::runtime_error("File not found!");
 82 | 
 83 | 	std::string buff;
 84 | 	std::getline(infile, buff);
 85 | 	std::getline(infile, buff);
 86 | 
 87 | 	std::string dummy;
 88 | 	std::getline(infile, buff);
 89 | 	std::stringstream ss(buff);
 90 | 	int count;
 91 | 	ss >> dummy >> dummy >> count;
 92 | 
 93 | 	while (std::getline(infile, buff))
 94 | 		if (buff.compare("end_header") == 0)
 95 | 			break;
 96 | 
 97 | 	gaussians.resize(count - skyboxpoints);
 98 | 
 99 | 	std::vector<RichPoint> points(count);
100 | 	infile.read((char*)points.data(), count * sizeof(RichPoint));
101 | 
102 | 	for (int i = 0; i < gaussians.size(); i++)
103 | 	{
104 | 		Gaussian& g = gaussians[i];
105 | 		RichPoint& p = points[skyboxpoints + i];
106 | 
107 | 		g.opacity = sigmoid(p.opacity);
108 | 		g.position = p.position;
109 | 		g.rotation = Eigen::Vector4f(p.rotation[0], p.rotation[1], p.rotation[2], p.rotation[3]).normalized();
110 | 		g.scale = p.scale.array().exp();
111 | 		for (int j = 0; j < 3; j++)
112 | 			g.shs[j] = p.shs[j];
113 | 		for (int j = 1; j < 16; j++)
114 | 		{
115 | 			g.shs[j * 3 + 0] = p.shs[(j - 1) + 3];
116 | 			g.shs[j * 3 + 1] = p.shs[(j - 1) + 18];
117 | 			g.shs[j * 3 + 2] = p.shs[(j - 1) + 33];
118 | 		}
119 | 		computeCovariance(g.scale, g.rotation, g.covariance);
120 | 	}
121 | }
122 | 
123 | void Loader::loadBin(const char* filename, std::vector<Gaussian>& gaussians, int skyboxpoints)
124 | {
125 | 	std::ifstream infile(filename, std::ios_base::binary);
126 | 
127 | 	std::cout << filename << std::endl;
128 | 	if (!infile.good())
129 | 		throw std::runtime_error("File not found!");
130 | 
131 | 	int count;
132 | 	infile.read((char*)&count, sizeof(int));
133 | 
134 | 	std::vector<Eigen::Vector3f> pos(count);
135 | 	std::vector<SHs> shs(count);
136 | 	std::vector<Eigen::Vector3f> scales(count);
137 | 	std::vector<Eigen::Vector4f> rot(count);
138 | 	std::vector<float> alphas(count);
139 | 
140 | 	gaussians.resize(count - skyboxpoints);
141 | 
142 | 	infile.read((char*)pos.data(), sizeof(float) * 3 * count);
143 | 	infile.read((char*)shs.data(), sizeof(SHs) * count);
144 | 	infile.read((char*)alphas.data(), sizeof(float) * count);
145 | 	infile.read((char*)scales.data(), sizeof(float) * 3 * count);
146 | 	infile.read((char*)rot.data(), sizeof(float) * 4 * count);
147 | 
148 | 	for (int i = 0; i < gaussians.size(); i++)
149 | 	{
150 | 		int k = i + skyboxpoints;
151 | 		Gaussian& g = gaussians[i];
152 | 
153 | 		g.opacity = sigmoid(alphas[k]);
154 | 		g.position = pos[k];
155 | 		g.rotation = Eigen::Vector4f(rot[k][0], rot[k][1], rot[k][2], rot[k][3]).normalized();
156 | 		g.scale = scales[k].array().exp();
157 | 		g.shs = shs[k];
158 | 		computeCovariance(g.scale, g.rotation, g.covariance);
159 | 	}
160 | }


--------------------------------------------------------------------------------
/loader.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "common.h"
15 | #include <vector>
16 | 
17 | class Loader
18 | {
19 | public:
20 | 	static void loadPlyDir(const char* filename, std::vector<Gaussian>& gaussian);
21 | 
22 | 	static void loadPly(const char* filename, std::vector<Gaussian>& gaussian, int skyboxpoints = 0);
23 | 
24 | 	static void loadBin(const char* filename, std::vector<Gaussian>& gaussian, int skyboxpoints = 0);
25 | };


--------------------------------------------------------------------------------
/mainHierarchyCreator.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #include "loader.h"
 13 | #include "writer.h"
 14 | #include <vector>
 15 | #include <iostream>
 16 | #include "FlatGenerator.h"
 17 | #include "PointbasedKdTreeGenerator.h"
 18 | #include "AvgMerger.h"
 19 | #include "ClusterMerger.h"
 20 | #include "common.h"
 21 | #include <fstream>
 22 | #include <filesystem>
 23 | #include "appearance_filter.h"
 24 | #include "rotation_aligner.h"
 25 | 
 26 | void recTraverse(ExplicitTreeNode* node, int& zerocount)
 27 | {
 28 | 	if (node->depth == 0)
 29 | 		zerocount++;
 30 | 	if (node->children.size() > 0 && node->depth == 0)
 31 | 		throw std::runtime_error("Leaf nodes should never have children!");
 32 | 
 33 | 	for (auto c : node->children)
 34 | 	{
 35 | 		recTraverse(c, zerocount);
 36 | 	}
 37 | }
 38 | 
 39 | int main(int argc, char* argv[])
 40 | {
 41 | 	if (argc < 3)
 42 | 		throw std::runtime_error("Failed to pass args <plyfile> <source dir> [scaffold dir]");
 43 | 
 44 | 	int skyboxpoints = 0;
 45 | 	if (argc > 4)
 46 | 	{
 47 | 		const char* scaffoldpath = nullptr;
 48 | 		scaffoldpath = argv[4];
 49 | 		std::ifstream scaffoldfile(std::string(scaffoldpath) + "/pc_info.txt");
 50 | 
 51 | 		std::string line;
 52 | 		std::getline(scaffoldfile, line);
 53 | 		skyboxpoints = std::atoi(line.c_str());
 54 | 	}
 55 | 
 56 | 	std::vector<Gaussian> gaussians_unfiltered;
 57 | 	try
 58 | 	{
 59 | 		Loader::loadPly(argv[1], gaussians_unfiltered, skyboxpoints);
 60 | 	}
 61 | 	catch(const std::runtime_error& e)
 62 | 	{
 63 | 		std::cout << "Could not load .ply. Attempt loading .bin\n";
 64 | 		std::string filename(argv[1]);
 65 | 		filename.pop_back();  
 66 | 		filename.pop_back();
 67 | 		filename.pop_back();
 68 | 		filename = filename + "bin";
 69 | 		std::cout << filename << std::endl;
 70 | 		Loader::loadBin(filename.c_str(), gaussians_unfiltered, skyboxpoints);
 71 | 	}
 72 | 
 73 | 	int valid = 0;
 74 | 	bool not_warned = true;
 75 | 	std::vector<Gaussian> gaussians(gaussians_unfiltered.size());
 76 | 	for (int i = 0; i < gaussians_unfiltered.size(); i++)
 77 | 	{
 78 | 		Gaussian& g = gaussians_unfiltered[i];
 79 | 		if (std::isinf(g.opacity))
 80 | 		{
 81 | 			if (not_warned)
 82 | 				std::cout << "Found Inf opacity";
 83 | 			not_warned = false;
 84 | 			continue;
 85 | 		}
 86 | 		if (std::isnan(g.opacity))
 87 | 		{
 88 | 			if (not_warned)
 89 | 				std::cout << "Found NaN opacity";
 90 | 			not_warned = false;
 91 | 			continue;
 92 | 		}
 93 | 		if (g.scale.hasNaN())
 94 | 		{
 95 | 			if (not_warned)
 96 | 				std::cout << "Found NaN scale";
 97 | 			not_warned = false;
 98 | 			continue;
 99 | 		}
100 | 		if (g.rotation.hasNaN())
101 | 		{
102 | 			if (not_warned)
103 | 				std::cout << "Found NaN rot";
104 | 			not_warned = false;
105 | 			continue;
106 | 		}
107 | 		if (g.position.hasNaN())
108 | 		{
109 | 			if (not_warned)
110 | 				std::cout << "Found NaN pos";
111 | 			not_warned = false;
112 | 			continue;
113 | 		}
114 | 		if (g.shs.hasNaN())
115 | 		{
116 | 			if(not_warned)
117 | 				std::cout << "Found NaN sh";
118 | 			not_warned = false;
119 | 			continue;
120 | 		}
121 | 		if (g.opacity == 0)
122 | 		{
123 | 			if (not_warned)
124 | 				std::cout << "Found 0 opacity";
125 | 			not_warned = false;
126 | 			continue;
127 | 		}
128 | 
129 | 		gaussians[valid] = g;
130 | 		valid++;
131 | 	}
132 | 
133 | 	gaussians.resize(valid);
134 | 
135 | 	std::cout << "Generating" << std::endl;
136 | 
137 | 	PointbasedKdTreeGenerator generator;
138 | 	auto root = generator.generate(gaussians);
139 | 
140 | 	std::cout << "Merging" << std::endl;
141 | 
142 | 	ClusterMerger merger;
143 | 	merger.merge(root, gaussians);
144 | 
145 | 	std::cout << "Fixing rotations" << std::endl;
146 | 	RotationAligner::align(root, gaussians);
147 | 
148 | 	std::cout << "Filtering" << std::endl;
149 | 	float limit = 0.0005f;
150 | 
151 | 	char* filterpath = "";
152 | 	filterpath = argv[2];
153 | 
154 | 	AppearanceFilter filter;
155 | 	filter.init(filterpath);
156 | 	filter.filter(root, gaussians, limit, 2.0f);
157 | 
158 | 	filter.writeAnchors((std::string(argv[3]) + "/anchors.bin").c_str(), root, gaussians, limit);
159 | 	
160 | 	std::cout << "Writing" << std::endl;
161 | 
162 | 	Writer::writeHierarchy((std::string(argv[3]) + "/hierarchy.hier").c_str(), gaussians, root, true);
163 | }


--------------------------------------------------------------------------------
/mainHierarchyMerger.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | 
 13 | #include "loader.h"
 14 | #include "writer.h"
 15 | #include "FlatGenerator.h"
 16 | #include "PointbasedKdTreeGenerator.h"
 17 | #include "AvgMerger.h"
 18 | #include "ClusterMerger.h"
 19 | #include "common.h"
 20 | #include "dependencies/json.hpp"
 21 | #include "hierarchy_explicit_loader.h"
 22 | #include <vector>
 23 | #include <iostream>
 24 | #include <fstream>
 25 | #include <filesystem>
 26 | #include "appearance_filter.h"
 27 | #include "rotation_aligner.h"
 28 | 
 29 | using json = nlohmann::json;
 30 | 
 31 | void recTraverse(ExplicitTreeNode* node, int& zerocount)
 32 | {
 33 | 	if (node->depth == 0)
 34 | 		zerocount++;
 35 | 	if (node->children.size() > 0 && node->depth == 0)
 36 | 		throw std::runtime_error("Leaf nodes should never have children!");
 37 | 
 38 | 	for (auto c : node->children)
 39 | 	{
 40 | 		recTraverse(c, zerocount);
 41 | 	}
 42 | }
 43 | 
 44 | int main(int argc, char* argv[])
 45 | {
 46 | 	if (argc < 5)
 47 | 		throw std::runtime_error("Failed to pass filename");
 48 | 
 49 | 	int chunk_count(argc - 5);
 50 | 	std::string rootpath(argv[1]);
 51 | 	if (std::stoi(argv[2]))
 52 | 	{
 53 | 		for (int chunk_id(0); chunk_id < chunk_count; chunk_id++)
 54 | 		{
 55 | 			int argidx(chunk_id + 5);
 56 | 			std::cout << "Building hierarchy for chunk " << argv[chunk_id + 5] << std::endl;
 57 | 
 58 | 			AppearanceFilter filter;
 59 | 			std::string inpath(argv[3]);
 60 | 			filter.init((inpath + "/" + argv[argidx]).c_str());
 61 | 
 62 | 			std::vector<Gaussian> gaussians;
 63 | 			ExplicitTreeNode *root;
 64 | 			Loader::loadPly((rootpath + "/" + argv[argidx] + "/point_cloud/iteration_30000/point_cloud.ply").c_str(), gaussians);
 65 | 
 66 | 			std::cout << "Generating" << std::endl;
 67 | 			PointbasedKdTreeGenerator generator;
 68 | 			root = generator.generate(gaussians);
 69 | 
 70 | 			std::cout << "Merging" << std::endl;
 71 | 			ClusterMerger merger;
 72 | 			merger.merge(root, gaussians);
 73 | 
 74 | 			std::cout << "Fixing rotations" << std::endl;
 75 | 			RotationAligner::align(root, gaussians);
 76 | 
 77 | 			std::cout << "Filtering" << std::endl;
 78 | 			filter.filter(root, gaussians, 0.0005f, 4.0f);
 79 | 
 80 | 			std::cout << "Writing" << std::endl;
 81 | 
 82 | 			Writer::writeHierarchy(
 83 | 				(rootpath + "/" + argv[argidx] + "/chunk.hier").c_str(), 
 84 | 				gaussians, 
 85 | 				root);
 86 | 		}
 87 | 	}
 88 | 	else
 89 | 	{
 90 | 		// Read chunk centers
 91 | 		std::string inpath(argv[3]);
 92 | 		std::vector<Eigen::Vector3f> chunk_centers(chunk_count);
 93 | 		for (int chunk_id(0); chunk_id < chunk_count; chunk_id++)
 94 | 		{
 95 | 			int argidx(chunk_id + 5);
 96 | 			std::ifstream f(inpath + "/" + argv[argidx] + "/center.txt");
 97 | 			Eigen::Vector3f chunk_center(0.f, 0.f, 0.f);
 98 | 			f >> chunk_center[0]; f >> chunk_center[1]; f >> chunk_center[2];
 99 | 			chunk_centers[chunk_id] = chunk_center;
100 | 		}
101 | 
102 | 		// Read per chunk hierarchies and discard unwanted primitives 
103 | 		// based on the distance to the chunk's center
104 | 		std::vector<Gaussian> gaussians; 
105 | 		ExplicitTreeNode* root = new ExplicitTreeNode;
106 | 
107 | 		for (int chunk_id(0); chunk_id < chunk_count; chunk_id++)
108 | 		{
109 | 			int argidx(chunk_id + 5);
110 | 			std::cout << "Adding hierarchy for chunk " << argv[argidx] << std::endl;
111 | 			
112 | 			ExplicitTreeNode* chunkRoot = new ExplicitTreeNode;
113 | 			HierarchyExplicitLoader::loadExplicit(
114 | 				(rootpath + "/" + argv[argidx] + "/hierarchy.hier_opt").c_str(),
115 | 				gaussians, chunkRoot, chunk_id, chunk_centers);
116 | 
117 | 			if (chunk_id == 0)
118 | 			{
119 | 				root->bounds = chunkRoot->bounds;
120 | 			}
121 | 			else
122 | 			{	
123 | 				for (int idx(0); idx < 3; idx++)
124 | 				{
125 | 					root->bounds.minn[idx] = std::min(root->bounds.minn[idx], chunkRoot->bounds.minn[idx]);
126 | 					root->bounds.maxx[idx] = std::max(root->bounds.maxx[idx], chunkRoot->bounds.maxx[idx]);
127 | 				}
128 | 			}
129 | 			root->depth = std::max(root->depth, chunkRoot->depth + 1);
130 | 			root->children.push_back(chunkRoot);
131 | 			root->merged.push_back(chunkRoot->merged[0]);
132 | 			root->bounds.maxx[3] = 1e9f;
133 | 			root->bounds.minn[3] = 1e9f;
134 | 		}
135 | 
136 | 		Writer::writeHierarchy(
137 | 			(argv[4]), 
138 | 			gaussians, root, true);
139 | 	}
140 | }


--------------------------------------------------------------------------------
/rotation_aligner.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #include "rotation_aligner.h"
 13 | 
 14 | float frobenius(Eigen::Matrix3f& A, Eigen::Matrix3f& B)
 15 | {
 16 | 	float sum = 0.f;
 17 | 	for (int i = 0; i < 3; i++)
 18 | 		for (int j = 0; j < 3; j++)
 19 | 			sum += A(i, j) * B(i, j);
 20 | 	return sum;
 21 | }
 22 | 
 23 | void matchExhaustive(const Gaussian& ref, Gaussian& match)
 24 | {
 25 | 	auto Rref = matrixFromQuat(ref.rotation);
 26 | 	auto Rmatch = matrixFromQuat(match.rotation);
 27 | 
 28 | 	auto ror = match.rotation;
 29 | 
 30 | 	Cov pre;
 31 | 	computeCovariance(match.scale, match.rotation, pre);
 32 | 
 33 | 	bool noscore = true;
 34 | 	float best_score = 0;
 35 | 	int best_i, best_j, best_k;
 36 | 	Eigen::Matrix3f base, test, best;
 37 | 	for (int i = 0; i < 3; i++)
 38 | 	{
 39 | 		base.col(0) = Rmatch.col(i);
 40 | 		for (int j = 0; j < 3; j++)
 41 | 		{
 42 | 			if (j == i)
 43 | 				continue;
 44 | 
 45 | 			base.col(1) = Rmatch.col(j);
 46 | 			for (int k = 0; k < 3; k++)
 47 | 			{
 48 | 				if (k == j || k == i)
 49 | 					continue;
 50 | 
 51 | 				base.col(2) = Rmatch.col(k);
 52 | 
 53 | 				int state = 0;
 54 | 				while (state < 8)
 55 | 				{
 56 | 					for (int w = 0; w < 3; w++)
 57 | 						if ((state & (1 << w)))
 58 | 							test.col(w) = -base.col(w);
 59 | 						else
 60 | 							test.col(w) = base.col(w);
 61 | 
 62 | 					state++;
 63 | 
 64 | 					if (test.col(0).cross(test.col(1)).dot(test.col(2)) < 0)
 65 | 						continue;
 66 | 
 67 | 					float score = frobenius(test, Rref);
 68 | 					if (score > best_score)
 69 | 					{
 70 | 						best_i = i;
 71 | 						best_j = j;
 72 | 						best_k = k;
 73 | 						best = test;
 74 | 						best_score = score;
 75 | 						noscore = false;
 76 | 					}
 77 | 				}
 78 | 			}
 79 | 		}
 80 | 	}
 81 | 
 82 | 	auto q = Eigen::Quaternionf(best);
 83 | 	match.rotation = { q.w(), q.x(), q.y(), q.z() };
 84 | 	match.scale = Eigen::Vector3f(match.scale[best_i], match.scale[best_j], match.scale[best_k]);
 85 | 
 86 | 	Cov cov;
 87 | 	computeCovariance(match.scale, match.rotation, cov);
 88 | }
 89 | 
 90 | void topDownAlign(ExplicitTreeNode* node, std::vector<Gaussian>& gaussians)
 91 | {
 92 | 	if (node->merged.size() != 0)
 93 | 	{
 94 | 		Gaussian& ref = node->merged[0];
 95 | 
 96 | 		for (auto& child : node->children)
 97 | 		{
 98 | 			for (auto& m : child->merged)
 99 | 			{
100 | 				matchExhaustive(ref, m);
101 | 			}
102 | 			for (auto& i : child->leaf_indices)
103 | 			{
104 | 				matchExhaustive(ref, gaussians[i]);
105 | 			}
106 | 
107 | 			topDownAlign(child, gaussians);
108 | 		}
109 | 	}
110 | }
111 | 
112 | void RotationAligner::align(ExplicitTreeNode* root, std::vector<Gaussian>& gaussians)
113 | {
114 | 	topDownAlign(root, gaussians);
115 | }


--------------------------------------------------------------------------------
/rotation_aligner.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | #include "common.h"
14 | 
15 | class RotationAligner
16 | {
17 | public:
18 | 	static void align(ExplicitTreeNode* root, std::vector<Gaussian>& gaussians);
19 | };


--------------------------------------------------------------------------------
/runtime_maintenance.cu:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | 
 13 | #include "cuda_runtime.h"
 14 | #include "device_launch_parameters.h"
 15 | #include <cub/cub.cuh>
 16 | #include <cub/device/device_radix_sort.cuh>
 17 | #include <iostream>
 18 | #include <algorithm>
 19 | #include <numeric>
 20 | #include <thrust/sequence.h>
 21 | #include <thrust/device_vector.h>
 22 | #include <cuda.h>
 23 | #include <fstream>
 24 | #include <cooperative_groups.h>
 25 | #include <cooperative_groups/reduce.h>
 26 | #include <nvtx3/nvToolsExt.h>
 27 | #include <thrust/host_vector.h>
 28 | #include <tuple>
 29 | #include "types.h"
 30 | #include "runtime_maintenance.h"
 31 | 
 32 | __device__ int safeexc(int* data, int index)
 33 | {
 34 | 	if (index == 0)
 35 | 		return 0;
 36 | 	return data[index - 1];
 37 | }
 38 | 
 39 | __global__ void rearrange(
 40 | 	int N,
 41 | 	int* active_nodes,
 42 | 	int* new_active_nodes,
 43 | 	const Node* nodes,
 44 | 	const Box* boxes,
 45 | 	const float3* pos_space,
 46 | 	const float4* rot_space,
 47 | 	const float* sh_space,
 48 | 	const float* alpha_space,
 49 | 	const float3* scale_space,
 50 | 	int* split_space,
 51 | 	Node* new_nodes,
 52 | 	Box* new_boxes,
 53 | 	float3* new_pos_space,
 54 | 	float4* new_rot_space,
 55 | 	float* new_sh_space,
 56 | 	float* new_alpha_space,
 57 | 	float3* new_scale_space,
 58 | 	int* new_split_space,
 59 | 	int* cuda2cpu_src,
 60 | 	int* cuda2cpu_dst,
 61 | 	int* node_indices,
 62 | 	int* gaussian_indices)
 63 | {
 64 | 	int idx = threadIdx.x + blockDim.x * blockIdx.x;
 65 | 	if ((idx/15) >= N)
 66 | 		return;
 67 | 
 68 | 	int role = idx / N;
 69 | 	idx = idx % N;
 70 | 	bool actor = role == 0;
 71 | 
 72 | 	int node_id = active_nodes[idx];
 73 | 	int target_id = safeexc(node_indices, node_id);
 74 | 
 75 | 	Node node = nodes[node_id];
 76 | 
 77 | 	if (actor)
 78 | 	{
 79 | 		new_boxes[target_id] = boxes[node_id];
 80 | 		new_split_space[target_id] = split_space[node_id];
 81 | 
 82 | 		if (split_space[node_id] == 0) // Every unexpanded node is gone
 83 | 			node.start_children = -1;
 84 | 		split_space[node_id] = 0; // Clean up after yourself
 85 | 
 86 | 		int new_parent = node.parent == -1 ? -1 : safeexc(node_indices, node.parent);
 87 | 		int new_start_children = node.start_children == -1 ? -1 : safeexc(node_indices, node.start_children);
 88 | 
 89 | 		node.parent = new_parent;
 90 | 		node.start_children = new_start_children;
 91 | 	}
 92 | 
 93 | 	int new_start = safeexc(gaussian_indices, node_id);
 94 | 	for (int i = 0; i < node.count_leafs + node.count_merged; i++)
 95 | 	{
 96 | 		int dst = new_start + i;
 97 | 		int src = node.start + i;
 98 | 		if (role == 0)
 99 | 			new_pos_space[dst] = pos_space[src];
100 | 		else if (role == 1)
101 | 			new_rot_space[dst] = rot_space[src];
102 | 		else if (role >= 2 && role < 14)
103 | 			*(((float4*)(new_sh_space + dst * 48)) + role - 2) = *(((float4*)(sh_space + src * 48)) + role - 2);
104 | 		else
105 | 		{
106 | 			new_alpha_space[dst] = alpha_space[src];
107 | 			new_scale_space[dst] = scale_space[src];
108 | 		}
109 | 	}
110 | 
111 | 	if (actor)
112 | 	{
113 | 		node.start = new_start;
114 | 		new_nodes[target_id] = node;
115 | 		new_active_nodes[idx] = idx;
116 | 		cuda2cpu_dst[target_id] = cuda2cpu_src[node_id];
117 | 	}
118 | }
119 | 
120 | void Maintenance::reorder(
121 | 	int N,
122 | 	int* active_nodes,
123 | 	int* new_active_nodes,
124 | 	const int* nodes,
125 | 	const float* boxes,
126 | 	const float* pos_space,
127 | 	const float* rot_space,
128 | 	const float* sh_space,
129 | 	const float* alpha_space,
130 | 	const float* scale_space,
131 | 	int* split_space,
132 | 	int* new_nodes,
133 | 	float* new_boxes,
134 | 	float* new_pos_space,
135 | 	float* new_rot_space,
136 | 	float* new_sh_space,
137 | 	float* new_alpha_space,
138 | 	float* new_scale_space,
139 | 	int* new_split_space,
140 | 	int* cuda2cpu_src,
141 | 	int* cuda2cpu_dst,
142 | 	int* node_indices,
143 | 	int* gaussian_indices,
144 | 	void* streamy
145 | )
146 | {
147 | 	cudaStream_t stream = (cudaStream_t)streamy;
148 | 	int num_blocks = (N * 15 + 255) / 256;
149 | 	rearrange << <num_blocks, 256, 0, stream >> > (
150 | 		N,
151 | 		active_nodes,
152 | 		new_active_nodes,
153 | 		(Node*)nodes,
154 | 		(Box*)boxes,
155 | 		(float3*)pos_space,
156 | 		(float4*)rot_space,
157 | 		sh_space,
158 | 		alpha_space,
159 | 		(float3*)scale_space,
160 | 		split_space,
161 | 		(Node*)new_nodes,
162 | 		(Box*)new_boxes,
163 | 		(float3*)new_pos_space,
164 | 		(float4*)new_rot_space,
165 | 		new_sh_space,
166 | 		new_alpha_space,
167 | 		(float3*)new_scale_space,
168 | 		new_split_space,
169 | 		cuda2cpu_src,
170 | 		cuda2cpu_dst,
171 | 		node_indices,
172 | 		gaussian_indices
173 | 		);
174 | }
175 | 
176 | __global__ void mark(
177 | 	int N,
178 | 	const int* indices,
179 | 	Node* nodes,
180 | 	int* nodes_count,
181 | 	int* gaussians_count)
182 | {
183 | 	int idx = threadIdx.x + blockDim.x * blockIdx.x;
184 | 	if (idx >= N)
185 | 		return;
186 | 
187 | 	int node_id = indices[idx];
188 | 	Node node = nodes[node_id];
189 | 	nodes_count[node_id] = 1;
190 | 	gaussians_count[node_id] = node.count_merged + node.count_leafs;
191 | }
192 | 
193 | __global__ void zero(int N, int* a, int* b)
194 | {
195 | 	int idx = threadIdx.x + blockDim.x * blockIdx.x;
196 | 	if (idx >= N)
197 | 		return;
198 | 
199 | 	a[idx] = 0;
200 | 	b[idx] = 0;
201 | }
202 | 
203 | void Maintenance::compactPart1(
204 | 	int topN,
205 | 	int N,
206 | 	int* active_nodes,
207 | 	int* new_active_nodes,
208 | 	const int* nodes,
209 | 	const float* boxes,
210 | 	const float* pos_space,
211 | 	const float* rot_space,
212 | 	const float* sh_space,
213 | 	const float* alpha_space,
214 | 	const float* scale_space,
215 | 	int* split_space,
216 | 	int* new_nodes,
217 | 	float* new_boxes,
218 | 	float* new_pos_space,
219 | 	float* new_rot_space,
220 | 	float* new_sh_space,
221 | 	float* new_alpha_space,
222 | 	float* new_scale_space,
223 | 	int* new_split_space,
224 | 	int* cuda2cpu_src,
225 | 	int* cuda2cpu_dst,
226 | 	int* NsrcI,
227 | 	int* NsrcI2,
228 | 	int* NdstI,
229 | 	int* NdstI2,
230 | 	char*& scratchspace,
231 | 	size_t& scratchspacesize,
232 | 	void* streamy,
233 | 	int* count
234 | )
235 | {
236 | 	cudaStream_t stream = (cudaStream_t)streamy;
237 | 
238 | 	zero << <(topN + 255) / 256, 256, 0, stream >> > (topN, NsrcI, NsrcI2);
239 | 	mark << <(N + 255) / 256, 256, 0, stream >> > (N, active_nodes, (Node*)nodes, NsrcI, NsrcI2);
240 | 
241 | 	cub::DeviceScan::InclusiveSum(scratchspace, scratchspacesize, NsrcI, NdstI, topN, stream);
242 | 	cudaMemcpyAsync(count, NdstI + topN - 1, sizeof(int), cudaMemcpyDeviceToHost, stream);
243 | 	cudaStreamSynchronize(stream);
244 | }
245 | 
246 | __global__ void compressCUDA(
247 | 	int N,
248 | 	Node* nodes,
249 | 	float3* scales,
250 | 	float4* rots,
251 | 	float* shs,
252 | 	float* opacities
253 | )
254 | {
255 | 	int idx = threadIdx.x + blockIdx.x * blockDim.x;
256 | 	if (idx >= N)
257 | 		return;
258 | 
259 | 	Node node = nodes[idx];
260 | 
261 | 	int g = node.start;
262 | 
263 | 	int parent = g;
264 | 	if (node.parent != -1)
265 | 	{
266 | 		parent = nodes[node.parent].start;
267 | 	}
268 | 
269 | 	__half2 mp;
270 | 	float* scale = (float*)&scales[g];
271 | 	float* pscale = (float*)&scales[parent];
272 | 	for (int i = 0; i < 3; i++)
273 | 	{
274 | 		mp.x = scale[i];
275 | 		mp.y = pscale[i];
276 | 		scale[i] = *((float*)&mp); // We are overwriting ourselves
277 | 	}
278 | }
279 | 
280 | void Maintenance::compress(
281 | 	int N,
282 | 	int* nodes,
283 | 	float* scales,
284 | 	float* rots,
285 | 	float* shs,
286 | 	float* opacs
287 | )
288 | {
289 | 	int num_blocks = (N + 255) / 256;
290 | 	compressCUDA << <num_blocks, 256 >> > (N,
291 | 		(Node*)nodes,
292 | 		(float3*)scales,
293 | 		(float4*)rots,
294 | 		shs,
295 | 		opacs
296 | 		);
297 | }
298 | 
299 | void Maintenance::compactPart2(
300 | 	int topN,
301 | 	int N,
302 | 	int* active_nodes,
303 | 	int* new_active_nodes,
304 | 	const int* nodes,
305 | 	const float* boxes,
306 | 	const float* pos_space,
307 | 	const float* rot_space,
308 | 	const float* sh_space,
309 | 	const float* alpha_space,
310 | 	const float* scale_space,
311 | 	int* split_space,
312 | 	int* new_nodes,
313 | 	float* new_boxes,
314 | 	float* new_pos_space,
315 | 	float* new_rot_space,
316 | 	float* new_sh_space,
317 | 	float* new_alpha_space,
318 | 	float* new_scale_space,
319 | 	int* new_split_space,
320 | 	int* cuda2cpu_src,
321 | 	int* cuda2cpu_dst,
322 | 	int* NsrcI,
323 | 	int* NsrcI2,
324 | 	int* NdstI,
325 | 	int* NdstI2,
326 | 	char*& scratchspace,
327 | 	size_t& scratchspacesize,
328 | 	void* streamy,
329 | 	int* count
330 | )
331 | {
332 | 	cudaStream_t stream = (cudaStream_t)streamy;
333 | 
334 | 	cub::DeviceScan::InclusiveSum(scratchspace, scratchspacesize, NsrcI2, NdstI2, topN, stream);
335 | 
336 | 	cudaMemcpyAsync(count, NdstI2 + topN - 1, sizeof(int), cudaMemcpyDeviceToHost, stream);
337 | 	cudaStreamSynchronize(stream);
338 | 
339 | 	reorder(
340 | 		N,
341 | 		active_nodes,
342 | 		new_active_nodes,
343 | 		nodes,
344 | 		boxes,
345 | 		pos_space,
346 | 		rot_space,
347 | 		sh_space,
348 | 		alpha_space,
349 | 		scale_space,
350 | 		split_space,
351 | 		new_nodes,
352 | 		new_boxes,
353 | 		new_pos_space,
354 | 		new_rot_space,
355 | 		new_sh_space,
356 | 		new_alpha_space,
357 | 		new_scale_space,
358 | 		new_split_space,
359 | 		cuda2cpu_src,
360 | 		cuda2cpu_dst,
361 | 		NdstI,
362 | 		NdstI2,
363 | 		stream
364 | 	);
365 | }
366 | 
367 | __global__ void setStarts(Node* nodes, int N, int* indices, int* starts)
368 | {
369 | 	int idx = threadIdx.x + blockDim.x * blockIdx.x;
370 | 	if (idx >= N)
371 | 		return;
372 | 
373 | 	int parent_id = indices[idx];
374 | 	nodes[parent_id].start_children = starts[idx];
375 | }
376 | 
377 | void Maintenance::updateStarts(
378 | 	int* nodes,
379 | 	int num_indices,
380 | 	int* indices,
381 | 	int* starts,
382 | 	void* streamy
383 | )
384 | {
385 | 	cudaStream_t stream = (cudaStream_t)streamy;
386 | 	int num_blocks = (num_indices + 255) / 256;
387 | 	setStarts << <num_blocks, 256, 0, stream >> > ((Node*)nodes, num_indices, indices, starts);
388 | }


--------------------------------------------------------------------------------
/runtime_maintenance.h:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #pragma once
 13 | #include <cstdio>
 14 | #include <tuple>
 15 | #include <string>
 16 | 
 17 | class Maintenance
 18 | {
 19 | public:
 20 | 
 21 | 	static void reorder(
 22 | 		int N,
 23 | 		int* active_nodes,
 24 | 		int* new_active_nodes,
 25 | 		const int* nodes,
 26 | 		const float* boxes,
 27 | 		const float* pos_space,
 28 | 		const float* rot_space,
 29 | 		const float* sh_space,
 30 | 		const float* alpha_space,
 31 | 		const float* scale_space,
 32 | 		int* split_space,
 33 | 		int* new_nodes,
 34 | 		float* new_boxes,
 35 | 		float* new_pos_space,
 36 | 		float* new_rot_space,
 37 | 		float* new_sh_space,
 38 | 		float* new_alpha_space,
 39 | 		float* new_scale_space,
 40 | 		int* new_split_space,
 41 | 		int* cuda2cpu_src,
 42 | 		int* cuda2cpu_dst,
 43 | 		int* node_indices,
 44 | 		int* gaussian_indices,
 45 | 		void* streamy
 46 | 	);
 47 | 
 48 | 	static void compress(
 49 | 		int N,
 50 | 		int* nodes,
 51 | 		float* scales,
 52 | 		float* rots,
 53 | 		float* shs,
 54 | 		float* opacs
 55 | 	);
 56 | 
 57 | 	static void compactPart1(
 58 | 		int topN,
 59 | 		int N,
 60 | 		int* active_nodes,
 61 | 		int* new_active_nodes,
 62 | 		const int* nodes,
 63 | 		const float* boxes,
 64 | 		const float* pos_space,
 65 | 		const float* rot_space,
 66 | 		const float* sh_space,
 67 | 		const float* alpha_space,
 68 | 		const float* scale_space,
 69 | 		int* split_space,
 70 | 		int* new_nodes,
 71 | 		float* new_boxes,
 72 | 		float* new_pos_space,
 73 | 		float* new_rot_space,
 74 | 		float* new_sh_space,
 75 | 		float* new_alpha_space,
 76 | 		float* new_scale_space,
 77 | 		int* new_split_space,
 78 | 		int* cuda2cpu_src,
 79 | 		int* cuda2cpu_dst,
 80 | 		int* NsrcI,
 81 | 		int* NsrcI2,
 82 | 		int* NdstI,
 83 | 		int* NdstI2,
 84 | 		char*& scratchspace,
 85 | 		size_t& scratchspacesize,
 86 | 		void* streamy,
 87 | 		int* count
 88 | 	);
 89 | 
 90 | 	static void compactPart2(
 91 | 		int topN,
 92 | 		int N,
 93 | 		int* active_nodes,
 94 | 		int* new_active_nodes,
 95 | 		const int* nodes,
 96 | 		const float* boxes,
 97 | 		const float* pos_space,
 98 | 		const float* rot_space,
 99 | 		const float* sh_space,
100 | 		const float* alpha_space,
101 | 		const float* scale_space,
102 | 		int* split_space,
103 | 		int* new_nodes,
104 | 		float* new_boxes,
105 | 		float* new_pos_space,
106 | 		float* new_rot_space,
107 | 		float* new_sh_space,
108 | 		float* new_alpha_space,
109 | 		float* new_scale_space,
110 | 		int* new_split_space,
111 | 		int* cuda2cpu_src,
112 | 		int* cuda2cpu_dst,
113 | 		int* NsrcI,
114 | 		int* NsrcI2,
115 | 		int* NdstI,
116 | 		int* NdstI2,
117 | 		char*& scratchspace,
118 | 		size_t& scratchspacesize,
119 | 		void* streamy,
120 | 		int* count
121 | 	);
122 | 
123 | 	static void updateStarts(
124 | 		int* nodes,
125 | 		int num_indices,
126 | 		int* indices,
127 | 		int* starts,
128 | 		void* streamy
129 | 	);
130 | };


--------------------------------------------------------------------------------
/runtime_switching.cu:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | #include "cuda_runtime.h"
 13 | #include "device_launch_parameters.h"
 14 | #include <float.h>
 15 | #include <cub/cub.cuh>
 16 | #include <cub/device/device_radix_sort.cuh>
 17 | #include <iostream>
 18 | #include <algorithm>
 19 | #include <numeric>
 20 | #include <thrust/sequence.h>
 21 | #include <thrust/device_vector.h>
 22 | #include <cuda.h>
 23 | #include <fstream>
 24 | #include <cooperative_groups.h>
 25 | #include <cooperative_groups/reduce.h>
 26 | #include <nvtx3/nvToolsExt.h>
 27 | #include <thrust/host_vector.h>
 28 | #include <tuple>
 29 | #include "types.h"
 30 | #include "runtime_switching.h"
 31 | 
 32 | __global__ void markTargetNodes(Node* nodes, int N, int target, int* node_counts)
 33 | {
 34 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
 35 | 	if (idx >= N)
 36 | 		return;
 37 | 
 38 | 	int count = 0;
 39 | 	Node node = nodes[idx];
 40 | 	if (node.depth > target)
 41 | 		count = node.count_leafs;
 42 | 	else if (node.parent != -1)
 43 | 	{
 44 | 		Node parentnode = nodes[node.parent];
 45 | 		if (parentnode.depth > target)
 46 | 		{
 47 | 			count = node.count_leafs;
 48 | 			if (node.depth != 0)
 49 | 				count += node.count_merged;
 50 | 		}
 51 | 	}
 52 | 	node_counts[idx] = count;
 53 | }
 54 | 
 55 | __global__ void putRenderIndices(Node* nodes, int N, int* node_counts, int* node_offsets, int* render_indices, int* parent_indices = nullptr, int* nodes_for_render_indices = nullptr)
 56 | {
 57 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
 58 | 	if (idx >= N)
 59 | 		return;
 60 | 
 61 | 	Node node = nodes[idx];
 62 | 	int count = node_counts[idx];
 63 | 	int offset = idx == 0 ? 0 : node_offsets[idx - 1];
 64 | 	int start = node.start;
 65 | 	
 66 | 	int parentgaussian = -1;
 67 | 	if (node.parent != -1)
 68 | 	{
 69 | 		parentgaussian = nodes[node.parent].start;
 70 | 	}
 71 | 
 72 | 	for (int i = 0; i < count; i++)
 73 | 	{
 74 | 		render_indices[offset + i] = node.start + i;
 75 | 		if (parent_indices)
 76 | 			parent_indices[offset + i] = parentgaussian; 
 77 | 		if (nodes_for_render_indices)
 78 | 			nodes_for_render_indices[offset + i] = idx;
 79 | 	}
 80 | }
 81 | 
 82 | int Switching::expandToTarget(
 83 | 	int N,
 84 | 	int target,
 85 | 	int* nodes,
 86 | 	int* render_indices
 87 | )
 88 | {
 89 | 	thrust::device_vector<int> render_counts(N);
 90 | 	thrust::device_vector<int> render_offsets(N);
 91 | 
 92 | 	int num_blocks = (N + 255) / 256;
 93 | 	markTargetNodes << <num_blocks, 256 >> > ((Node*)nodes, N, target, render_counts.data().get());
 94 | 
 95 | 	size_t temp_storage_bytes;
 96 | 	thrust::device_vector<char> temp_storage;
 97 | 	cub::DeviceScan::InclusiveSum(nullptr, temp_storage_bytes, render_counts.data().get(), render_offsets.data().get(), N);
 98 | 	temp_storage.resize(temp_storage_bytes);
 99 | 	cub::DeviceScan::InclusiveSum(temp_storage.data().get(), temp_storage_bytes, render_counts.data().get(), render_offsets.data().get(), N);
100 | 
101 | 	putRenderIndices << <num_blocks, 256 >> > ((Node*)nodes, N, render_counts.data().get(), render_offsets.data().get(), render_indices);
102 | 
103 | 	int count = 0;
104 | 	cudaMemcpy(&count, render_offsets.data().get() + N - 1, sizeof(int), cudaMemcpyDeviceToHost);
105 | 	return count;
106 | }
107 | 
108 | __device__ bool inboxCUDA(Box& box, Point viewpoint)
109 | {
110 | 	bool inside = true;
111 | 	for (int i = 0; i < 3; i++)
112 | 	{
113 | 		inside &= viewpoint.xyz[i] >= box.minn.xyz[i] && viewpoint.xyz[i] <= box.maxx.xyz[i];
114 | 	}
115 | 	return inside;
116 | }
117 | 
118 | __device__ float pointboxdistCUDA(Box& box, Point viewpoint)
119 | {
120 | 	Point closest = {
121 | 		max(box.minn.xyz[0], min(box.maxx.xyz[0], viewpoint.xyz[0])),
122 | 		max(box.minn.xyz[1], min(box.maxx.xyz[1], viewpoint.xyz[1])),
123 | 		max(box.minn.xyz[2], min(box.maxx.xyz[2], viewpoint.xyz[2]))
124 | 	};
125 | 
126 | 	Point diff = {
127 | 		viewpoint.xyz[0] - closest.xyz[0],
128 | 		viewpoint.xyz[1] - closest.xyz[1],
129 | 		viewpoint.xyz[2] - closest.xyz[2]
130 | 	};
131 | 
132 | 	return sqrt(diff.xyz[0] * diff.xyz[0] + diff.xyz[1] * diff.xyz[1] + diff.xyz[2] * diff.xyz[2]);
133 | }
134 | 
135 | __device__ float computeSizeGPU(Box& box, Point viewpoint, Point zdir)
136 | {
137 | 	if (inboxCUDA(box, viewpoint))
138 | 		return FLT_MAX;
139 | 
140 | 	float min_dist = pointboxdistCUDA(box, viewpoint);
141 | 
142 | 	return box.minn.xyz[3] / min_dist;
143 | }
144 | 
145 | __global__ void changeNodesOnce(
146 | 	Node* nodes,
147 | 	int N,
148 | 	int* indices,
149 | 	Box* boxes,
150 | 	Point* viewpoint,
151 | 	Point zdir,
152 | 	float target_size,
153 | 	int* split,
154 | 	int* node_counts,
155 | 	int* node_ids,
156 | 	char* needs_children
157 | )
158 | {
159 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
160 | 	if (idx >= N)
161 | 		return;
162 | 
163 | 	int node_id = indices[idx];
164 | 	Node node = nodes[node_id];
165 | 	float size = computeSizeGPU(boxes[node_id], *viewpoint, zdir);
166 | 
167 | 	int count = 1; // repeat yourself
168 | 	char need_child = 0;
169 | 	if (size >= target_size)
170 | 	{
171 | 		if (node.depth > 0 && split[node_id] == 0) // split
172 | 		{
173 | 			if (node.start_children == -1)
174 | 			{
175 | 				node_ids[idx] = node_id;
176 | 				need_child = 1;
177 | 			}
178 | 			else
179 | 			{
180 | 				count += node.count_children;
181 | 				split[node_id] = 1;
182 | 			}
183 | 		}
184 | 	}
185 | 	else
186 | 	{
187 | 		int parent_node_id = node.parent;
188 | 		if (parent_node_id != -1)
189 | 		{
190 | 			Node parent_node = nodes[parent_node_id];
191 | 			float parent_size = computeSizeGPU(boxes[parent_node_id], *viewpoint, zdir);
192 | 			if (parent_size < target_size) // collapse
193 | 			{
194 | 				split[parent_node_id] = 0;
195 | 				count = 0; // forget yourself
196 | 			}
197 | 		}
198 | 	}
199 | 	needs_children[idx] = need_child;
200 | 	node_counts[idx] = count;
201 | }
202 | 
203 | __global__ void putNodes(
204 | 	Node* nodes,
205 | 	int N,
206 | 	int* indices,
207 | 	int* node_counts,
208 | 	int* node_offsets,
209 | 	int* next_nodes)
210 | {
211 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
212 | 	if (idx >= N)
213 | 		return;
214 | 
215 | 	int count = node_counts[idx];
216 | 	if (count == 0)
217 | 		return;
218 | 
219 | 	int node_id = indices[idx];
220 | 	Node node = nodes[node_id];
221 | 	int offset = idx == 0 ? 0 : node_offsets[idx - 1];
222 | 
223 | 	next_nodes[offset] = node_id;
224 | 	for (int i = 1; i < count; i++)
225 | 		next_nodes[offset + i] = node.start_children + i - 1;
226 | }
227 | 
228 | __global__ void countRenderIndicesIndexed(Node* nodes, int* split, int N, int* node_indices, int* render_counts)
229 | {
230 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
231 | 	if (idx >= N)
232 | 		return;
233 | 
234 | 	int node_idx = node_indices[idx];
235 | 
236 | 	Node node = nodes[node_idx];
237 | 	int count = node.count_leafs;
238 | 	if (node.depth > 0 && split[node_idx] == 0)
239 | 		count += node.count_merged;
240 | 
241 | 	render_counts[idx] = count;
242 | }
243 | 
244 | __global__ void putRenderIndicesIndexed(Node* nodes, int N, int* node_indices, int* render_counts, int* render_offsets, int* render_indices, int* parent_indices, int* nodes_of_render_indices, Box* boxes, float3* debug)
245 | {
246 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
247 | 	if (idx >= N)
248 | 		return;
249 | 
250 | 	int node_idx = node_indices[idx];
251 | 
252 | 	Node node = nodes[node_idx];
253 | 	int count = render_counts[idx];
254 | 	int offset = idx == 0 ? 0 : render_offsets[idx - 1];
255 | 	int start = node.start;
256 | 
257 | 	int parentgaussian = -1;
258 | 	if (node.parent != -1)
259 | 	{
260 | 		parentgaussian = nodes[node.parent].start;
261 | 	}
262 | 
263 | 	for (int i = 0; i < count; i++)
264 | 	{
265 | 		render_indices[offset + i] = node.start + i;
266 | 		parent_indices[offset + i] = parentgaussian;
267 | 		nodes_of_render_indices[offset + i] = node_idx;
268 | 	}
269 | 
270 | 	if (debug != nullptr)
271 | 	{
272 | 		Box box = boxes[node_idx];
273 | 		for (int i = 0; i < count; i++)
274 | 		{
275 | 			float red = min(1.0f, node.depth / 10.0f);
276 | 			debug[offset + i] = { red, 1.0f - red, 0 };
277 | 			if (node.depth == 0)
278 | 				debug[offset + i] = { 0, 0, 1.0f };
279 | 		}
280 | 	}
281 | }
282 | 
283 | void Switching::changeToSizeStep(
284 | 	float target_size,
285 | 	int N,
286 | 	int* node_indices,
287 | 	int* new_node_indices,
288 | 	int* nodes,
289 | 	float* boxes,
290 | 	float* viewpoint,
291 | 	float x, float y, float z,
292 | 	int* split,
293 | 	int* render_indices,
294 | 	int* parent_indices,
295 | 	int* nodes_of_render_indices,
296 | 	int* nodes_to_expand,
297 | 	float* debug,
298 | 	char*& scratchspace,
299 | 	size_t& scratchspacesize,
300 | 	int* NsrcI,
301 | 	int* NdstI,
302 | 	char* NdstC,
303 | 	int* numI,
304 | 	int maxN,
305 | 	int& add_success,
306 | 	int* new_N,
307 | 	int* new_R,
308 | 	int* need_expansion,
309 | 	void* maintenanceStream)
310 | {
311 | 	cudaStream_t stream = (cudaStream_t)maintenanceStream;
312 | 
313 | 	int num_node_blocks = (N + 255) / 256;
314 | 
315 | 	Point zdir = { x, y, z };
316 | 
317 | 	int* num_to_expand = numI;
318 | 	int* node_counts = NsrcI, * node_offsets = NdstI, * node_ids = NdstI;
319 | 	char* need_children = NdstC;
320 | 	if (scratchspacesize == 0)
321 | 	{
322 | 		size_t testsize;
323 | 
324 | 		cub::DeviceScan::InclusiveSum(nullptr, testsize, node_counts, node_offsets, maxN, stream);
325 | 		scratchspacesize = testsize;
326 | 		cub::DeviceSelect::Flagged(nullptr, testsize, node_ids, need_children, nodes_to_expand, num_to_expand, maxN, stream);
327 | 		scratchspacesize = std::max(testsize, scratchspacesize);
328 | 
329 | 		if (scratchspace)
330 | 			cudaFree(scratchspace);
331 | 		scratchspacesize = testsize;
332 | 		cudaMalloc(&scratchspace, scratchspacesize);
333 | 	}
334 | 
335 | 	changeNodesOnce << <num_node_blocks, 256, 0, stream >> > (
336 | 		(Node*)nodes, 
337 | 		N, 
338 | 		node_indices, 
339 | 		(Box*)boxes, 
340 | 		(Point*)viewpoint, 
341 | 		zdir, 
342 | 		target_size, 
343 | 		split, 
344 | 		node_counts, 
345 | 		node_ids, 
346 | 		need_children
347 | 		);
348 | 
349 | 	cub::DeviceSelect::Flagged(scratchspace, scratchspacesize, node_ids, need_children, nodes_to_expand, num_to_expand, N, stream);
350 | 	cub::DeviceScan::InclusiveSum(scratchspace, scratchspacesize, node_counts, node_offsets, N, stream);
351 | 
352 | 	cudaMemcpyAsync(need_expansion, num_to_expand, sizeof(int), cudaMemcpyDeviceToHost, stream);
353 | 	cudaMemcpyAsync(new_N, node_offsets + N - 1, sizeof(int), cudaMemcpyDeviceToHost, stream);
354 | 	cudaStreamSynchronize(stream);
355 | 
356 | 	if (*new_N > maxN)
357 | 	{
358 | 		add_success = 0;
359 | 		return;
360 | 	}
361 | 
362 | 	putNodes << <num_node_blocks, 256, 0, stream>> > (
363 | 		(Node*)nodes,
364 | 		N, 
365 | 		node_indices, 
366 | 		node_counts, 
367 | 		node_offsets, 
368 | 		new_node_indices
369 | 		);
370 | 
371 | 	int num_render_blocks = (*new_N + 255) / 256;
372 | 	int* render_counts = NsrcI, * render_offsets = NdstI;
373 | 
374 | 	countRenderIndicesIndexed << <num_render_blocks, 256, 0, stream >> > (
375 | 		(Node*)nodes, 
376 | 		split, 
377 | 		*new_N, 
378 | 		new_node_indices, 
379 | 		render_counts
380 | 		);
381 | 
382 | 	cub::DeviceScan::InclusiveSum(scratchspace, scratchspacesize, render_counts, render_offsets, *new_N, stream);
383 | 
384 | 	putRenderIndicesIndexed << <num_render_blocks, 256, 0, stream >> > (
385 | 		(Node*)nodes, 
386 | 		*new_N, 
387 | 		new_node_indices, 
388 | 		render_counts, 
389 | 		render_offsets, 
390 | 		render_indices, 
391 | 		parent_indices, 
392 | 		nodes_of_render_indices, 
393 | 		(Box*)boxes,
394 | 		(float3*)debug
395 | 		);
396 | 
397 | 	cudaMemcpyAsync(new_R, render_offsets + *new_N - 1, sizeof(int), cudaMemcpyDeviceToHost, stream);
398 | 
399 | 	add_success = 1;
400 | }
401 | 
402 | __global__ void markNodesForSize(Node* nodes, Box* boxes, int N, Point* viewpoint, Point zdir, float target_size, int* render_counts, int* node_markers)
403 | {
404 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
405 | 	if (idx >= N)
406 | 		return;
407 | 
408 | 	int node_id = idx;
409 | 	Node node = nodes[node_id];
410 | 	float size = computeSizeGPU(boxes[node_id], *viewpoint, zdir);
411 | 
412 | 	int count = 0;
413 | 	if (size >= target_size)
414 | 		count = node.count_leafs;
415 | 	else if (node.parent != -1)
416 | 	{
417 | 		float parent_size = computeSizeGPU(boxes[node.parent], *viewpoint, zdir);
418 | 		if (parent_size >= target_size)
419 | 		{
420 | 			count = node.count_leafs;
421 | 			if (node.depth != 0)
422 | 				count += node.count_merged;
423 | 		}
424 | 	}
425 | 
426 | 	if (count != 0 && node_markers != nullptr)
427 | 		node_markers[node_id] = 1;
428 | 
429 | 	if (render_counts != nullptr)
430 | 		render_counts[node_id] = count;
431 | }
432 | 
433 | __global__ void computeTsIndexed(
434 | 	Node* nodes,
435 | 	Box* boxes,
436 | 	int N,
437 | 	int* indices,
438 | 	Point viewpoint,
439 | 	Point zdir,
440 | 	float target_size,
441 | 	float* ts,
442 | 	int* kids
443 | )
444 | {
445 | 	int idx = blockDim.x * blockIdx.x + threadIdx.x;
446 | 	if (idx >= N)
447 | 		return;
448 | 
449 | 	int node_id = indices[idx];
450 | 	Node node = nodes[node_id];
451 | 
452 | 	float t;
453 | 	if (node.parent == -1)
454 | 		t = 1.0f;
455 | 	else
456 | 	{
457 | 		float parentsize = computeSizeGPU(boxes[node.parent], viewpoint, zdir);
458 | 
459 | 		if (parentsize > 2.0f * target_size)
460 | 			t = 1.0f;
461 | 		else
462 | 		{
463 | 			float size = computeSizeGPU(boxes[node_id], viewpoint, zdir);
464 | 			float start = max(0.5f * parentsize, size);
465 | 			float diff = parentsize - start;
466 | 
467 | 			if (diff <= 0)
468 | 				t = 1.0f;
469 | 			else
470 | 			{
471 | 				float tdiff = max(0.0f, target_size - start);
472 | 				t = max(1.0f - (tdiff / diff), 0.0f);
473 | 			}
474 | 		}
475 | 	}
476 | 
477 | 	ts[idx] = t;
478 | 	kids[idx] = (node.parent == -1) ? 1 : nodes[node.parent].count_children;
479 | }
480 | 
481 | void Switching::getTsIndexed(
482 | 	int N,
483 | 	int* indices,
484 | 	float target_size,
485 | 	int* nodes,
486 | 	float* boxes,
487 | 	float vx, float vy, float vz,
488 | 	float x, float y, float z,
489 | 	float* ts,
490 | 	int* kids,
491 | 	void* stream
492 | )
493 | {
494 | 	Point zdir = { x, y, z };
495 | 	Point cam = { vx, vy, vz };
496 | 	int num_blocks = (N + 255) / 256;
497 | 	computeTsIndexed<<<num_blocks, 256, 0, (cudaStream_t)stream >>>(
498 | 		(Node*)nodes, 
499 | 		(Box*)boxes, 
500 | 		N, 
501 | 		indices, 
502 | 		cam,
503 | 		zdir, 
504 | 		target_size, 
505 | 		ts, 
506 | 		kids);
507 | }
508 | 
509 | int Switching::expandToSize(
510 | 	int N,
511 | 	float target_size,
512 | 	int* nodes,
513 | 	float* boxes,
514 | 	float* viewpoint,
515 | 	float x, float y, float z,
516 | 	int* render_indices,
517 | 	int* node_markers,
518 | 	int* parent_indices,
519 | 	int* nodes_for_render_indices)
520 | {
521 | 	size_t temp_storage_bytes;
522 | 	thrust::device_vector<char> temp_storage;
523 | 	thrust::device_vector<int> render_counts(N);
524 | 	thrust::device_vector<int> render_offsets(N);
525 | 
526 | 	Point zdir = { x, y, z };
527 | 
528 | 	int num_blocks = (N + 255) / 256;
529 | 	markNodesForSize << <num_blocks, 256 >> > ((Node*)nodes, (Box*)boxes, N, (Point*)viewpoint, zdir, target_size, render_counts.data().get(), node_markers);
530 | 
531 | 	cub::DeviceScan::InclusiveSum(nullptr, temp_storage_bytes, render_counts.data().get(), render_offsets.data().get(), N);
532 | 	temp_storage.resize(temp_storage_bytes);
533 | 	cub::DeviceScan::InclusiveSum(temp_storage.data().get(), temp_storage_bytes, render_counts.data().get(), render_offsets.data().get(), N);
534 | 
535 | 	putRenderIndices << <num_blocks, 256 >> > ((Node*)nodes, N, render_counts.data().get(), render_offsets.data().get(), render_indices, parent_indices, nodes_for_render_indices);
536 | 
537 | 	int count = 0;
538 | 	cudaMemcpy(&count, render_offsets.data().get() + N - 1, sizeof(int), cudaMemcpyDeviceToHost);
539 | 	return count;
540 | }
541 | 
542 | void Switching::markVisibleForAllViewpoints(
543 | 	float target_size,
544 | 	int* nodes,
545 | 	int num_nodes,
546 | 	float* boxes,
547 | 	float* viewpoints,
548 | 	int num_viewpoints,
549 | 	int* seen,
550 | 	float zx,
551 | 	float zy,
552 | 	float zz
553 | )
554 | {
555 | 	thrust::device_vector<int> seen_cuda(num_nodes);
556 | 	thrust::device_vector<Point> viewpoint_cuda(1);
557 | 	thrust::device_vector<Node> nodes_cuda(num_nodes);
558 | 	thrust::device_vector<Box> boxes_cuda(num_nodes);
559 | 
560 | 	cudaMemcpy(nodes_cuda.data().get(), nodes, sizeof(Node) * num_nodes, cudaMemcpyHostToDevice);
561 | 	cudaMemcpy(boxes_cuda.data().get(), boxes, sizeof(Box) * num_nodes, cudaMemcpyHostToDevice);
562 | 
563 | 	Point zdir = { zx, zy, zz };
564 | 
565 | 	Point* points = (Point*)viewpoints;
566 | 	int num_blocks = (num_nodes + 255) / 256;
567 | 	for (int i = 0; i < num_viewpoints; i++)
568 | 	{
569 | 		Point viewpoint = points[i];
570 | 		cudaMemcpy(viewpoint_cuda.data().get(), &viewpoint, sizeof(Point), cudaMemcpyHostToDevice);
571 | 
572 | 		markNodesForSize << <num_blocks, 256 >> > (
573 | 			nodes_cuda.data().get(),
574 | 			boxes_cuda.data().get(),
575 | 			num_nodes,
576 | 			viewpoint_cuda.data().get(),
577 | 			zdir,
578 | 			target_size,
579 | 			nullptr,
580 | 			seen_cuda.data().get());
581 | 	}
582 | 	cudaMemcpy(seen, seen_cuda.data().get(), sizeof(int) * num_nodes, cudaMemcpyDeviceToHost);
583 | 
584 | 	if (cudaDeviceSynchronize())
585 | 		std::cout << "Errors: " << cudaDeviceSynchronize() << std::endl;
586 | }


--------------------------------------------------------------------------------
/runtime_switching.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | #include <cstdio>
14 | #include <tuple>
15 | #include <string>
16 | 
17 | class Switching
18 | {
19 | public:
20 | 
21 | 	static int expandToTarget(
22 | 		int N,
23 | 		int target,
24 | 		int* nodes,
25 | 		int* render_indices
26 | 	);
27 | 
28 | 	static int expandToSize(
29 | 		int N,
30 | 		float target_size,
31 | 		int* nodes,
32 | 		float* boxes,
33 | 		float* viewpoint,
34 | 		float x, float y, float z,
35 | 		int* render_indices,
36 | 		int* node_markers,
37 | 		int* parent_indices = nullptr,
38 | 		int* nodes_for_render_indices=nullptr);
39 | 
40 | 	static void getTsIndexed(
41 | 		int N,
42 | 		int* indices,
43 | 		float target_size,
44 | 		int* nodes,
45 | 		float* boxes,
46 | 		float vx, float vy, float vz,
47 | 		float x, float y, float z,
48 | 		float* ts,
49 | 		int* kids,
50 | 		void* stream
51 | 	);
52 | 
53 | 	static void changeToSizeStep(
54 | 		float target_size,
55 | 		int N,
56 | 		int* node_indices,
57 | 		int* new_node_indices,
58 | 		int* nodes,
59 | 		float* boxes,
60 | 		float* viewpoint,
61 | 		float x, float y, float z,
62 | 		int* split,
63 | 		int* render_indices,
64 | 		int* parent_indices,
65 | 		int* nodes_of_render_indices,
66 | 		int* nodes_to_expand,
67 | 		float* debug,
68 | 		char*& scratchspace,
69 | 		size_t& scratchspacesize,
70 | 		int* NsrcI,
71 | 		int* NdstI,
72 | 		char* NdstC,
73 | 		int* numI,
74 | 		int maxN,
75 | 		int& add_success,
76 | 		int* new_N,
77 | 		int* new_R,
78 | 		int* need_expansion,
79 | 		void* maintenanceStream);
80 | 
81 | 	static void markVisibleForAllViewpoints(
82 | 		float target_size,
83 | 		int* nodes,
84 | 		int num_nodes,
85 | 		float* boxes,
86 | 		float* viewpoints,
87 | 		int num_viewpoints,
88 | 		int* seen,
89 | 		float zx,
90 | 		float zy,
91 | 		float zz
92 | 	);
93 | };


--------------------------------------------------------------------------------
/setup.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 setuptools import setup
13 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension
14 | import os
15 | os.path.dirname(os.path.abspath(__file__))
16 | 
17 | setup(
18 |     name="gaussian_hierarchy",
19 |     ext_modules=[
20 |         CUDAExtension(
21 |             name="gaussian_hierarchy._C",
22 |             sources=[
23 |             "hierarchy_loader.cpp",
24 |             "hierarchy_writer.cpp",
25 |             "traversal.cpp",
26 |             "runtime_switching.cu",
27 |             "torch/torch_interface.cpp",
28 |             "ext.cpp"],
29 |             extra_compile_args={"cxx": ["-I" + os.path.join(os.path.dirname(os.path.abspath(__file__)), "dependencies/eigen/")]}
30 |             )
31 |         ],
32 |     cmdclass={
33 |         'build_ext': BuildExtension
34 |     }
35 | )
36 | 


--------------------------------------------------------------------------------
/torch/torch_interface.cpp:
--------------------------------------------------------------------------------
  1 | #include "torch_interface.h"
  2 | #include "../hierarchy_loader.h"
  3 | #include "../hierarchy_writer.h"
  4 | #include "../traversal.h"
  5 | #include "../runtime_switching.h"
  6 | 
  7 | std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
  8 | LoadHierarchy(std::string filename)
  9 | {
 10 | 	HierarchyLoader loader;
 11 | 	
 12 | 	std::vector<Eigen::Vector3f> pos;
 13 | 	std::vector<SHs> shs;
 14 | 	std::vector<float> alphas;
 15 | 	std::vector<Eigen::Vector3f> scales;
 16 | 	std::vector<Eigen::Vector4f> rot;
 17 | 	std::vector<Node> nodes;
 18 | 	std::vector<Box> boxes;
 19 | 	
 20 | 	loader.load(filename.c_str(), pos, shs, alphas, scales, rot, nodes, boxes);
 21 | 	
 22 | 	int P = pos.size();
 23 | 	
 24 | 	torch::TensorOptions options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
 25 | 	torch::Tensor pos_tensor = torch::from_blob(pos.data(), {P, 3}, options).clone();
 26 | 	torch::Tensor shs_tensor = torch::from_blob(shs.data(), {P, 16, 3}, options).clone();
 27 | 	torch::Tensor alpha_tensor = torch::from_blob(alphas.data(), {P, 1}, options).clone();
 28 | 	torch::Tensor scale_tensor = torch::from_blob(scales.data(), {P, 3}, options).clone();
 29 | 	torch::Tensor rot_tensor = torch::from_blob(rot.data(), {P, 4}, options).clone();
 30 | 	
 31 | 	int N = nodes.size();
 32 | 	torch::TensorOptions intoptions = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
 33 | 	
 34 | 	torch::Tensor nodes_tensor = torch::from_blob(nodes.data(), {N, 7}, intoptions).clone();
 35 | 	torch::Tensor box_tensor = torch::from_blob(boxes.data(), {N, 2, 4}, options).clone();
 36 | 	
 37 | 	return std::make_tuple(pos_tensor, shs_tensor, alpha_tensor, scale_tensor, rot_tensor, nodes_tensor, box_tensor);
 38 | }
 39 | 
 40 | void WriteHierarchy(
 41 | 					std::string filename,
 42 | 					torch::Tensor& pos,
 43 | 					torch::Tensor& shs,
 44 | 					torch::Tensor& opacities,
 45 | 					torch::Tensor& log_scales,
 46 | 					torch::Tensor& rotations,
 47 | 					torch::Tensor& nodes,
 48 | 					torch::Tensor& boxes)
 49 | {
 50 | 	HierarchyWriter writer;
 51 | 	
 52 | 	int allP = pos.size(0);
 53 | 	int allN = nodes.size(0);
 54 | 	
 55 | 	writer.write(
 56 | 		filename.c_str(),
 57 | 		allP,
 58 | 		allN,
 59 | 		(Eigen::Vector3f*)pos.cpu().contiguous().data_ptr<float>(),
 60 | 		(SHs*)shs.cpu().contiguous().data_ptr<float>(),
 61 | 		opacities.cpu().contiguous().data_ptr<float>(),
 62 | 		(Eigen::Vector3f*)log_scales.cpu().contiguous().data_ptr<float>(),
 63 | 		(Eigen::Vector4f*)rotations.cpu().contiguous().data_ptr<float>(),
 64 | 		(Node*)nodes.cpu().contiguous().data_ptr<int>(),
 65 | 		(Box*)boxes.cpu().contiguous().data_ptr<float>()
 66 | 	);
 67 | }
 68 | 
 69 | torch::Tensor
 70 | ExpandToTarget(torch::Tensor& nodes, int target)
 71 | {
 72 | 	std::vector<int> indices = Traversal::expandToTarget((Node*)nodes.cpu().contiguous().data_ptr<int>(), target);
 73 | 	torch::TensorOptions intoptions = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
 74 | 	return torch::from_blob(indices.data(), {(int)indices.size()}, intoptions).clone();
 75 | }
 76 | 
 77 | int ExpandToSize(
 78 | torch::Tensor& nodes, 
 79 | torch::Tensor& boxes, 
 80 | float size, 
 81 | torch::Tensor& viewpoint, 
 82 | torch::Tensor& viewdir, 
 83 | torch::Tensor& render_indices,
 84 | torch::Tensor& parent_indices,
 85 | torch::Tensor& nodes_for_render_indices)
 86 | {
 87 | 	return Switching::expandToSize(
 88 | 	nodes.size(0), 
 89 | 	size,
 90 | 	nodes.contiguous().data_ptr<int>(), 
 91 | 	boxes.contiguous().data_ptr<float>(),
 92 | 	viewpoint.contiguous().data_ptr<float>(),
 93 | 	viewdir.data_ptr<float>()[0], viewdir.data_ptr<float>()[1], viewdir.data_ptr<float>()[2],
 94 | 	render_indices.contiguous().data_ptr<int>(),
 95 | 	nullptr,
 96 | 	parent_indices.contiguous().data_ptr<int>(),
 97 | 	nodes_for_render_indices.contiguous().data_ptr<int>());
 98 | }
 99 | 
100 | void GetTsIndexed(
101 | torch::Tensor& indices,
102 | float size,
103 | torch::Tensor& nodes,
104 | torch::Tensor& boxes,
105 | torch::Tensor& viewpoint, 
106 | torch::Tensor& viewdir, 
107 | torch::Tensor& ts,
108 | torch::Tensor& num_kids)
109 | {
110 | 	Switching::getTsIndexed(
111 | 	indices.size(0),
112 | 	indices.contiguous().data_ptr<int>(),
113 | 	size,
114 | 	nodes.contiguous().data_ptr<int>(), 
115 | 	boxes.contiguous().data_ptr<float>(),
116 | 	viewpoint.data_ptr<float>()[0], viewpoint.data_ptr<float>()[1], viewpoint.data_ptr<float>()[2],
117 | 	viewdir.data_ptr<float>()[0], viewdir.data_ptr<float>()[1], viewdir.data_ptr<float>()[2],
118 | 	ts.contiguous().data_ptr<float>(),
119 | 	num_kids.contiguous().data_ptr<int>(), 0);
120 | }


--------------------------------------------------------------------------------
/torch/torch_interface.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | #include <torch/extension.h>
 3 | #include <cstdio>
 4 | #include <tuple>
 5 | #include <string>
 6 | 
 7 | std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 8 | LoadHierarchy(std::string filename);
 9 | 
10 | void WriteHierarchy(
11 | 					std::string filename,
12 | 					torch::Tensor& pos,
13 | 					torch::Tensor& shs,
14 | 					torch::Tensor& opacities,
15 | 					torch::Tensor& log_scales,
16 | 					torch::Tensor& rotations,
17 | 					torch::Tensor& nodes,
18 | 					torch::Tensor& boxes);
19 | 
20 | torch::Tensor
21 | ExpandToTarget(torch::Tensor& nodes, int target);
22 | 
23 | int ExpandToSize(
24 | torch::Tensor& nodes, 
25 | torch::Tensor& boxes, 
26 | float size, 
27 | torch::Tensor& viewpoint, 
28 | torch::Tensor& viewdir, 
29 | torch::Tensor& render_indices,
30 | torch::Tensor& parent_indices,
31 | torch::Tensor& nodes_for_render_indices);
32 | 
33 | void GetTsIndexed(
34 | torch::Tensor& indices,
35 | float size,
36 | torch::Tensor& nodes,
37 | torch::Tensor& boxes,
38 | torch::Tensor& viewpoint, 
39 | torch::Tensor& viewdir, 
40 | torch::Tensor& ts,
41 | torch::Tensor& num_kids);


--------------------------------------------------------------------------------
/traversal.cpp:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | 
13 | #include "traversal.h"
14 | 
15 | void recExpand(int node_id, Node* nodes, std::vector<int>& indices, int target)
16 | {
17 | 	Node& node = nodes[node_id];
18 | 
19 | 	for (int i = 0; i < node.count_leafs; i++)
20 | 		indices.push_back(node.start + i);
21 | 
22 | 	if (node.depth <= target) // We are below target. Nodes will not be expanded, content approximated
23 | 	{
24 | 		for (int i = 0; i < node.count_merged; i++)
25 | 			indices.push_back(node.start + node.count_leafs + i);
26 | 	}
27 | 	else // we keep expanding and adding visible leaves
28 | 	{
29 | 		for (int i = 0; i < node.count_children; i++)
30 | 			recExpand(node.start_children + i, nodes, indices, target);
31 | 	}
32 | }
33 | 
34 | std::vector<int> Traversal::expandToTarget(Node* nodes, int target)
35 | {
36 | 	std::vector<int> indices;
37 | 	recExpand(0, nodes, indices, target);
38 | 	return indices;
39 | }


--------------------------------------------------------------------------------
/traversal.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include <vector>
15 | #include "common.h"
16 | 
17 | class Traversal
18 | {
19 | public:
20 | 	static std::vector<int>  expandToTarget(Node* nodes, int target);
21 | };


--------------------------------------------------------------------------------
/types.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #ifndef __CUDACC__
15 | 
16 | #include <Eigen/Core>
17 | 
18 | struct Box
19 | {
20 | 	Box(Eigen::Vector3f minn, Eigen::Vector3f maxx) : minn(minn.x(), minn.y(), minn.z(), 0), maxx(maxx.x(), maxx.y(), maxx.z(), 0)
21 | 	{}
22 | 
23 | 	Box() {};
24 | 
25 | 	Eigen::Vector4f minn;
26 | 	Eigen::Vector4f maxx;
27 | };
28 | 
29 | typedef Eigen::Matrix<float, 48, 1> SHs;
30 | 
31 | #else
32 | 
33 | struct Point
34 | {
35 | 	float xyz[3];
36 | };
37 | 
38 | struct Point4
39 | {
40 | 	float xyz[4];
41 | };
42 | 
43 | struct Box
44 | {
45 | 	Point4 minn;
46 | 	Point4 maxx;
47 | };
48 | 
49 | #endif
50 | 
51 | struct Node
52 | {
53 | 	int depth = -1;
54 | 	int parent = -1;
55 | 	int start;
56 | 	int count_leafs;
57 | 	int count_merged;
58 | 	int start_children;
59 | 	int count_children;
60 | };
61 | 
62 | struct HalfNode
63 | {
64 | 	int parent = -1;
65 | 	int start;
66 | 	int start_children;
67 | 	short dccc[4];
68 | };


--------------------------------------------------------------------------------
/writer.cpp:
--------------------------------------------------------------------------------
  1 | /*
  2 |  * Copyright (C) 2024, 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 | 
 13 | #include "writer.h"
 14 | #include <iostream>
 15 | #include <fstream>
 16 | #include "hierarchy_writer.h"
 17 | #include <map>
 18 | 
 19 | void populateRec(
 20 | 	const ExplicitTreeNode* treenode,
 21 | 	int id,
 22 | 	const std::vector<Gaussian>& gaussians, 
 23 | 	std::vector<Eigen::Vector3f>& positions,
 24 | 	std::vector<Eigen::Vector4f>& rotations,
 25 | 	std::vector<Eigen::Vector3f>& log_scales,
 26 | 	std::vector<float>& opacities,
 27 | 	std::vector<SHs>& shs,
 28 | 	std::vector<Node>& basenodes,
 29 | 	std::vector<Box>& boxes,
 30 | 	std::map<int, const ExplicitTreeNode*>* base2tree = nullptr)
 31 | {
 32 | 	if(base2tree)
 33 | 		base2tree->insert(std::make_pair(id, treenode));
 34 | 
 35 | 	boxes[id] = treenode->bounds;
 36 | 	basenodes[id].start = positions.size();
 37 | 	for (auto& i : treenode->leaf_indices)
 38 | 	{
 39 | 		const Gaussian& g = gaussians[i];
 40 | 		positions.push_back(g.position);
 41 | 		rotations.push_back(g.rotation);
 42 | 		log_scales.push_back(g.scale.array().log());
 43 | 		opacities.push_back(g.opacity);
 44 | 		shs.push_back(g.shs);
 45 | 	}
 46 | 	basenodes[id].count_leafs = treenode->leaf_indices.size();
 47 | 
 48 | 	for (auto& g : treenode->merged)
 49 | 	{
 50 | 		positions.push_back(g.position);
 51 | 		rotations.push_back(g.rotation);
 52 | 		log_scales.push_back(g.scale.array().log());
 53 | 		opacities.push_back(g.opacity);
 54 | 		shs.push_back(g.shs);
 55 | 	}
 56 | 	basenodes[id].count_merged = treenode->merged.size();
 57 | 
 58 | 	basenodes[id].start_children = basenodes.size();
 59 | 	for (int n = 0; n < treenode->children.size(); n++)
 60 | 	{
 61 | 		basenodes.push_back(Node());
 62 | 		basenodes.back().parent = id;
 63 | 		boxes.push_back(Box());
 64 | 	}
 65 | 	basenodes[id].count_children = treenode->children.size();
 66 | 
 67 | 	basenodes[id].depth = treenode->depth;
 68 | 
 69 | 	for (int n = 0; n < treenode->children.size(); n++)
 70 | 	{
 71 | 		populateRec(
 72 | 			treenode->children[n],
 73 | 			basenodes[id].start_children + n,
 74 | 			gaussians, 
 75 | 			positions, 
 76 | 			rotations, 
 77 | 			log_scales, 
 78 | 			opacities,
 79 | 			shs, 
 80 | 			basenodes,
 81 | 			boxes,
 82 | 			base2tree);
 83 | 	}
 84 | }
 85 | 
 86 | void recTraverse(int id, std::vector<Node>& nodes, int& count)
 87 | {
 88 | 	if (nodes[id].depth == 0)
 89 | 		count++;
 90 | 	if (nodes[id].count_children != 0 && nodes[id].depth == 0)
 91 | 		throw std::runtime_error("An error occurred in traversal");
 92 | 	for (int i = 0; i < nodes[id].count_children; i++)
 93 | 	{
 94 | 		recTraverse(nodes[id].start_children + i, nodes, count);
 95 | 	}
 96 | }
 97 | 
 98 | void Writer::makeHierarchy(
 99 | 	const std::vector<Gaussian>& gaussians,
100 | 	const ExplicitTreeNode* root,
101 | 	std::vector<Eigen::Vector3f>& positions,
102 | 	std::vector<Eigen::Vector4f>& rotations,
103 | 	std::vector<Eigen::Vector3f>& log_scales,
104 | 	std::vector<float>& opacities,
105 | 	std::vector<SHs>& shs,
106 | 	std::vector<Node>& basenodes,
107 | 	std::vector<Box>& boxes,
108 | 	std::map<int, const ExplicitTreeNode*>* base2tree)
109 | {
110 | 	basenodes.resize(1);
111 | 	boxes.resize(1);
112 | 
113 | 	populateRec(
114 | 		root,
115 | 		0,
116 | 		gaussians,
117 | 		positions, rotations, log_scales, opacities, shs, basenodes, boxes,
118 | 		base2tree);
119 | }
120 | 
121 | void Writer::writeHierarchy(const char* filename, const std::vector<Gaussian>& gaussians, const ExplicitTreeNode* root, bool compressed)
122 | {
123 | 	std::vector<Eigen::Vector3f> positions;
124 | 	std::vector<Eigen::Vector4f> rotations;
125 | 	std::vector<Eigen::Vector3f> log_scales;
126 | 	std::vector<float> opacities;
127 | 	std::vector<SHs> shs;
128 | 	std::vector<Node> basenodes;
129 | 	std::vector<Box> boxes;
130 | 
131 | 	makeHierarchy(gaussians, root, positions, rotations, log_scales, opacities, shs, basenodes, boxes);
132 | 
133 | 	HierarchyWriter writer;
134 | 	writer.write(
135 | 		filename,
136 | 		positions.size(),
137 | 		basenodes.size(),
138 | 		positions.data(),
139 | 		shs.data(),
140 | 		opacities.data(),
141 | 		log_scales.data(),
142 | 		rotations.data(),
143 | 		basenodes.data(),
144 | 		boxes.data(),
145 | 		compressed
146 | 	);
147 | }


--------------------------------------------------------------------------------
/writer.h:
--------------------------------------------------------------------------------
 1 | /*
 2 |  * Copyright (C) 2024, 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 | #pragma once
13 | 
14 | #include "common.h"
15 | #include <map>
16 | 
17 | class Writer
18 | {
19 | public:
20 | 	static void writeHierarchy(const char* filename, const std::vector<Gaussian>& gaussian, const ExplicitTreeNode* root, bool compressed = true);
21 | 
22 | 	static void makeHierarchy(
23 | 		const std::vector<Gaussian>& gaussians,
24 | 		const ExplicitTreeNode* root,
25 | 		std::vector<Eigen::Vector3f>& positions,
26 | 		std::vector<Eigen::Vector4f>& rotations,
27 | 		std::vector<Eigen::Vector3f>& log_scales,
28 | 		std::vector<float>& opacities,
29 | 		std::vector<SHs>& shs,
30 | 		std::vector<Node>& basenodes,
31 | 		std::vector<Box>& boxes,
32 | 		std::map<int, const ExplicitTreeNode*>* base2tree = nullptr);
33 | };


--------------------------------------------------------------------------------