Vector data skeleton line extraction algorithm based on C++, GDAL, OpenCV

Vector data skeleton line extraction algorithm based on C++, GDAL, OpenCV

CGAL has already implemented this function, but because CGAL relies on the Boost library, it is too large after compilation, so the method used in this article implements the skeleton line extraction function.

Effect:

Ideas:
1. Split the imported shp into new shp one by one according to elements
2. Convert all split shp to rasters respectively, and use OpenCV to extract skeleton lines
3. Convert all skeleton lines to shp and merge the output

The detailed code is as follows:

transfer

    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;
		}