salmanoff/buildmach/meshFromPcd.cpp

#include <chrono>
#include <cstdint>
#include <cmath>
#include <filesystem>
#include <iomanip>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

#include <pcl/PolygonMesh.h>
#include <pcl/common/io.h>
#include <pcl/features/normal_3d.h>
#include <pcl/io/pcd_io.h>
#include <pcl/io/vtk_io.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/search/kdtree.h>
#include <pcl/surface/gp3.h>
#include <pcl/surface/organized_fast_mesh.h>

namespace {

enum class MeshAlgorithm
{
	Ofm,
	Gp3,
};

struct ToolOptions
{
	std::vector<std::filesystem::path> inputPaths;
	std::filesystem::path outputDirectory;
	bool hasOutputDirectory = false;
	MeshAlgorithm algorithm = MeshAlgorithm::Ofm;
	int ofmTrianglePixelSize = 1;
	float ofmMaxEdgeLength = 0.25f;
	pcl::OrganizedFastMesh<pcl::PointXYZ>::TriangulationType ofmTriangulationType =
		pcl::OrganizedFastMesh<pcl::PointXYZ>::TRIANGLE_ADAPTIVE_CUT;
	double gp3SearchRadius = 0.05;
	double gp3Mu = 2.5;
	int gp3MaxNeighbors = 100;
	double gp3MaxSurfaceAngleDeg = 45.0;
	double gp3MinAngleDeg = 10.0;
	double gp3MaxAngleDeg = 120.0;
	int gp3NormalK = 20;
	bool gp3NormalConsistency = false;
};

struct ConversionStats
{
	std::filesystem::path inputPath;
	std::filesystem::path outputPath;
	MeshAlgorithm algorithm = MeshAlgorithm::Ofm;
	std::size_t pointCount = 0;
	std::size_t finitePointCount = 0;
	std::size_t polygonCount = 0;
	double normalDurationMs = 0.0;
	double meshDurationMs = 0.0;
	double totalDurationMs = 0.0;
	bool success = false;
	std::string errorMessage;
};

constexpr double kPi = 3.14159265358979323846;
constexpr std::size_t kMinimumTriangulationPoints = 3;

double degreesToRadians(double degrees)
{
	return degrees * kPi / 180.0;
}

std::string algorithmToString(MeshAlgorithm algorithm)
{
	switch (algorithm)
	{
	case MeshAlgorithm::Ofm:
		return "ofm";
	case MeshAlgorithm::Gp3:
		return "gp3";
	}

	throw std::runtime_error("Unsupported mesh algorithm enum");
}

void printUsage(const char* argv0)
{
	std::cerr
		<< "Usage: " << argv0
		<< " [--algorithm ofm|gp3]"
		<< " [--output-dir DIR]"
		<< " [--ofm-triangle-pixel-size N]"
		<< " [--ofm-max-edge-length F]"
		<< " [--ofm-triangulation adaptive|left|right|quad]"
		<< " [--gp3-search-radius F]"
		<< " [--gp3-mu F]"
		<< " [--gp3-max-neighbors N]"
		<< " [--gp3-max-surface-angle-deg F]"
		<< " [--gp3-min-angle-deg F]"
		<< " [--gp3-max-angle-deg F]"
		<< " [--gp3-normal-k N]"
		<< " [--gp3-normal-consistency true|false]"
		<< " input1.pcd [input2.pcd ...]\n";
}

bool parseBooleanValue(const std::string& value)
{
	if (value == "true" || value == "1")
	{
		return true;
	}
	if (value == "false" || value == "0")
	{
		return false;
	}

	throw std::runtime_error(
		"Expected boolean value true|false|1|0, got: " + value);
}

MeshAlgorithm parseAlgorithm(const std::string& value)
{
	if (value == "ofm")
	{
		return MeshAlgorithm::Ofm;
	}
	if (value == "gp3")
	{
		return MeshAlgorithm::Gp3;
	}

	throw std::runtime_error("Unsupported algorithm: " + value);
}

bool parseTriangulationType(
	const std::string& value,
	pcl::OrganizedFastMesh<pcl::PointXYZ>::TriangulationType& ofmTriangulationType)
{
	if (value == "adaptive")
	{
		ofmTriangulationType =
			pcl::OrganizedFastMesh<pcl::PointXYZ>::TRIANGLE_ADAPTIVE_CUT;
		return true;
	}
	if (value == "left")
	{
		ofmTriangulationType =
			pcl::OrganizedFastMesh<pcl::PointXYZ>::TRIANGLE_LEFT_CUT;
		return true;
	}
	if (value == "right")
	{
		ofmTriangulationType =
			pcl::OrganizedFastMesh<pcl::PointXYZ>::TRIANGLE_RIGHT_CUT;
		return true;
	}
	if (value == "quad")
	{
		ofmTriangulationType =
			pcl::OrganizedFastMesh<pcl::PointXYZ>::QUAD_MESH;
		return true;
	}

	return false;
}

void parseToolOption(
	ToolOptions& options,
	const std::string& arg,
	int& i,
	int argc,
	char** argv)
{
	auto requireValue = [&](const char* optionName) -> std::string {
		if (i + 1 >= argc)
		{
			throw std::runtime_error(std::string(optionName) + " requires a value");
		}
		return argv[++i];
	};

	if (arg == "--algorithm")
	{
		options.algorithm = parseAlgorithm(requireValue("--algorithm"));
		return;
	}
	if (arg == "--output-dir")
	{
		options.outputDirectory = requireValue("--output-dir");
		options.hasOutputDirectory = true;
		return;
	}
	if (arg == "--ofm-triangle-pixel-size")
	{
		options.ofmTrianglePixelSize = std::stoi(
			requireValue("--ofm-triangle-pixel-size"));
		return;
	}
	if (arg == "--ofm-max-edge-length")
	{
		options.ofmMaxEdgeLength = std::stof(
			requireValue("--ofm-max-edge-length"));
		return;
	}
	if (arg == "--ofm-triangulation")
	{
		std::string triangulationValue = requireValue("--ofm-triangulation");
		if (!parseTriangulationType(
			triangulationValue, options.ofmTriangulationType))
		{
			throw std::runtime_error(
				"Unsupported triangulation type: " + triangulationValue);
		}
		return;
	}
	if (arg == "--gp3-search-radius")
	{
		options.gp3SearchRadius = std::stod(
			requireValue("--gp3-search-radius"));
		return;
	}
	if (arg == "--gp3-mu")
	{
		options.gp3Mu = std::stod(requireValue("--gp3-mu"));
		return;
	}
	if (arg == "--gp3-max-neighbors")
	{
		options.gp3MaxNeighbors = std::stoi(
			requireValue("--gp3-max-neighbors"));
		return;
	}
	if (arg == "--gp3-max-surface-angle-deg")
	{
		options.gp3MaxSurfaceAngleDeg = std::stod(
			requireValue("--gp3-max-surface-angle-deg"));
		return;
	}
	if (arg == "--gp3-min-angle-deg")
	{
		options.gp3MinAngleDeg = std::stod(
			requireValue("--gp3-min-angle-deg"));
		return;
	}
	if (arg == "--gp3-max-angle-deg")
	{
		options.gp3MaxAngleDeg = std::stod(
			requireValue("--gp3-max-angle-deg"));
		return;
	}
	if (arg == "--gp3-normal-k")
	{
		options.gp3NormalK = std::stoi(requireValue("--gp3-normal-k"));
		return;
	}
	if (arg == "--gp3-normal-consistency")
	{
		options.gp3NormalConsistency = parseBooleanValue(
			requireValue("--gp3-normal-consistency"));
		return;
	}

	throw std::runtime_error("Unknown option: " + arg);
}

ToolOptions parseArgs(int argc, char** argv)
{
	ToolOptions options;

	for (int i = 1; i < argc; ++i)
	{
		std::string arg = argv[i];
		if (!arg.empty() && arg[0] == '-')
		{
			parseToolOption(options, arg, i, argc, argv);
			continue;
		}

		options.inputPaths.emplace_back(arg);
	}

	if (options.inputPaths.empty())
	{
		throw std::runtime_error("At least one input .pcd file is required");
	}

	return options;
}

std::filesystem::path computeOutputPath(
	const ToolOptions& options,
	const std::filesystem::path& inputPath)
{
	std::filesystem::path outputDirectory = options.hasOutputDirectory
		? options.outputDirectory
		: inputPath.parent_path();
	std::filesystem::path outputFileName = inputPath.filename();
	outputFileName.replace_extension(".vtk");
	return outputDirectory / outputFileName;
}

pcl::PointCloud<pcl::PointXYZ>::Ptr loadInputCloud(
	const std::filesystem::path& inputPath)
{
	pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(
		new pcl::PointCloud<pcl::PointXYZ>());
	if (pcl::io::loadPCDFile(inputPath.string(), *cloud) != 0)
	{
		throw std::runtime_error("Failed to load input PCD");
	}

	return cloud;
}

void ensureOutputDirectoryExists(const std::filesystem::path& outputPath)
{
	std::filesystem::create_directories(outputPath.parent_path());
}

void ensureOrganizedCloudForOfm(const pcl::PointCloud<pcl::PointXYZ>& cloud)
{
	if (!cloud.isOrganized())
	{
		throw std::runtime_error("Input point cloud is not organized");
	}
}

pcl::PointCloud<pcl::PointXYZ>::Ptr buildFinitePointCloud(
	const pcl::PointCloud<pcl::PointXYZ>& cloud)
{
	pcl::PointCloud<pcl::PointXYZ>::Ptr finiteCloud(
		new pcl::PointCloud<pcl::PointXYZ>());
	finiteCloud->reserve(cloud.size());

	for (const auto& point : cloud.points)
	{
		if (!std::isfinite(point.x) || !std::isfinite(point.y) ||
			!std::isfinite(point.z))
		{
			continue;
		}

		finiteCloud->push_back(point);
	}

	finiteCloud->width = static_cast<std::uint32_t>(finiteCloud->size());
	finiteCloud->height = 1;
	finiteCloud->is_dense = false;
	return finiteCloud;
}

void ensureEnoughFinitePoints(
	const pcl::PointCloud<pcl::PointXYZ>& finiteCloud,
	int normalK)
{
	if (finiteCloud.size() < kMinimumTriangulationPoints)
	{
		throw std::runtime_error("Too few finite points to triangulate");
	}

	if (finiteCloud.size() <= static_cast<std::size_t>(normalK))
	{
		throw std::runtime_error(
			"Too few finite points for requested GP3 normal-k");
	}
}

pcl::PointCloud<pcl::Normal>::Ptr estimateNormals(
	const pcl::PointCloud<pcl::PointXYZ>::Ptr& finiteCloud,
	int normalK)
{
	pcl::search::KdTree<pcl::PointXYZ>::Ptr searchTree(
		new pcl::search::KdTree<pcl::PointXYZ>());
	searchTree->setInputCloud(finiteCloud);

	pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> normalEstimation;
	normalEstimation.setInputCloud(finiteCloud);
	normalEstimation.setSearchMethod(searchTree);
	normalEstimation.setKSearch(normalK);
	normalEstimation.setViewPoint(0.0f, 0.0f, 0.0f);

	pcl::PointCloud<pcl::Normal>::Ptr normals(new pcl::PointCloud<pcl::Normal>());
	normalEstimation.compute(*normals);
	return normals;
}

pcl::PointCloud<pcl::PointNormal>::Ptr buildPointNormalsCloud(
	const pcl::PointCloud<pcl::PointXYZ>& finiteCloud,
	const pcl::PointCloud<pcl::Normal>& normals)
{
	pcl::PointCloud<pcl::PointNormal>::Ptr pointNormals(
		new pcl::PointCloud<pcl::PointNormal>());
	pcl::concatenateFields(finiteCloud, normals, *pointNormals);
	return pointNormals;
}

void ensureGp3NormalsAreFinite(
	const pcl::PointCloud<pcl::PointNormal>& pointNormals)
{
	for (const auto& point : pointNormals.points)
	{
		if (std::isfinite(point.normal_x) && std::isfinite(point.normal_y) &&
			std::isfinite(point.normal_z))
		{
			return;
		}
	}

	throw std::runtime_error("Normal estimation produced no finite normals");
}

void saveMeshOrThrow(
	const std::filesystem::path& outputPath,
	const pcl::PolygonMesh& mesh)
{
	if (mesh.polygons.empty())
	{
		throw std::runtime_error("Reconstruction produced no polygons");
	}

	if (pcl::io::saveVTKFile(outputPath.string(), mesh) != 0)
	{
		throw std::runtime_error("Failed to save output VTK");
	}
}

void configureOfm(
	pcl::OrganizedFastMesh<pcl::PointXYZ>& ofm,
	const ToolOptions& options,
	const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud)
{
	ofm.setInputCloud(cloud);
	ofm.setTriangulationType(options.ofmTriangulationType);
	ofm.setTrianglePixelSize(options.ofmTrianglePixelSize);
	ofm.setViewpoint(Eigen::Vector3f::Zero());
	ofm.setMaxEdgeLength(options.ofmMaxEdgeLength);
}

void configureGp3(
	pcl::GreedyProjectionTriangulation<pcl::PointNormal>& gp3,
	const ToolOptions& options,
	const pcl::PointCloud<pcl::PointNormal>::Ptr& pointNormals,
	const pcl::search::KdTree<pcl::PointNormal>::Ptr& searchTree)
{
	gp3.setInputCloud(pointNormals);
	gp3.setSearchMethod(searchTree);
	gp3.setSearchRadius(options.gp3SearchRadius);
	gp3.setMu(options.gp3Mu);
	gp3.setMaximumNearestNeighbors(options.gp3MaxNeighbors);
	gp3.setMaximumSurfaceAngle(
		degreesToRadians(options.gp3MaxSurfaceAngleDeg));
	gp3.setMinimumAngle(degreesToRadians(options.gp3MinAngleDeg));
	gp3.setMaximumAngle(degreesToRadians(options.gp3MaxAngleDeg));
	gp3.setNormalConsistency(options.gp3NormalConsistency);
}

ConversionStats convertWithOfm(
	const ToolOptions& options,
	const std::filesystem::path& inputPath)
{
	ConversionStats stats;
	stats.algorithm = MeshAlgorithm::Ofm;
	stats.inputPath = inputPath;
	stats.outputPath = computeOutputPath(options, inputPath);

	auto cloud = loadInputCloud(inputPath);
	ensureOrganizedCloudForOfm(*cloud);
	ensureOutputDirectoryExists(stats.outputPath);

	pcl::OrganizedFastMesh<pcl::PointXYZ> ofm;
	configureOfm(ofm, options, cloud);

	pcl::PolygonMesh mesh;
	auto meshStart = std::chrono::steady_clock::now();
	ofm.reconstruct(mesh);
	auto meshEnd = std::chrono::steady_clock::now();

	saveMeshOrThrow(stats.outputPath, mesh);

	stats.pointCount = cloud->size();
	stats.polygonCount = mesh.polygons.size();
	stats.meshDurationMs = std::chrono::duration<double, std::milli>(
		meshEnd - meshStart).count();
	stats.totalDurationMs = stats.meshDurationMs;
	stats.success = true;
	return stats;
}

ConversionStats convertWithGp3(
	const ToolOptions& options,
	const std::filesystem::path& inputPath)
{
	ConversionStats stats;
	stats.algorithm = MeshAlgorithm::Gp3;
	stats.inputPath = inputPath;
	stats.outputPath = computeOutputPath(options, inputPath);

	auto cloud = loadInputCloud(inputPath);
	auto finiteCloud = buildFinitePointCloud(*cloud);
	ensureEnoughFinitePoints(*finiteCloud, options.gp3NormalK);
	ensureOutputDirectoryExists(stats.outputPath);

	auto normalStart = std::chrono::steady_clock::now();
	auto normals = estimateNormals(finiteCloud, options.gp3NormalK);
	auto normalEnd = std::chrono::steady_clock::now();

	auto pointNormals = buildPointNormalsCloud(*finiteCloud, *normals);
	ensureGp3NormalsAreFinite(*pointNormals);

	pcl::search::KdTree<pcl::PointNormal>::Ptr searchTree(
		new pcl::search::KdTree<pcl::PointNormal>());
	searchTree->setInputCloud(pointNormals);

	pcl::GreedyProjectionTriangulation<pcl::PointNormal> gp3;
	configureGp3(gp3, options, pointNormals, searchTree);

	pcl::PolygonMesh mesh;
	auto meshStart = std::chrono::steady_clock::now();
	gp3.reconstruct(mesh);
	auto meshEnd = std::chrono::steady_clock::now();

	saveMeshOrThrow(stats.outputPath, mesh);

	stats.pointCount = cloud->size();
	stats.finitePointCount = finiteCloud->size();
	stats.polygonCount = mesh.polygons.size();
	stats.normalDurationMs = std::chrono::duration<double, std::milli>(
		normalEnd - normalStart).count();
	stats.meshDurationMs = std::chrono::duration<double, std::milli>(
		meshEnd - meshStart).count();
	stats.totalDurationMs = stats.normalDurationMs + stats.meshDurationMs;
	stats.success = true;
	return stats;
}

ConversionStats convertSingleFile(
	const ToolOptions& options,
	const std::filesystem::path& inputPath)
{
	try
	{
		if (options.algorithm == MeshAlgorithm::Ofm)
		{
			return convertWithOfm(options, inputPath);
		}

		return convertWithGp3(options, inputPath);
	}
	catch (const std::exception& e)
	{
		ConversionStats stats;
		stats.algorithm = options.algorithm;
		stats.inputPath = inputPath;
		stats.outputPath = computeOutputPath(options, inputPath);
		stats.errorMessage = e.what();
		return stats;
	}
}

void printOfmStats(const ConversionStats& stats)
{
	std::cout << std::fixed << std::setprecision(3)
		<< "OK"
		<< " algorithm=ofm"
		<< " input=" << stats.inputPath
		<< " output=" << stats.outputPath
		<< " points=" << stats.pointCount
		<< " polygons=" << stats.polygonCount
		<< " mesh_ms=" << stats.meshDurationMs
		<< "\n";
}

void printGp3Stats(const ConversionStats& stats)
{
	std::cout << std::fixed << std::setprecision(3)
		<< "OK"
		<< " algorithm=gp3"
		<< " input=" << stats.inputPath
		<< " output=" << stats.outputPath
		<< " points=" << stats.pointCount
		<< " finite_points=" << stats.finitePointCount
		<< " polygons=" << stats.polygonCount
		<< " normal_ms=" << stats.normalDurationMs
		<< " mesh_ms=" << stats.meshDurationMs
		<< " total_ms=" << stats.totalDurationMs
		<< "\n";
}

void printPerFileStats(const ConversionStats& stats)
{
	if (!stats.success)
	{
		std::cout << "FAIL"
			<< " algorithm=" << algorithmToString(stats.algorithm)
			<< " input=" << stats.inputPath
			<< " error=\"" << stats.errorMessage << "\"\n";
		return;
	}

	if (stats.algorithm == MeshAlgorithm::Ofm)
	{
		printOfmStats(stats);
		return;
	}

	printGp3Stats(stats);
}

void printAggregateStats(const ToolOptions& options,
	const std::vector<ConversionStats>& statsList)
{
	std::size_t successCount = 0;
	std::size_t failureCount = 0;
	double totalNormalMs = 0.0;
	double totalMeshMs = 0.0;
	double totalConversionMs = 0.0;
	double minTotalMs = std::numeric_limits<double>::max();
	double maxTotalMs = 0.0;

	for (const auto& stats : statsList)
	{
		if (!stats.success)
		{
			++failureCount;
			continue;
		}

		++successCount;
		totalNormalMs += stats.normalDurationMs;
		totalMeshMs += stats.meshDurationMs;
		totalConversionMs += stats.totalDurationMs;
		minTotalMs = std::min(minTotalMs, stats.totalDurationMs);
		maxTotalMs = std::max(maxTotalMs, stats.totalDurationMs);
	}

	double averageTotalMs = successCount == 0
		? 0.0
		: totalConversionMs / static_cast<double>(successCount);
	if (successCount == 0)
	{
		minTotalMs = 0.0;
	}

	std::cout << std::fixed << std::setprecision(3)
		<< "SUMMARY"
		<< " algorithm=" << algorithmToString(options.algorithm)
		<< " processed=" << statsList.size()
		<< " succeeded=" << successCount
		<< " failed=" << failureCount;

	if (options.algorithm == MeshAlgorithm::Gp3)
	{
		std::cout
			<< " total_normal_ms=" << totalNormalMs
			<< " total_mesh_ms=" << totalMeshMs
			<< " total_conversion_ms=" << totalConversionMs
			<< " avg_total_ms=" << averageTotalMs
			<< " min_total_ms=" << minTotalMs
			<< " max_total_ms=" << maxTotalMs
			<< "\n";
		return;
	}

	std::cout
		<< " total_mesh_ms=" << totalMeshMs
		<< " avg_mesh_ms=" << averageTotalMs
		<< " min_mesh_ms=" << minTotalMs
		<< " max_mesh_ms=" << maxTotalMs
		<< "\n";
}

} // namespace

int main(int argc, char** argv)
{
	try
	{
		ToolOptions options = parseArgs(argc, argv);
		std::vector<ConversionStats> statsList;
		statsList.reserve(options.inputPaths.size());

		for (const auto& inputPath : options.inputPaths)
		{
			ConversionStats stats = convertSingleFile(options, inputPath);
			printPerFileStats(stats);
			statsList.push_back(std::move(stats));
		}

		printAggregateStats(options, statsList);
		return 0;
	}
	catch (const std::exception& e)
	{
		printUsage(argv[0]);
		std::cerr << e.what() << std::endl;
		return 1;
	}
}