基于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/sample_consensus/ransac.h>
#include <pcl/sample_consensus/sac_model_line.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/conditional_removal.h>
#include <opencv2/opencv.hpp>
#include <cmath>
#include <pcl/features/principal_curvatures.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;
};

double calculateEuclideanDistance(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);
}

// 类表示一个连通的、无向的、加权图
class Graph
{
public:
	int V, E;
	std::vector<Edge> edges;
	std::vector<Vertex> vertexes;

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

	// 结构体表示一个顶点
	struct Vertex
	{
		pcl::PointXYZ point;
	};

	int addVertex(pcl::PointXYZ point)
	{
		Vertex vertex;
		vertex.point = point;
		vertexes.push_back(vertex);
		return vertexes.size() - 1;
	}

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

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

	// 合并两个集合 x 和 y
	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 算法查找 MST
	void KruskalMST(PointCloudT::Ptr cloud, std::vector<Edge>& result)
	{
		// 将所有边按权重非递减顺序排序
		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;  // 用于选择下一个要包含在 MST 中的边的索引

		// 要取的边数等于 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 < 0.043)
			{
				if (x != y)
				{
					result.push_back(next_edge);
					Union(subsets, x, y);
					++e;
				}
			}
		}
		// 删除与权重<0.013 的单次出现的或三次出现的点相连的边
		std::vector<Edge> filtered_result;
		std::unordered_map<int, int> countMap2;
		for (const auto& edge : result) {
			countMap2[edge.src]++;
			countMap2[edge.tgt]++;
		}
		for (const auto& edge : result) {
			if (!(countMap2[edge.src] == 1 || countMap2[edge.src] == 3) ||
				!(countMap2[edge.tgt] == 1 || countMap2[edge.tgt] == 3) ||
				edge.weight <= 0.013) {
				filtered_result.push_back(edge);
			}
		}
		result = filtered_result;
		// 可视化结果最小生成树
		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());
		}
		// 连接两个断开的边
		pcl::PointXYZ point1 = cloud->points[16];
		pcl::PointXYZ point2 = cloud->points[33];
		pcl::PointXYZ point3 = cloud->points[32];
		pcl::PointXYZ point4 = cloud->points[38];
		viewer.addLine<pcl::PointXYZ, pcl::PointXYZ>(point1, point2, 0, 255, 0, "line1");
		viewer.addLine<pcl::PointXYZ, pcl::PointXYZ>(point3, point4, 0, 255, 0, "line2");
		pcl::PointXYZ point5 = cloud->points[61];
		viewer.addSphere(point5, 0.004, 0, 1, 0, "sphere", 0);

		// 查找 y 方向上具有最大值的节点并将其标记为绿色
		std::unordered_map<int, int> countMap;
		for (const auto& edge : result)
		{
			countMap[edge.src]++;
			countMap[edge.tgt]++;
		}
		// 出现三次的节点用绿色表示，表示为节点
		pcl::PointCloud<pcl::PointXYZ>::Ptr threejiedian(new pcl::PointCloud<pcl::PointXYZ>);
		int begin = 0;
		for (const auto& pair : countMap)
		{
			if (pair.second >= 3)
			{
				//std::cout << "该点坐标： " << pair.first << " 出现 " << pair.second << " 次" << std::endl;
				//std::cout << cloud->points[pair.first] << std::endl;
				pcl::PointXYZ point6 = cloud->points[pair.first];
				std::stringstream sa;
				//std::cout << "节点索引号: " << pair.first << std::endl;
				sa << begin;
				begin++;
				viewer.addSphere(pcl::PointXYZ(point6), 0.0023, 0, 1, 0, sa.str());
				threejiedian->push_back(point3);
			}
		}
		// 出现一次的节点用红色表示，表示为叶尖
		pcl::PointCloud<pcl::PointXYZ>::Ptr singlejiedian(new pcl::PointCloud<pcl::PointXYZ>);
		int begin2 = 0;
		for (const auto& pair : countMap)
		{
			if (pair.second == 1)
			{
				pcl::PointXYZ point6 = cloud->points[pair.first];
				std::stringstream sb;
				sb << begin2;
				begin2++;
				viewer.addSphere(pcl::PointXYZ(point6), 0.0023, 1, 0, 0, sb.str());
				singlejiedian->push_back(point6);
			}
		}
		// 遍历所有边
		for (int i = 0; i < edges.size(); i++)
		{
			Edge next_edge = edges[i];

			// 获取边的起点和终点索引
			int src = next_edge.src;
			int tgt = next_edge.tgt;

			// 检查起点和终点是否与 `singlejiedian` 相连
			bool isSingleJiedianSrc = false;
			bool isSingleJiedianTgt = false;
			for (const auto& point : *singlejiedian)
			{
				float distanceSrc = calculateEuclideanDistance(cloud->points[src], point);
				float distanceTgt = calculateEuclideanDistance(cloud->points[tgt], point);
				if (distanceSrc < 0.013)
				{
					isSingleJiedianSrc = true;
				}
				if (distanceTgt < 0.013)
				{
					isSingleJiedianTgt = true;
				}
			}
			// 如果边的起点和终点有一个点与 `singlejiedian` 相连且权重小于0.014，则删除该边
			if ((isSingleJiedianSrc || isSingleJiedianTgt) && next_edge.weight < 0.013)
			{
				// 不将边添加到最小生成树结果中，即删除该边
				continue;
			}

			// 将边添加到最小生成树结果中
			result.push_back(next_edge);
		}
		while (!viewer.wasStopped())
		{
			viewer.spinOnce();
		}
	}
	void processPointCloud(PointCloudT::Ptr cloud) {
		// 计算点云的质心和法线
		pcl::PointXYZ centroid;
		double search_radius = 0.1;
		pcl::computeCentroid(*cloud, centroid);
		pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> ne;
		ne.setInputCloud(cloud);
		pcl::search::KdTree<pcl::PointXYZ>::Ptr search_method(new pcl::search::KdTree<pcl::PointXYZ>);
		ne.setSearchMethod(search_method);
		ne.setRadiusSearch(search_radius);
		pcl::PointCloud<pcl::Normal>::Ptr cloud_normals(new pcl::PointCloud<pcl::Normal>);
		ne.compute(*cloud_normals);
		pcl::PrincipalCurvaturesEstimation<pcl::PointXYZ, pcl::Normal, pcl::PrincipalCurvatures> ce;
		ce.setInputCloud(cloud);
		ce.setInputNormals(cloud_normals);
		ce.setSearchMethod(search_method);
		ce.setRadiusSearch(search_radius);
		pcl::PointCloud<pcl::PrincipalCurvatures>::Ptr cloud_curvatures(new pcl::PointCloud<pcl::PrincipalCurvatures>);
		ce.compute(*cloud_curvatures);

		// 根据曲率阈值分类点云
		std::vector<pcl::PointXYZ> leaf_points;
		std::vector<pcl::PointXYZ> stem_points;
		float curvature_threshold = 0.1f; // 自定义的曲率阈值
		for (int i = 0; i < cloud->size(); ++i) {
			if (cloud_curvatures->at(i).principal_curvature[0] < curvature_threshold) {
				leaf_points.push_back(cloud->at(i));
			}
			else {
				stem_points.push_back(cloud->at(i));
			}
		}

		// 将分类后的点添加到骨架化连接中
		std::vector<pcl::PointXYZ> skeleton_points;
		skeleton_points.insert(skeleton_points.end(), leaf_points.begin(), leaf_points.end());
		skeleton_points.insert(skeleton_points.end(), stem_points.begin(), stem_points.end());
		for (int i = 0; i < skeleton_points.size(); ++i) {
			addVertex(skeleton_points[i]);
		}
		for (int i = 0; i < edges.size(); ++i) {
			pcl::PointXYZ p1 = edges[i].p1;
			pcl::PointXYZ p2 = edges[i].p2;
			float weight = calculateEuclideanDistance(p1, p2);
			addEdge(findVertexIndex(p1), findVertexIndex(p2), weight);
		}
		std::vector<Edge> result;
		KruskalMST(cloud, result);
		for (int i = 0; i < result.size(); ++i) {
			int index1 = result[i].src;
			int index2 = result[i].tgt;
			float weight = result[i].weight;
			pcl::PointXYZ p1 = vertexes[index1].point;
			pcl::PointXYZ p2 = vertexes[index2].point;
			addEdge(index1, index2, weight);
		}
	}
};