点云最小生成树可视化：Kruskal算法实现与节点高亮

#include <iostream>
#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 <vector>
#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<PointT> PointCloudT;

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

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

// 表示一个连通的、无向的、带权重的图的类
class Graph
{
public:
	int V, E;
	std::vector<Edge> 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;
		// 将与threejiedian相连的边的顶点设置为绿色
		pcl::PointCloud<pcl::PointXYZ>::Ptr green_cloud(new pcl::PointCloud<pcl::PointXYZ>);
		*green_cloud = *cloud;
		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())
			{
				const auto& src_point = cloud->points[edge.src];
				const auto& tgt_point = cloud->points[edge.tgt];
				green_cloud->points[edge.src].r = 0;
				green_cloud->points[edge.src].g = 255;
				green_cloud->points[edge.src].b = 0;
				green_cloud->points[edge.tgt].r = 0;
				green_cloud->points[edge.tgt].g = 255;
				green_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::PointCloudColorHandlerCustom<pcl::PointXYZ> green_color(green_cloud, 0, 255, 0);
		pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> red_color(cloud, 255, 0, 0);
		viewer.addPointCloud<pcl::PointXYZ>(cloud, red_color, "original_cloud");
		viewer.addPointCloud<pcl::PointXYZ>(green_cloud, green_color, "green_cloud");
		viewer.setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 10, "green_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.01;  // 设置球的半径为0.01
		viewer.addSphere(sphere_coeff, "max_y_sphere");
		viewer.spin();
		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::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
	pcl::io::loadPLYFile<pcl::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::PointCloud<pcl::PointXYZ>::Ptr new_cloud(new pcl::PointCloud<pcl::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;
	}

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