OpenSplat/opensplat.cpp at main · YHZN-REPOS/OpenSplat · GitHub

Repository navigation

History

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

#include <filesystem>

#include <nlohmann/json.hpp>

#include "opensplat.hpp"

#include "input_data.hpp"

#include "utils.hpp"

#include "cv_utils.hpp"

#include "constants.hpp"

#include "zip_utils.hpp"

#include <cxxopts.hpp>

#ifdef USE_VISUALIZATION

#include "visualizer.hpp"

#endif

namespace fs = std::filesystem;

using namespace torch::indexing;

int main(int argc, char *argv[]){

cxxopts::Options options("opensplat", "Open Source 3D Gaussian Splats generator - " APP_VERSION);

options.add_options()

("i,input", "Path to nerfstudio project", cxxopts::value<std::string>())

("o,output", "Path where to save output scene (default: splat.ply next to the input)", cxxopts::value<std::string>()->default_value("splat.ply"))

("output-cameras", "Path where to save a cameras JSON file", cxxopts::value<std::string>()->default_value(""))

("s,save-every", "Save output scene every these many steps (set to -1 to disable)", cxxopts::value<int>()->default_value("-1"))

("resume", "Resume training from this PLY file", cxxopts::value<std::string>()->default_value(""))

("val", "Withhold a camera shot for validating the scene loss")

("val-image", "Filename of the image to withhold for validating scene loss", cxxopts::value<std::string>()->default_value("random"))

("val-render", "Path of the directory where to render validation images", cxxopts::value<std::string>()->default_value(""))

("center", "Center the model at the origin")

("cpu", "Force CPU execution")

("n,num-iters", "Number of iterations to run", cxxopts::value<int>()->default_value("30000"))

("d,downscale-factor", "Scale input images by this factor.", cxxopts::value<float>()->default_value("1"))

("num-downscales", "Number of images downscales to use. After being scaled by [downscale-factor], images are initially scaled by a further (2^[num-downscales]) and the scale is increased every [resolution-schedule]", cxxopts::value<int>()->default_value("0"))

("resolution-schedule", "Double the image resolution every these many steps", cxxopts::value<int>()->default_value("3000"))

("sh-degree", "Maximum spherical harmonics degree (must be > 0)", cxxopts::value<int>()->default_value("3"))

("sh-degree-interval", "Increase the number of spherical harmonics degree after these many steps (will not exceed [sh-degree])", cxxopts::value<int>()->default_value("1000"))

("ssim-weight", "Weight to apply to the structural similarity loss. Set to zero to use least absolute deviation (L1) loss only", cxxopts::value<float>()->default_value("0.2"))

("refine-every", "Densify/prune gaussians every these many steps", cxxopts::value<int>()->default_value("500"))

("densify-from", "Start densifying gaussians after these many steps", cxxopts::value<int>()->default_value("500"))

("densify-until", "Stop densifying gaussians after these many steps (-1 = min(15000, half of num-iters))", cxxopts::value<int>()->default_value("-1"))

("loss-thresh", "High-error pixel threshold on the normalized L1 map for multi-view scoring", cxxopts::value<float>()->default_value("0.1"))

("no-edge-guidance", "Disable Canny edge weighting of the densification importance", cxxopts::value<bool>()->default_value("false"))

("max-gaussians", "Maximum number of gaussians (0 = unlimited)", cxxopts::value<int>()->default_value("5000000"))

("no-masks", "Ignore image masks even when present", cxxopts::value<bool>()->default_value("false"))

("no-gpu-cache", "Do not cache images/masks on the GPU (reduces VRAM usage, slower)", cxxopts::value<bool>()->default_value("false"))

("colmap-image-path", "Override the default image path for COLMAP-based input", cxxopts::value<std::string>()->default_value(""))

("opensfm-image-path", "Override the default image path for OpenSfM-based input", cxxopts::value<std::string>()->default_value(""))

#ifdef USE_VISUALIZATION

("has-visualization", "Show the visualization steps of training", cxxopts::value<bool>()->default_value("0"))

#endif

("h,help", "Print usage")

("version", "Print version")

;

options.parse_positional({ "input" });

options.positional_help("[colmap/nerfstudio/opensfm/odx/openmvg project path or .zip archive]");

cxxopts::ParseResult result;

try {

result = options.parse(argc, argv);

}

catch (const std::exception &e) {

std::cerr << e.what() << std::endl;

std::cerr << options.help() << std::endl;

return EXIT_FAILURE;

}

if (result.count("version")){

std::cout << APP_VERSION << std::endl;

return EXIT_SUCCESS;

}

if (result.count("help") || !result.count("input")) {

std::cout << options.help() << std::endl;

return EXIT_SUCCESS;

}

const std::string projectRoot = result["input"].as<std::string>();

std::string outputScene = result["output"].as<std::string>();

if (result.count("output") == 0){

// Output next to the input scene, for a .zip, next to the archive

fs::path in = fs::absolute(fs::path(projectRoot));

if (in.filename().empty()) in = in.parent_path();

outputScene = (in.parent_path() / outputScene).string();

}

const std::string outputCameras = result["output-cameras"].as<std::string>();

const int saveEvery = result["save-every"].as<int>();

const std::string resume = result["resume"].as<std::string>();

const bool validate = result.count("val") > 0 || result.count("val-render") > 0;

const std::string valImage = result["val-image"].as<std::string>();

const std::string valRender = result["val-render"].as<std::string>();

if (!valRender.empty() && !fs::exists(valRender)) fs::create_directories(valRender);

const bool keepCrs = result.count("center") == 0;

const float downScaleFactor = (std::max)(result["downscale-factor"].as<float>(), 1.0f);

const int numIters = result["num-iters"].as<int>();

const int numDownscales = result["num-downscales"].as<int>();

const int resolutionSchedule = result["resolution-schedule"].as<int>();

const int shDegree = result["sh-degree"].as<int>();

const int shDegreeInterval = result["sh-degree-interval"].as<int>();

const float ssimWeight = result["ssim-weight"].as<float>();

const int refineEvery = result["refine-every"].as<int>();

const int densifyFrom = result["densify-from"].as<int>();

int densifyUntil = result["densify-until"].as<int>();

if (densifyUntil < 0) densifyUntil = (std::min)(15000, result["num-iters"].as<int>() / 2);

const float lossThresh = result["loss-thresh"].as<float>();

const int maxGaussians = result["max-gaussians"].as<int>();

const bool noMasks = result["no-masks"].as<bool>();

const std::string colmapImageSourcePath = result["colmap-image-path"].as<std::string>();

const std::string opensfmImageSourcePath = result["opensfm-image-path"].as<std::string>();

#ifdef USE_VISUALIZATION

const bool hasVisualization = result["has-visualization"].as<bool>();

#endif

torch::Device device = torch::kCPU;

int displayStep = 10;

if (torch::hasCUDA() && result.count("cpu") == 0) {

std::cout << "Using CUDA" << std::endl;

device = torch::kCUDA;

} else if (torch::hasMPS() && result.count("cpu") == 0) {

std::cout << "Using MPS" << std::endl;

device = torch::kMPS;

}else{

std::cout << "Using CPU" << std::endl;

displayStep = 1;

}

#ifdef USE_VISUALIZATION

Visualizer visualizer;

if (hasVisualization)

visualizer.Initialize(numIters);

#endif

try{

std::string projectPath = projectRoot;

if (isZipArchive(projectRoot)) projectPath = extractZipToCache(projectRoot);

InputData inputData = inputDataFromX(projectPath, colmapImageSourcePath, opensfmImageSourcePath);

int numMasks = 0;

if (!noMasks){

for (Camera &cam : inputData.cameras){

cam.maskPath = findMaskPath(cam.filePath, projectPath);

if (!cam.maskPath.empty()) numMasks++;

}

if (numMasks > 0) std::cout << "Found " << numMasks << " masks" << std::endl;

parallel_for(inputData.cameras.begin(), inputData.cameras.end(), [&downScaleFactor](Camera &cam){

cam.loadImage(downScaleFactor);

});

// Withhold a validation camera if necessary

auto t = inputData.getCameras(validate, valImage);

std::vector<Camera> cams = std::get<0>(t);

Camera *valCam = std::get<1>(t);

// Discard cameras that see none of the sparse points

{

torch::Tensor xyz = inputData.points.xyz;

std::vector<Camera> kept;

kept.reserve(cams.size());

for (Camera &cam : cams){

torch::Tensor R = cam.camToWorld.index({Slice(None, 3), Slice(None, 3)});

torch::Tensor T = cam.camToWorld.index({Slice(None, 3), Slice(3, 4)});

R = torch::matmul(R, torch::diag(torch::tensor({1.0f, -1.0f, -1.0f})));

torch::Tensor pCam = torch::matmul(xyz - T.transpose(0, 1), R);

torch::Tensor z = pCam.index({Slice(), 2});

torch::Tensor px = pCam.index({Slice(), 0}) / z * cam.fx + cam.cx;

torch::Tensor py = pCam.index({Slice(), 1}) / z * cam.fy + cam.cy;

long long inView = ((z > 0.01f) & (px >= 0.0f) & (px < static_cast<float>(cam.width)) &

(py >= 0.0f) & (py < static_cast<float>(cam.height))).sum().item<int64_t>();

if (inView > 0){

kept.push_back(cam);

}else{

std::cout << "Discarding " << cam.filePath << " (no points in view)" << std::endl;

}

if (kept.empty()) throw std::runtime_error("No cameras see any sparse points");

cams = std::move(kept);

}

Model model(inputData,

cams.size(),

numDownscales, resolutionSchedule, shDegree, shDegreeInterval,

refineEvery, densifyFrom, densifyUntil, maxGaussians,

lossThresh,

numIters, keepCrs,

device);

model.trainCams = &cams;

model.edgeGuidance = !result["no-edge-guidance"].as<bool>();

Camera::gpuCacheEnabled = !result["no-gpu-cache"].as<bool>();

std::vector< size_t > camIndices( cams.size() );

std::iota( camIndices.begin(), camIndices.end(), 0 );

InfiniteRandomIterator<size_t> camsIter( camIndices );

int imageSize = -1;

size_t step = 1;

if (resume != ""){

step = model.loadPly(resume) + 1;

}

for (; step <= numIters; step++){

Camera& cam = cams[ camsIter.next() ];

torch::Tensor rgb = model.forward(cam, step);

torch::Tensor gt = cam.getImageGpu(model.getDownscaleFactor(step), device);

torch::Tensor mask = cam.getMaskGpu(model.getDownscaleFactor(step), device);

torch::Tensor mainLoss = model.mainLoss(rgb, gt, mask, ssimWeight);

mainLoss.backward();

if (step % displayStep == 0) {

const float percentage = static_cast<float>(step) / numIters;

std::cout << "Step " << step << ": " << mainLoss.item<float>() << " [" << floor(percentage * 100) << "%]" << std::endl;

}

model.afterTrain(step);

model.optimizerStepCadence(step);

model.schedulersStep(step);

if (saveEvery > 0 && step % saveEvery == 0){

fs::path p(outputScene);

model.save(p.replace_filename(fs::path(p.stem().string() + "_" + std::to_string(step) + p.extension().string())).string(), step);

}

if (!valRender.empty() && step % 10 == 0){

torch::Tensor rgb = model.forward(*valCam, step);

cv::Mat image = tensorToImage(rgb.detach().cpu());

cv::cvtColor(image, image, cv::COLOR_RGB2BGR);

cv::imwrite((fs::path(valRender) / (std::to_string(step) + ".png")).string(), image);

}

#ifdef USE_VISUALIZATION

if (hasVisualization) {

visualizer.SetInitialGaussianNum(inputData.points.xyz.size(0));

visualizer.SetLoss(step, mainLoss.item<float>());

visualizer.SetGaussians(model.means, model.scales, model.featuresDc,

model.opacities);

visualizer.SetImage(rgb, gt);

if (visualizer.QuitApp())

step = numIters + 1;

visualizer.Draw();

}

#endif

}

if (!outputCameras.empty()) inputData.saveCameras(outputCameras, keepCrs);

model.save(outputScene, numIters);

// model.saveDebugPly("debug.ply", numIters);

// Validate

if (valCam != nullptr){

torch::Tensor rgb = model.forward(*valCam, numIters);

torch::Tensor gt = valCam->getImageGpu(model.getDownscaleFactor(numIters), device);

torch::Tensor valMask = valCam->getMaskGpu(model.getDownscaleFactor(numIters), device);

std::cout << valCam->filePath << " validation loss: " << model.mainLoss(rgb, gt, valMask, ssimWeight).item<float>() << std::endl;

torch::Tensor mse;

if (valMask.defined() && valMask.numel() > 0){

mse = (valMask.unsqueeze(-1) * (rgb - gt).pow(2)).sum() / (valMask.sum() * gt.size(2) + 1e-8f);

}else{

mse = (rgb - gt).pow(2).mean();

}

std::cout << valCam->filePath << " validation PSNR: " << (10.0f * torch::log10(1.0f / mse)).item<float>() << std::endl;

}

}catch(const std::exception &e){

std::cerr << e.what() << std::endl;

exit(1);

}

OpenSplat/opensplat.cpp at main · YHZN-REPOS/OpenSplat · GitHub

Repository navigation

Expand file tree

Breadcrumbs

opensplat.cpp

Latest commit

History

Breadcrumbs

opensplat.cpp

File metadata and controls