基于C++、GDAL、OpenCV的矢量数据骨架线提取算法

基于C++、GDAL、OpenCV的矢量数据骨架线提取算法

CGAL已经实现了该功能，但由于CGAL依赖于Boost库，编译后过大，因此本文所采用的这套方式实现骨架线提取功能。

效果：

思路：
1、将导入shp按照要素逐一拆分成新的shp
2、将所有拆分后的shp分别转栅格，利用OpenCV提取骨架线
3、将所有骨架线转为shp，并合并输出

详细代码如下：

调用

    basePolygonAlgorithm::SkeletonExtractor extract2; 
	extract2.polygon2Skelton("originFile.shp", "outputFile.shp");

.h

#include "opencv2/core/core.hpp"
#include "opencv2/opencv.hpp"
#include"opencv2/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp" 
#include "ogrsf_frmts.h"
#include "gdal_priv.h"
#include "ogr_geometry.h"
#include "ogr_attrind.h"
#include "ogr_srs_api.h" 

		//提取骨架线
		class SkeletonExtractor
		{
		public:
			SkeletonExtractor();
			void polygon2Skelton(string src_path, string dst_path);//src_path待提取数据  dst_path：提取后数据
			;
			
		private:
			cv::Mat thinning(cv::Mat img);//Mat骨架线提取
			cv::Mat readRasterData(GDALDataset* pDstDataset);//pDstDataset转Mat
			void polygon2subPolygon(const string& src_path, int* featureNum, double adfGeoTransform[6]);
			void subPolygon2skelton(const string& dst_path, const int& featureNum, double adfGeoTransform[6]);
			void mergeShp(vector<string>vecSrcFiles, string pszDstFile);
			void writeShp(string dst_path, cv::Mat skel, double adfGeoTransform[], OGRSpatialReference* oSRS);
		};

.cpp

        SkeletonExtractor::SkeletonExtractor()
		{
		}

		void SkeletonExtractor::polygon2Skelton(string src_path, string dst_path)
		{
			int featureNum = 0;
			double adfGeoTransform[6];
			polygon2subPolygon(src_path, &featureNum, adfGeoTransform);
			subPolygon2skelton(dst_path, featureNum, adfGeoTransform);
			return;
		}
		
		cv::Mat SkeletonExtractor::thinning(cv::Mat img)
		{
			cv::Mat skel(img.size(), CV_8UC1, cv::Scalar(0));
			cv::Mat temp(img.size(), CV_8UC1);

			cv::Mat element = cv::getStructuringElement(cv::MORPH_CROSS, cv::Size(3, 3));

			bool done;
			do
			{
				cv::morphologyEx(img, temp, cv::MORPH_OPEN, element);
				cv::bitwise_not(temp, temp);
				cv::bitwise_and(img, temp, temp);
				cv::bitwise_or(skel, temp, skel);
				cv::erode(img, img, element);

				double maxVal = 0;
				cv::minMaxLoc(img, nullptr, &maxVal);
				done = (maxVal == 0);
			} while (!done);
 
			img.release();
			temp.release();
			return skel;  
		}

		cv::Mat SkeletonExtractor::readRasterData(GDALDataset* pDstDataset)
		{
			// Read the first band of the raster dataset
			GDALRasterBand* pBand = pDstDataset->GetRasterBand(1);

			// Allocate memory for the raster data
			int nXSize = pBand->GetXSize();
			int nYSize = pBand->GetYSize();
			GByte* pData = new GByte[nXSize * nYSize];

			// Read the raster data
			pBand->RasterIO(GF_Read, 0, 0, nXSize, nYSize, pData, nXSize, nYSize, GDT_Byte, 0, 0);
			//转化为Mat
			cv::Mat _Mymat(nYSize, nXSize, CV_8UC1);
			for (int i = 0; i < nXSize; i++)
			{
				for (int j = 0; j < nYSize; j++)
				{
					_Mymat.at<uchar>(j, i) = (uchar)pData[j * nXSize + i];
				}
			}
			// 执行骨架化
			cv::Mat skel;
			skel = thinning(_Mymat);

			// 定义结构元素（如3x3的矩形）
			int morph_size = 1;
			cv::Mat element = cv::getStructuringElement(cv::MORPH_RECT,
				cv::Size(2 * morph_size + 1, 2 * morph_size + 1),
				cv::Point(morph_size, morph_size));

			// 膨胀操作
			cv::dilate(skel, skel, element);

			// 腐蚀操作
			cv::erode(skel, skel, element);

			delete[] pData;
			return skel;
		}

		void SkeletonExtractor::polygon2subPolygon(const string& src_path, int* featureNum, double adfGeoTransform[6])
		{
			// Register GDAL drivers
			GDALAllRegister();
			OGRRegisterAll();
			//	NULL, NULL, NULL)); 
			GDALDataset* pSrcDataset = (GDALDataset*)GDALOpenEx(src_path.c_str(), GDAL_OF_VECTOR, NULL, NULL, NULL);
			// The one layer.
			OGRLayer* pSrcLayer = pSrcDataset->GetLayer(0);
			OGRSpatialReference* oSRS = NULL;
			oSRS = new OGRSpatialReference;
			oSRS = pSrcLayer->GetSpatialRef();
			*featureNum = pSrcLayer->GetFeatureCount();
			pSrcDataset->GetGeoTransform(adfGeoTransform);
			for (int i = 0; i < *featureNum; i++) {
				// 获取要素
				OGRFeature* pSrcFeature = pSrcLayer->GetFeature(i);

				// 创建新的GDALDataset和图层
				char outputFilename[256];
				sprintf(outputFilename, "temp_%d.shp", i);
				GDALDriver* pDriver = GetGDALDriverManager()->GetDriverByName("ESRI Shapefile");
				GDALDataset* temp_pDstDataset = pDriver->Create(outputFilename, 0, 0, 0, GDT_Unknown, NULL);
				OGRLayer* pDstLayer = temp_pDstDataset->CreateLayer(pSrcLayer->GetName(), oSRS, pSrcLayer->GetGeomType(), NULL);

				// 复制要素到新的图层中
				OGRFeature* pDstFeature = pSrcFeature->Clone();
				pDstLayer->CreateFeature(pDstFeature);
				OGRFeature::DestroyFeature(pDstFeature);
				GDALFlushCache(temp_pDstDataset);
				// 释放资源
				GDALClose(temp_pDstDataset);
				OGRFeature::DestroyFeature(pSrcFeature);
			}

			GDALClose(pSrcDataset); 
		}

		void SkeletonExtractor::mergeShp(vector<string> vecSrcFiles, string pszDstFile)
		{
			GDALAllRegister();

			// 获取Shapefile驱动
			GDALDriver* poDriver = GetGDALDriverManager()->GetDriverByName("ESRI Shapefile");
			if (poDriver == nullptr) {
				cerr << "ESRI Shapefile driver not available." << endl;
				return;
			}

			// 创建目标shapefile
			GDALDataset* poDstDS = poDriver->Create(pszDstFile.c_str(), 0, 0, 0, GDT_Unknown, nullptr);
			if (poDstDS == nullptr) {
				cerr << "Creation of output file failed." << endl;
				return;
			}

			OGRLayer* poDstLayer = nullptr;

			// 遍历所有源shapefile
			for (const auto& srcFile : vecSrcFiles) {
				GDALDataset* poSrcDS = static_cast<GDALDataset*>(GDALOpenEx(srcFile.c_str(), GDAL_OF_VECTOR, nullptr, nullptr, nullptr));

				if (poSrcDS == nullptr) {
					cerr << "Open failed: " << srcFile << endl;
					continue;
				}

				OGRLayer* poSrcLayer = poSrcDS->GetLayer(0);
				if (poSrcLayer == nullptr) {
					cerr << "Failed to get layer from source file: " << srcFile << endl;
					GDALClose(poSrcDS);
					continue;
				}

				// 如果目标图层不存在，则根据源图层创建一个
				if (poDstLayer == nullptr) {
					poDstLayer = poDstDS->CreateLayer(poSrcLayer->GetName(), poSrcLayer->GetSpatialRef(), poSrcLayer->GetGeomType(), nullptr);
					if (poDstLayer == nullptr) {
						cerr << "Failed to create destination layer." << endl;
						GDALClose(poSrcDS);
						GDALClose(poDstDS);
						return;
					}

					// 复制属性表结构
					OGRFeatureDefn* poSrcFDefn = poSrcLayer->GetLayerDefn();
					for (int i = 0; i < poSrcFDefn->GetFieldCount(); ++i) {
						poDstLayer->CreateField(poSrcFDefn->GetFieldDefn(i));
					}
				}

				// 复制要素
				OGRFeature* poFeature;
				poSrcLayer->ResetReading();
				while ((poFeature = poSrcLayer->GetNextFeature()) != nullptr) {
					OGRFeature* poDstFeature = OGRFeature::CreateFeature(poDstLayer->GetLayerDefn());
					poDstFeature->SetFrom(poFeature);
					poDstFeature->SetGeometry(poFeature->GetGeometryRef());

					if (poDstLayer->CreateFeature(poDstFeature) != OGRERR_NONE) {
						cerr << "Failed to create feature in destination shapefile." << endl;
						OGRFeature::DestroyFeature(poFeature);
						OGRFeature::DestroyFeature(poDstFeature);
						GDALClose(poSrcDS);
						GDALClose(poDstDS);
						return;
					}

					OGRFeature::DestroyFeature(poFeature);
					OGRFeature::DestroyFeature(poDstFeature);
				}

				GDALClose(poSrcDS);
			}

			GDALClose(poDstDS);
		}

		void SkeletonExtractor::writeShp(string dst_path, cv::Mat skel, double adfGeoTransform[], OGRSpatialReference* oSRS)
		{
			GDALAllRegister();
			OGRRegisterAll();
			GDALDriver* poDriver = GetGDALDriverManager()->GetDriverByName("ESRI Shapefile");
			GDALDataset* poDS = poDriver->Create(dst_path.c_str(), 0, 0, 0, GDT_Unknown, NULL);
			OGRLayer* poLayer = poDS->CreateLayer("layer", oSRS, wkbUnknown, NULL);
			std::vector<cv::Vec4i> hierarchy;
			std::vector<std::vector<cv::Point>> contours;
			cv::findContours(skel, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);
			OGRMultiLineString multiLineString;
			OGRFeature* poFeature = OGRFeature::CreateFeature(poLayer->GetLayerDefn());
			for (int i = 0; i < contours.size(); i++)
			{
				OGRLineString lineString;
				for (int j = 0; j < contours[i].size(); j++)
				{
					double x = adfGeoTransform[0] + contours[i][j].x * adfGeoTransform[1] + contours[i][j].y * adfGeoTransform[2];
					double y = adfGeoTransform[3] + contours[i][j].x * adfGeoTransform[4] + contours[i][j].y * adfGeoTransform[5];
					lineString.addPoint(x, y);
				}
				multiLineString.addGeometry(&lineString);
			}
			poFeature->SetGeometry(&multiLineString);
			skelton_result = poFeature->GetGeometryRef()->clone();//clone将skelton_result复制到一个新的 OGRGeometry 对象中，防止因poFeature释放而变成空
			poLayer->CreateFeature(poFeature);
			OGRFeature::DestroyFeature(poFeature);
			GDALClose(poDS);
			hierarchy.clear();
			contours.clear(); 
		}

		void SkeletonExtractor::subPolygon2skelton(const string& dst_path, const int& featureNum, double adfGeoTransform[6])
		{
			GDALAllRegister();
			OGRRegisterAll();

			//栅格参数设置
			char** argv = NULL;
			//nodata
			argv = CSLAddString(argv, "-a_nodata");
			argv = CSLAddString(argv, "-9999");
			
            //初始值
			argv = CSLAddString(argv, "-init");
			argv = CSLAddString(argv, "0");

            //分辨率
			argv = CSLAddString(argv, "-ts");
			argv = CSLAddString(argv, "1024");
			argv = CSLAddString(argv, "1024");

            //存储类型
			argv = CSLAddString(argv, "-ot");
			argv = CSLAddString(argv, "Byte");

			GDALRasterizeOptions* pOptions = GDALRasterizeOptionsNew(argv, NULL);
			OGRSpatialReference* oSRS = NULL;
			vector<string>filePaths;
			for (int i = 0; i < featureNum; i++) { 
				//重新打开新的SHP文件
				char outputFilename[256];
				sprintf(outputFilename, "temp_%d.shp", i);
				GDALDataset* pNewDS = (GDALDataset*)GDALOpenEx(outputFilename, GDAL_OF_VECTOR, NULL, NULL, NULL);
				OGRLayer* pNewLayer = pNewDS->GetLayer(0);
				oSRS = new OGRSpatialReference;
				oSRS = pNewLayer->GetSpatialRef();
				//获取转换6参数
				int pbUsageError = FALSE;
				GDALDataset* pImageDataset = static_cast<GDALDataset*>(GDALRasterize("output.tif", NULL, pNewDS, pOptions, &pbUsageError));

				double adfGeoTransform[6];
				pImageDataset->GetGeoTransform(adfGeoTransform);

				//提取骨架线
				cv::Mat skel = readRasterData(pImageDataset);

				//写入结果
				outputFilename[256];
				sprintf(outputFilename, "res_%d.shp", i);
				writeShp(outputFilename, skel, adfGeoTransform, oSRS);
				filePaths.push_back(outputFilename);
				GDALClose(pNewDS);
				GDALClose(pImageDataset);
				skel.release();
			}

			// Close the datasets
			GDALRasterizeOptionsFree(pOptions);

			//合并shp
			mergeShp(filePaths, dst_path);

			//删除临时文件
			for (int i = 0; i < featureNum; i++) {
				char tempFile[256];
				char resFile[256];
				sprintf(tempFile, "temp_%d.shp", i);
				std::remove(tempFile);
				sprintf(tempFile, "temp_%d.shx", i);
				std::remove(tempFile);
				sprintf(tempFile, "temp_%d.dbf", i);
				std::remove(tempFile);
				sprintf(tempFile, "temp_%d.prj", i);
				std::remove(tempFile);

				sprintf(resFile, "res_%d.shp", i);
				std::remove(resFile);
				sprintf(resFile, "res_%d.shx", i);
				std::remove(resFile);
				sprintf(resFile, "res_%d.dbf", i);
				std::remove(resFile);
				sprintf(resFile, "res_%d.prj", i);
				std::remove(resFile);
			}
			std::remove("output.tif");
			return;
		}