PCL 点云最小生成树算法优化：去除连接三次节点的低权重边

#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"

VTK_MODULE_INIT(vtkRenderingOpenGL); // VTK was built with vtkRenderingOpenGL2

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

// Structure to represent an edge struct Edge { int src, tgt; float weight; };

// Structure to represent a subset for union-find struct Subset { int parent, rank; };

// Class to represent a connected, undirected and weighted graph class Graph { public: int V, E; std::vector edges;

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

// Add an edge to the graph void addEdge(int src, int tgt, float weight) { Edge edge; edge.src = src; edge.tgt = tgt; edge.weight = weight; edges.push_back(edge); }

// Find set of an element i int find(Subset subsets[], int i) { if (subsets[i].parent != i) subsets[i].parent = find(subsets, subsets[i].parent); return subsets[i].parent; }

// Union of two sets x and 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's algorithm to find MST void KruskalMST(PointCloudT::Ptr cloud, std::vector& result, float threshold) { // Sort all the edges in non-decreasing order of their weight std::sort(edges.begin(), edges.end(), [](const Edge& a, const Edge& b) { return a.weight < b.weight; });

// Allocate memory for creating V subsets Subset* subsets = new Subset[V]; for (int v = 0; v < V; ++v) { subsets[v].parent = v; subsets[v].rank = 0; }

int i = 0; // Index used to pick the next smallest edge int e = 0; // Index used to pick the next edge to include in MST

// Number of edges to be taken is equal to 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 including this edge doesn't cause a cycle and weight is above the threshold, include it in result and increment the index of result for next edge if (next_edge.weight >= threshold && x != y) { result.push_back(next_edge); Union(subsets, x, y); ++e; } }

// Find nodes that appear three times in the minimum spanning tree std::unordered_map<int, int> node_counts; for (const auto& edge : result) { node_counts[edge.src]++; node_counts[edge.tgt]++; }

std::vector threejiedian; for (const auto& node_count : node_counts) { if (node_count.second == 3) { threejiedian.push_back(node_count.first); } }

// Remove edges connected to threejiedian nodes with weight less than 0.0008 std::vector 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.0008) { pruned_result.push_back(edge); } } else { pruned_result.push_back(edge); } }

// Visualize the resulting minimum spanning tree pcl::visualization::PCLVisualizer viewer("Minimum Spanning Tree"); viewer.setBackgroundColor(0, 0, 0);

// Add the original point cloud to the viewer pcl::visualization::PointCloudColorHandlerCustompcl::PointXYZ single_color(cloud, 255, 255, 255); viewer.addPointCloudpcl::PointXYZ(cloud, single_color, "original_cloud");

// Add the minimum spanning tree edges to the viewer for (const auto& edge : pruned_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.addLinepcl::PointXYZ(src_point, tgt_point, ss.str()); }

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() { // Load the input point cloud from PLY file pcl::PointCloudpcl::PointXYZ::Ptr cloud(new pcl::PointCloudpcl::PointXYZ); pcl::io::loadPLYFilepcl::PointXYZ("D:\DIANYUNWENJIANJIA\newOUSHIJULEI_ply.ply", *cloud);

// Compute the centroid of the point cloud Eigen::Vector4f centroid; pcl::compute3DCentroid(*cloud, centroid);

// Compute the normals of the point cloud pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> ne; pcl::PointCloudpcl::Normal::Ptr cloud_normals(new pcl::PointCloudpcl::Normal); pcl::search::KdTreepcl::PointXYZ::Ptr tree(new pcl::search::KdTreepcl::PointXYZ); ne.setInputCloud(cloud); ne.setSearchMethod(tree); ne.setKSearch(40); ne.compute(*cloud_normals);

// Create a graph with V vertices and E edges int V = cloud->size(); int E = V * (V - 1) / 2; Graph graph(V, E);

// Calculate the edge weights based on Euclidean distance between points 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); } }

// Set the threshold for pruning float threshold = 0.00080;

// Perform Kruskal's algorithm to find the minimum spanning tree std::vector result; graph.KruskalMST(cloud, result, threshold);

// Find nodes that appear three times in the minimum spanning tree std::unordered_map<int, int> node_counts; for (const auto& edge : result) { node_counts[edge.src]++; node_counts[edge.tgt]++; }

std::vector threejiedian; for (const auto& node_count : node_counts) { if (node_count.second == 3) { threejiedian.push_back(node_count.first); } }

// Save threejiedian nodes to a file std::ofstream outfile("threejiedian.txt"); if (outfile.is_open()) { for (const auto& node : threejiedian) { outfile << node << std::endl; } outfile.close(); }

// Create a new point cloud object to save the result of the minimum spanning tree 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());

// Add the vertices of the minimum spanning tree to the new point cloud object for (const auto& edge : 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; }

// Save the new point cloud as a PLY file pcl::io::savePLYFile("D:\DIANYUNWENJIANJIA\newKRUSKAL_ply.ply", *new_cloud, true); return 0;