点云最小生成树可视化 - Kruskal 算法实现 - 常规

#include #include <pcl/point_types.h> #include <pcl/io/ply_io.h> #include <pcl/visualization/pcl_visualizer.h> #include <pcl/common/centroid.h> #include <pcl/features/normal_3d.h> #include <pcl/visualization/cloud_viewer.h> #include <pcl/segmentation/extract_clusters.h> #include #include <unordered_map> #include "vtkAutoInit.h" #include <pcl/visualization/point_cloud_color_handlers.h> VTK_MODULE_INIT(vtkRenderingOpenGL); // VTK was built with vtkRenderingOpenGL2

using namespace pcl; typedef pcl::PointXYZ PointT; typedef pcl::PointCloud PointCloudT;

// 表示一条边的结构体 struct Edge { int src, tgt; float weight; };

// 表示并查集的子集的结构体 struct Subset { int parent, rank; };

// 表示一个连通的、无向的、带权重的图的类 class Graph { public: int V, E; std::vector edges;

Graph(int v, int e)
{
	V = v;
	E = e;
}

// 添加一条边到图中
void addEdge(int src, int tgt, float weight)
{
	Edge edge;
	edge.src = src;
	edge.tgt = tgt;
	edge.weight = weight;
	edges.push_back(edge);
}

// 查找某个元素所在的集合
int find(Subset subsets[], int i)
{
	if (subsets[i].parent != i)
		subsets[i].parent = find(subsets, subsets[i].parent);
	return subsets[i].parent;
}

// 合并两个集合
void Union(Subset subsets[], int x, int y)
{
	int xroot = find(subsets, x);
	int yroot = find(subsets, y);
	if (subsets[xroot].rank < subsets[yroot].rank)
		subsets[xroot].parent = yroot;
	else if (subsets[xroot].rank > subsets[yroot].rank)
		subsets[yroot].parent = xroot;
	else {
		subsets[yroot].parent = xroot;
		subsets[xroot].rank++;
	}
}

// Kruskal算法找到最小生成树
void KruskalMST(PointCloudT::Ptr cloud, std::vector<Edge>& result, float threshold)
{
	// 将边按照权重从小到大排序
	std::sort(edges.begin(), edges.end(), [](const Edge& a, const Edge& b)
	{
		return a.weight < b.weight;
	});

	// 为V个元素创建并查集的子集
	Subset* subsets = new Subset[V];
	for (int v = 0; v < V; ++v)
	{
		subsets[v].parent = v;
		subsets[v].rank = 0;
	}

	int i = 0;  // 用于选择下一条边的索引
	int e = 0;  // 用于选择下一条边加入最小生成树的索引

	// 需要选择V-1条边
	while (e < V - 1 && i < E)
	{
		Edge next_edge = edges[i++];

		int x = find(subsets, next_edge.src);
		int y = find(subsets, next_edge.tgt);

		// 如果加入这条边不会形成环，并且权重大于阈值，则加入到结果中，并且增加已选择边的计数
		if (next_edge.weight >= threshold && x != y)
		{
			result.push_back(next_edge);
			Union(subsets, x, y);
			++e;
		}
	}

	// 可视化最小生成树的结果
	pcl::visualization::PCLVisualizer viewer("Minimum Spanning Tree");
	viewer.setBackgroundColor(0, 0, 0);

	// 将原始点云添加到可视化窗口中
	pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> single_color(cloud, 255, 255, 255);
	viewer.addPointCloud<pcl::PointXYZ>(cloud, single_color, "original_cloud");

	// 将最小生成树的边添加到可视化窗口中
	for (const auto& edge : result)
	{
		const auto& src_point = cloud->points[edge.src];
		const auto& tgt_point = cloud->points[edge.tgt];
		std::stringstream ss;
		ss << "edge_" << edge.src << "_" << edge.tgt;
		viewer.addLine<pcl::PointXYZ>(src_point, tgt_point, ss.str());
	}
	// Count the number of points in the minimum spanning tree
	int numPoints = 0;
	for (const auto& edge : result)
	{
		const auto& src_point = cloud->points[edge.src];
		const auto& tgt_point = cloud->points[edge.tgt];
		numPoints += 2;
	}

	std::cout << "Number of points in the minimum spanning tree: " << numPoints << std::endl;
	while (!viewer.wasStopped())
	{
		viewer.spinOnce();
	}
}

};

// 计算两个点之间的欧式距离 double euclideanDistance(PointXYZ p1, PointXYZ p2) { double dx = p2.x - p1.x; double dy = p2.y - p1.y; double dz = p2.z - p1.z; return std::sqrt(dx * dx + dy * dy + dz * dz); }

int main() { // 从PLY文件加载输入点云 pcl::PointCloudpcl::PointXYZ::Ptr cloud(new pcl::PointCloudpcl::PointXYZ); pcl::io::loadPLYFilepcl::PointXYZ("D:\DIANYUNWENJIANJIA\newOUSHIJULEI_ply.ply", *cloud);

// 计算点云的质心
Eigen::Vector4f centroid;
pcl::compute3DCentroid(*cloud, centroid);

// 计算点云的法线
pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> ne;
pcl::PointCloud<pcl::Normal>::Ptr cloud_normals(new pcl::PointCloud<pcl::Normal>);
pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);
ne.setInputCloud(cloud);
ne.setSearchMethod(tree);
ne.setKSearch(40);
ne.compute(*cloud_normals);

// 创建一个有V个顶点和E个边的图
int V = cloud->size();
int E = V * (V - 1) / 2;
Graph graph(V, E);

// 基于点之间的欧式距离计算边的权重
for (int i = 0; i < V - 1; ++i)
{
	const auto& src_point = cloud->points[i];
	for (int j = i + 1; j < V; ++j)
	{
		const auto& tgt_point = cloud->points[j];
		float distance = euclideanDistance(src_point, tgt_point);
		graph.addEdge(i, j, distance);
	}
}
// 设置修剪的阈值
float threshold = 0.00080;
// 执行Kruskal算法找到最小生成树
std::vector<Edge> result;
graph.KruskalMST(cloud, result, threshold);

// 创建一个新的点云对象保存修剪后的最小生成树结果 pcl::PointCloudpcl::PointXYZ::Ptr new_cloud(new pcl::PointCloudpcl::PointXYZ); new_cloud->width = cloud->width; new_cloud->height = cloud->height; new_cloud->points.resize(cloud->points.size());

// 找到在最小生成树中出现三次的节点
std::unordered_map<int, int> node_count;
for (const auto& edge : result)
{
	node_count[edge.src]++;
	node_count[edge.tgt]++;
}

std::vector<int> threejiedian;
for (const auto& pair : node_count)
{
	if (pair.second == 3)
	{
		threejiedian.push_back(pair.first);
	}
}
// 如果与threejiedian相连的边的权重小于0.0008，则从最小生成树中移除该边
std::vector<Edge> pruned_result;
for (const auto& edge : result)
{
	if (std::find(threejiedian.begin(), threejiedian.end(), edge.src) != threejiedian.end() ||
		std::find(threejiedian.begin(), threejiedian.end(), edge.tgt) != threejiedian.end())
	{
		if (edge.weight <= 0.008)
		{
			pruned_result.push_back(edge);
		}
	}
	else
	{
		pruned_result.push_back(edge);
	}
}
// 将修剪后的最小生成树的顶点添加到新的点云对象中
for (const auto& edge : pruned_result)
{
	const auto& src_point = cloud->points[edge.src];
	const auto& tgt_point = cloud->points[edge.tgt];
	new_cloud->points[edge.src] = src_point;
	new_cloud->points[edge.tgt] = tgt_point;
}
// 将与threejiedian相连的边的顶点设置为绿色
pcl::PointCloud<pcl::PointXYZRGB>::Ptr colored_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
colored_cloud->points.resize(cloud->points.size());
for (size_t i = 0; i < cloud->points.size(); ++i)
{
	colored_cloud->points[i].x = cloud->points[i].x;
	colored_cloud->points[i].y = cloud->points[i].y;
	colored_cloud->points[i].z = cloud->points[i].z;

	colored_cloud->points[i].r = 255;
	colored_cloud->points[i].g = 0;
	colored_cloud->points[i].b = 0;
}
for (const auto& edge : result)
{
	if (std::find(threejiedian.begin(), threejiedian.end(), edge.src) != threejiedian.end() ||
		std::find(threejiedian.begin(), threejiedian.end(), edge.tgt) != threejiedian.end())
	{
		colored_cloud->points[edge.src].r = 0;
		colored_cloud->points[edge.src].g = 255;
		colored_cloud->points[edge.src].b = 0;
		colored_cloud->points[edge.tgt].r = 0;
		colored_cloud->points[edge.tgt].g = 255;
		colored_cloud->points[edge.tgt].b = 0;
	}
}
// 找到y值最大的threejiedian节点
float max_y = std::numeric_limits<float>::min();
int max_y_index = -1;
for (const auto& index : threejiedian)
{
	const auto& point = cloud->points[index];
	if (point.y > max_y)
	{
		max_y = point.y;
		max_y_index = index;
	}
}
// 将y值最大的threejiedian节点用绿色小球显示
pcl::visualization::PointCloudColorHandlerRGBField<pcl::PointXYZRGB> colored_cloud_color_handler(colored_cloud);
pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> red_color(cloud, 255, 0, 0);
pcl::visualization::PCLVisualizer viewer("Minimum Spanning Tree");
viewer.setBackgroundColor(0, 0, 0);
viewer.addPointCloud<pcl::PointXYZRGB>(colored_cloud, colored_cloud_color_handler, "colored_cloud");
viewer.addPointCloud<pcl::PointXYZ>(cloud, red_color, "original_cloud");
viewer.setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 10, "colored_cloud");
viewer.spin();
// 在y值最大的threejiedian节点上添加一个绿色小球
pcl::ModelCoefficients sphere_coeff;
sphere_coeff.values.resize(4);
sphere_coeff.values[0] = cloud->points[max_y_index].x;
sphere_coeff.values[1] = cloud->points[max_y_index].y;
sphere_coeff.values[2] = cloud->points[max_y_index].z;
sphere_coeff.values[3] = 0.00001;  // 设置球的半径为0.01
viewer.addSphere(sphere_coeff, "max_y_sphere");
viewer.spin();

// 将新的点云保存为PLY文件
pcl::io::savePLYFile("D:\\DIANYUNWENJIANJIA\\newKRUSKAL_ply.ply", *new_cloud, true);

// 加载第一次生成的点云文件
pcl::PointCloud<pcl::PointXYZ>::Ptr second_cloud(new pcl::PointCloud<pcl::PointXYZ>);
pcl::io::loadPLYFile<pcl::PointXYZ>("D:\\DIANYUNWENJIANJIA\\newKRUSKAL_ply.ply", *second_cloud);

// 再次执行 Kruskal 算法，并设置新的阈值
float new_threshold = 0.00008; // 调整阈值，生成新的最小生成树
std::vector<Edge> second_result;
graph.KruskalMST(second_cloud, second_result, new_threshold);

// 创建新的点云对象，保存第二次最小生成树结果
pcl::PointCloud<pcl::PointXYZ>::Ptr new_second_cloud(new pcl::PointCloud<pcl::PointXYZ>);
new_second_cloud->width = second_cloud->width;
new_second_cloud->height = second_cloud->height;
new_second_cloud->points.resize(second_cloud->points.size());

// 将第二次最小生成树的顶点添加到新的点云对象中
for (const auto& edge : second_result)
{
	const auto& src_point = second_cloud->points[edge.src];
	const auto& tgt_point = second_cloud->points[edge.tgt];
	new_second_cloud->points[edge.src] = src_point;
	new_second_cloud->points[edge.tgt] = tgt_point;
}

// 将第二次最小生成树的顶点设置为蓝色
pcl::PointCloud<pcl::PointXYZRGB>::Ptr second_colored_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
second_colored_cloud->points.resize(second_cloud->points.size());
for (size_t i = 0; i < second_cloud->points.size(); ++i)
{
	second_colored_cloud->points[i].x = second_cloud->points[i].x;
	second_colored_cloud->points[i].y = second_cloud->points[i].y;
	second_colored_cloud->points[i].z = second_cloud->points[i].z;

	second_colored_cloud->points[i].r = 0;
	second_colored_cloud->points[i].g = 0;
	second_colored_cloud->points[i].b = 255;
}

// 可视化第二次最小生成树结果
pcl::visualization::PointCloudColorHandlerRGBField<pcl::PointXYZRGB> second_colored_cloud_color_handler(second_colored_cloud);
pcl::visualization::PCLVisualizer second_viewer("Second Minimum Spanning Tree");
second_viewer.setBackgroundColor(0, 0, 0);
second_viewer.addPointCloud<pcl::PointXYZRGB>(second_colored_cloud, second_colored_cloud_color_handler, "second_colored_cloud");
second_viewer.addPointCloud<pcl::PointXYZ>(second_cloud, red_color, "second_original_cloud");
second_viewer.setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 10, "second_colored_cloud");
second_viewer.spin();

// 保存第二次生成的点云文件
pcl::io::savePLYFile("D:\\DIANYUNWENJIANJIA\\newKRUSKAL2_ply.ply", *new_second_cloud, true);

return 0;

}