基于GDAL进行图像的读取, 保存, 滤波

基于GDAL进行图像的读取, 保存, 滤波

目前支持: 自定义算子的图像滤波, 图像读取,保存, 通道分离等功能

#pragma once

#ifdef __GDAL_UTILITY__
#include <gdal/gdal_priv.h>
#include <gdal/gdal.h>
#include <iostream>
#include <string>
#include <sstream>
#include <fstream>
#include <algorithm>
#include <list>
#include <vector>
#include <array>
#include <cassert>
#include <memory>
#include <cstring>
#include <initializer_list>
#include <random>
#include <thread>


namespace GD {

	using namespace std;
	void init() {
		GDALAllRegister();
		CPLSetConfigOption("GDAL_FILENAME_IS_UTF8", "NO");
	}

	template<typename T>
	T sum(vector<T> const& v) {
		T sum = 0;
		std::for_each(v.begin(), v.end(), [&](T const& e) {sum += e; });
		return sum;
	}
	template<typename T>
	int partition(vector<T>& vec, int left, int right) {
		int i = left, j = right + 1;
		T pivot = vec[left];
		while (1) {
			while (vec[++i] < pivot) {
				if (i == right) break;
			}
			while (vec[--j] > pivot) {
				if (j == left) break;
			}
			if (i >= j) break;
			swap(vec[i], vec[j]);
		}
		swap(vec[left], vec[j]);
		return j;
	}

	template<typename T>
	T qSelect(vector<T>& vec, size_t k) {
		int low = 0, high = vec.size() - 1;
		while (low <= high) {
			int j = partition(vec, low, high);
			if (j == k) {
				return vec[k];
			}
			else if (j > k) {
				high = j - 1;
			}
			else if (j < k) {
				low = j + 1;
			}
		}
		return vec[k];
	}


	template<typename _Ty>
	struct Matrix {
	public:
		static_assert(std::is_arithmetic<_Ty>::value, "_Ty must be numeric type");
		// data
		vector<_Ty*> Data;
		size_t m_width, m_height;
	public:
		Matrix(size_t width = 1, size_t height = 1) :m_width(width), m_height(height) {
			Data = vector<_Ty*>(height, nullptr);
			for (int i = 0; i < height; ++i) {
				Data[i] = new _Ty[width];
				memset((void*)Data[i], 0, width * sizeof(_Ty));
			}
			int centerx = width / 2, centery = height / 2;
			Data[centerx][centery] = 1;
		}
		Matrix(std::initializer_list<std::initializer_list<_Ty>> const& init_list)
		{
			size_t listW = init_list.begin()->size(), listH = init_list.size();
			m_width = listW;
			m_height = listH;
			if (listW != m_width || listH != m_height) {
				cerr << "Matrix init error" << endl;
				exit(1);
			}
			int i = 0;
			Data = vector<_Ty*>(m_height);
			for (auto const& list : init_list) {
				if (list.size() != m_width) {
					cerr << "Matrix init error" << endl;
					exit(1);
				}
				int j = 0;
				Data[i] = new _Ty[m_width];
				memset((void*)Data[i], 0, m_width * sizeof(_Ty));
				for (auto e : list) {
					Data[i][j++] = e;
				}
				++i;
			}
		}
	};

	template<typename T>
	Matrix<T> AverageKernel(size_t size) {
		if ((size & 1) == 0) {
			std::cout << "size must be odd" << endl;
		}

		Matrix<T> res(size, size);
		for (int i = 0; i < size; ++i) {
			for (int j = 0; j < size; j++) {
				res.Data[i][j] = 1.0 / (size * size);
			}
		}
		return res;
	}


	template<typename T>
	Matrix<T> SobelKernelx() {
		return Matrix<T>(
			{ {-1., -2., -1.},
			  { 0.,  0.,  0.},
			  { 1.,  2.,  1.} });
	}

	template<typename T>
	Matrix<T> SobelKernely() {
		return Matrix<T>(
			{ {-1., 0.,  1.},
			  {-2., 0.,  2.},
			  {-1., 0.,  1.} });
	}
	/**
	* { {-1., -2., -1.},
		{ 0.,  0.,  0.},
		{ 1.,  2.,  1.} });

	  { {-1., 0.,  1.},
		{-2., 0.,  2.},
		{-1., 0.,  1.} });
	*/
	template<typename T>
	Matrix<T> LaplaceKernel_8() {
		return Matrix<T>(
			{ {1,  1, 1},
			  {1, -8, 1},
			  {1,  1, 1} }
		);
	}

	template<typename T>
	Matrix<T> LaplaceKernel_4() {
		return Matrix<T>(
			{ {0,  1, 0},
			  {1, -4, 1},
			  {0,  1, 0} }
		);
	}

	template<typename T>
	Matrix<T> LaplaceSharpening_4() {
		return Matrix<T>(
			{ {0,  1, 0},
			  {1, -5, 1},
			  {0,  1, 0} }
		);
	}

	template<typename T>
	Matrix<T> LaplaceSharpening_8() {
		return Matrix<T>(
			{ {1,  1, 1},
			  {1, -9, 1},
			  {1,  1, 1} }
		);
	}
	template<typename T>
	Matrix<T> GaussianKernel(size_t size, T stdX, T stdY) {
		// 得到size*size的高斯核
		Matrix<T> ret(size, size);
		T sum = 0;
		for (int i = 0; i < size; ++i) {
			T x = i - size / 2;
			for (int j = 0; j < size; ++j) {
				T y = j - size / 2;
				ret.Data[i][j] = exp(-(x * x / (2 * stdX * stdX) + y * y / (2 * stdY * stdY)));
				sum += ret.Data[i][j];
			}
		}
		for (int i = 0; i < size; ++i) {
			for (int j = 0; j < size; j++) {
				ret.Data[i][j] /= sum;
			}
		}
		return ret;
	}

	class RasterBase {
		void* pData;
	public:
		RasterBase() = default;
		RasterBase(const string& filepath) {}
		RasterBase(const char* filepath) {}

		virtual size_t width() const { return 0; }
		virtual size_t height() const { return 0; }
		virtual size_t bands() const { return 0; }
		virtual void* data() const { return nullptr; }
		virtual GDALDataType type() const { return GDT_Unknown; }
		virtual void init(GDALDataset* pDatasetRead,
			size_t lWidth, size_t lHeight, size_t nBands,
			GDALDataType datatype) {};
		virtual void save(string const&, string const& DriverName = "GTIFF")const {};
		virtual vector<array<size_t, 256>> histgram() const { return vector<array<size_t, 256>>(); }
		virtual RasterBase* equalizeHist(vector<array<size_t, 256>> histgram) { return nullptr; }
		virtual RasterBase* equalizeHist() { return nullptr; }
		virtual RasterBase* channel_extraction(size_t channel) { return nullptr; };
		virtual RasterBase* filter(Matrix<float> const& kernel) { return nullptr; }
		virtual RasterBase* filter(Matrix<double> const& kernel) { return nullptr; }
		virtual RasterBase* MedianFilter() { return nullptr; }
		virtual RasterBase* addSalt(size_t num) { return nullptr; }
		virtual RasterBase* Sobel() { return nullptr; }
		virtual RasterBase* Laplace_4() { return nullptr; }
		virtual RasterBase* Laplace_8() { return nullptr; }
		virtual RasterBase* GaussBlur(size_t size, double stdx, double stdy) { return nullptr; }
		virtual RasterBase* addGaussNoise(size_t size, float stdy) { return nullptr; }
		virtual RasterBase* Laplace_4_Sharpening() { return nullptr; }
		virtual RasterBase* Laplace_8_Sharpening() { return nullptr; }
		virtual RasterBase* Sobel_Sharpening() { return nullptr; }
		virtual RasterBase* toGray() { return nullptr; }
		virtual ~RasterBase() {};
	};

	template<typename _T>
	_T* Malloc(size_t width, size_t height, size_t bands) {
		return (_T*)CPLMalloc(width * height * bands * sizeof(_T));
	}

	using RasterPtr = std::shared_ptr<RasterBase>;

	template<typename T>
	class Raster : virtual public RasterBase {
	private:
		GDALDataset* pDatasetRead;
		GDALDriver* pDriver;
		size_t lWidth, lHeight, nBands;
		T* pData;
		GDALDataType datatype;
		template<typename _T>
		_T& at(size_t x, size_t y) {
			return pData[x * width * nBands + y];
		}

		template<typename _Ty>
		vector<_Ty> conv(long long pos, Matrix<_Ty> const& mat) {
			// 对单个像素进行卷积, i为图像中的位置, 原图像边界采用0填充
			vector<_Ty> res;
			size_t total = lHeight * lWidth;
			long long newpos = pos % total;
			size_t x = newpos % lWidth, y = newpos / lWidth;
			long long n = pos / total;
			size_t matw = mat.m_width, math = mat.m_height;
			long long matx = matw / 2, maty = math / 2;
			for (int j = -maty; j <= maty; ++j) {
				int yj = y + j;
				for (int i = -matx; i <= matx; ++i) {
					int xi = x + i;
					if (xi < 0 || xi >= lHeight || yj < 0 || yj >= lWidth) {
						res.push_back(0);
						continue;
					}
					long long index = yj * lWidth + xi + n * total;
					res.push_back(((_Ty)pData[index]) * ((_Ty)mat.Data[maty + j][matx + i]));
				}
			}
			return res;
		}

		template<typename _Ty>
		RasterBase* Laplace(Matrix<_Ty> const& kernel) {
			RasterPtr pLaplace{ filter(kernel) };
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, 1, datatype) };
			size_t Area = lWidth * lHeight;
			for (size_t i = 0; i < Area; ++i) {
				double sum{ 0.0 };
				for (long long channel = 0; channel < nBands; ++channel) {
					long long index = i + channel * Area;
					sum += ((T*)pLaplace->data())[index];
				}
				sum /= 3;
				((T*)ret->data())[i] = (T)sum;
			}
			return ret;
		}

	public:
		size_t width() const override { return lWidth; }
		size_t height() const override { return lHeight; }
		size_t bands() const override { return nBands; }
		void* data() const override { return (void*)pData; }
		GDALDataType type() const override { return datatype; }
		void init(GDALDataset* pDatasetRead,
			size_t lWidth, size_t lHeight, size_t nBands,
			GDALDataType datatype)override
		{
			this->pDatasetRead = pDatasetRead;
			this->lWidth = lWidth; this->lHeight = lHeight; this->nBands = nBands;
			this->datatype = datatype;
			T* p(Malloc<T>(lWidth, lHeight, nBands));
			if (this->pDatasetRead)
				pDatasetRead->RasterIO(GF_Read, 0, 0, lWidth, lHeight, (void*)p, lWidth, lHeight, datatype, nBands, NULL, 0, 0, 0);//读取图像数据
			this->pData = p;
		}
		void save(string const& filepath, string const& DriverName = "GTIFF") const override {
			GDALDataset* pDatasetSave;
			auto pDriver = GetGDALDriverManager()->GetDriverByName(DriverName.c_str());
			pDatasetSave = pDriver->Create(filepath.c_str(), lWidth, lHeight, nBands, datatype,
				nullptr);
			assert(pDatasetSave != nullptr);
			pDatasetSave->RasterIO(GF_Write, 0, 0, lWidth, lHeight, (void*)pData, lWidth, lHeight, datatype, nBands, NULL, 0, 0, 0);//创建图像数据
			GDALClose(pDatasetSave);
		}

		vector<array<size_t, 256>> histgram() const override {
			vector<array<size_t, 256>> ret(nBands);
			for (int i = 0; i < nBands; ++i) {
				ret[i].fill(0);
			}
			long long total = this->lWidth * this->lHeight;
			for (int channel = 0; channel < nBands; ++channel) {
				auto p = pData + channel * total;
				for (int i = 0; i < total; ++i) {
					++ret[channel][p[i]];
				}
			}
			return ret;
		}

		RasterBase* equalizeHist(vector<array<size_t, 256>> histgram) override {
			RasterBase* ret{ new Raster<T> };
			ret->init(nullptr, lWidth, lHeight, nBands, datatype);
			size_t total = lHeight * lWidth;
			memcpy((T*)ret->data(), (T*)this->data(), nBands * total * sizeof(T));
			for (int channel = 0; channel < nBands; ++channel) {
				T* srcbegin = (T*)this->data() + lHeight * lWidth * channel;
				T* tarbegin = (T*)ret->data() + lHeight * lWidth * channel;
				for (int i = 1; i < 256; ++i) {
					histgram[channel][i] += histgram[channel][i - 1];
				}
				for (int i = 0; i < total; ++i) {
					tarbegin[i] = static_cast<T>(255.0 * histgram[channel][srcbegin[i]] / total);
				}
			}
			return ret;
		}

		RasterBase* equalizeHist()override {
			return equalizeHist(histgram());
		}

		RasterBase* channel_extraction(size_t channel) {
			assert(channel < nBands);
			RasterBase* ret{ new Raster<T> };
			ret->init(nullptr, lWidth, lHeight, nBands, datatype);
			size_t total = lHeight * lWidth;
			T* srcbegin = pData + lHeight * lWidth * channel;
			T* tarbegin = (T*)ret->data() + lHeight * lWidth * channel;
			memset(ret->data(), 0, total * sizeof(T));
			for (int i = 0; i < total; ++i) {
				tarbegin[i] = srcbegin[i];
			}
			return ret;
		}


		RasterBase* filter(Matrix<float> const& kernel) {
			if (kernel.m_width % 2 == 0 || kernel.m_height % 2 == 0) {
				cerr << "kernel size must be odd" << endl;
				exit(1);
			}
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, nBands, datatype) };
			size_t total = lHeight * lWidth * nBands;
			for (size_t i = 0; i < total; ++i) {
				auto res = conv(i, kernel);
				double s = sum(res);
				if (s < 0) s = fabs(s);
				if (s > 255) s = 255;
				((T*)ret->data())[i] = (T)(s);
			}
			return ret;
		}

		RasterBase* filter(Matrix<double> const& kernel) {
			if (kernel.m_width % 2 == 0 || kernel.m_height % 2 == 0) {
				cerr << "kernel size must be odd" << endl;
				exit(1);
			}
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, nBands, datatype) };
			size_t total = lHeight * lWidth * nBands;
			for (size_t i = 0; i < total; ++i) {
				auto res = conv(i, kernel);
				double s = sum(res);
				if (s < 0) s = fabs(s);
				if (s > 255) s = 255;
				((T*)ret->data())[i] = (T)(s);
			}
			return ret;
		}

		RasterBase* MedianFilter()override {
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, nBands, datatype) };
			size_t total = lHeight * lWidth * nBands;
			Matrix<int> kernel{ {1, 1, 1} ,{1, 1, 1}, {1, 1, 1} };
			for (size_t i = 0; i < total; ++i) {
				auto res = conv(i, kernel);
				((T*)ret->data())[i] = qSelect(res, res.size() / 2);
			}
			return ret;
		}

		RasterBase* addSalt(size_t num) {
			size_t limit = lWidth * lHeight;
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, nBands, datatype) };
			memcpy(ret->data(), this->data(), limit * nBands);
			for (long long i = 0; i < num; ++i) {
				int randomIndex = (rand() * rand()) % limit;
				bool is_white = rand() % 2;
				for (int channel = 0; channel < nBands; ++channel) {
					((T*)ret->data())[randomIndex + channel * limit] = is_white ? 255 : 0;
				}
			}
			return ret;
		}
		RasterBase* Sobel() {
			size_t total = lWidth * lHeight;
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, 1, datatype) };
			RasterPtr Sobelx(filter(SobelKernelx<float>()));
			RasterPtr Sobely(filter(SobelKernely<float>()));
			for (long long i = 0; i < total; ++i) {
				double x{ .0 }, y{ .0 };
				for (int channel = 0; channel < nBands; ++channel) {
					x += ((T*)Sobelx->data())[i + channel * total];
					y += ((T*)Sobely->data())[i + channel * total];
				}
				double s = sqrt(x * x + y * y);
				if (s > 255) s = 255;
				((T*)ret->data())[i] = (T)s;
			}
			return ret;
		}

		RasterBase* GaussBlur(size_t size, double stdx, double stdy)override {
			return filter(GaussianKernel(size, stdx, stdy));
		}
		RasterBase* addGaussNoise(size_t size, float std) {
			size_t total = lHeight * lWidth * nBands;
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, nBands, datatype) };
			memcpy(ret->data(), this->data(), total * sizeof(T));
			std::default_random_engine e;
			std::normal_distribution<float> x(0, std);
			for (long long i = 0; i < size; ++i) {
				int randomIndex = (rand() * rand()) % total;
				double s = ((T*)ret->data())[randomIndex] + x(e);
				if (s < 0) s = 0;
				if (s > 255) s = 255;
				((T*)ret->data())[randomIndex] = (T)(s);
			}
			return ret;
		}
		RasterBase* Laplace_4()override {
			return Laplace(LaplaceKernel_4<float>());
		}

		RasterBase* Laplace_8()override {
			return Laplace(LaplaceKernel_8<float>());
		}

		RasterBase* Laplace_4_Sharpening() {
			return filter(LaplaceSharpening_4<float>());
		}
		RasterBase* Laplace_8_Sharpening() {
			return filter(LaplaceSharpening_8<float>());
		}
		RasterBase* Sobel_Sharpening() {
			size_t total = lWidth * lHeight * nBands;
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, nBands, datatype) };
			RasterPtr Sobelx(filter(SobelKernelx<float>()));
			RasterPtr Sobely(filter(SobelKernely<float>()));
			for (long long i = 0; i < total; ++i) {
				T x = ((T*)Sobelx->data())[i];
				T y = ((T*)Sobely->data())[i];
				double s = sqrt(x * x + y * y) + pData[i];
				if (s < 0) s = fabs(s);
				if (s > 255) s = 255;
				((T*)ret->data())[i] = (T)(s);
			}
			return ret;
		}

		RasterBase* toGray() override {
			RasterBase* ret{ new Raster<T>(lWidth, lHeight, 1, datatype) };
			size_t total = lWidth * lHeight;
			double weights[] = { 0.299, 0.587, 0.114 };
			for (size_t i = 0; i < total; ++i) {
				double sum{ 0.0 };
				for (long long channel = 0; channel < nBands; ++channel) {
					long long index = i + channel * total;
					sum += weights[channel] * ((T*)pData)[index];
				}
				((T*)ret->data())[i] = (T)sum;
			}
			return ret;
		}

	public:
		Raster() {
			pData = nullptr;
			this->lHeight = -1;
			this->lHeight = -1;
			pDatasetRead = nullptr;
			pDriver = nullptr;
		}
		Raster(const string& filepath) :Raster(filepath.c_str()) {};
		Raster(const char* filepath) {
			assert(filepath != nullptr);
		}

		Raster(size_t width, size_t height, size_t nBands, GDALDataType datatype) {
			init(nullptr, width, height, nBands, datatype);
		}

		~Raster() {
			if (pDatasetRead)
				GDALClose(pDatasetRead);//关闭数据集
			if (pData)
				delete[] pData;
		}
	};

	RasterBase* getByType(GDALDataType datatype) {
		RasterBase* ret = nullptr;
		switch (datatype) {
		case GDT_Byte:
			ret = new Raster<unsigned char>();
			break;
		case GDT_UInt16:
			ret = new Raster<unsigned short>();
			break;
		case GDT_Int16:
			ret = new Raster<short>();
			break;
		case GDT_UInt32:
			ret = new Raster<unsigned int>();
			break;
		case GDT_Int32:
			ret = new Raster<int>();
			break;
		case GDT_Float32:
			ret = new Raster<float>();
			break;
		case GDT_Float64:
			ret = new Raster<double>();
			break;
		}
		assert(ret != nullptr);
		return ret;
	}

	RasterBase* open(const string& filepath) {
		GDALDataset* pDatasetRead{ nullptr };
		size_t lWidth{ 0 }, lHeight{ 0 }, nBands{ 0 };
		GDALDataType datatype{ GDT_Unknown };
		try {
			pDatasetRead = (GDALDataset*)GDALOpen(filepath.c_str(), GA_ReadOnly);
			lWidth = pDatasetRead->GetRasterXSize();//获取图像宽度
			lHeight = pDatasetRead->GetRasterYSize();//获取图像高度
			nBands = pDatasetRead->GetRasterCount();//获取图像波段数
			datatype = pDatasetRead->GetRasterBand(1)->GetRasterDataType();//获取图像格式类型
		}
		catch (std::exception& e) {
			cerr << e.what() << endl;
			exit(1);
		}
		RasterBase* ptr(getByType(datatype));
		ptr->init(pDatasetRead, lWidth, lHeight, nBands, datatype);
		return ptr;
	}


	RasterBase* operator-(RasterBase const& lhs, RasterBase const& rhs) {
		assert(lhs.width() == rhs.width() && lhs.height() == rhs.height() && lhs.bands() == rhs.bands());
		RasterBase* ret{ new Raster<unsigned char>(lhs.width(), lhs.height(), lhs.bands(), lhs.type()) };
		size_t total = lhs.width() * lhs.height() * lhs.bands();
		for (size_t i = 0; i < total; ++i) {
			double diff = ((unsigned char*)lhs.data())[i] - ((unsigned char*)rhs.data())[i];
			if (diff < 0) diff = 0;
			((unsigned char*)ret->data())[i] = (unsigned char)diff;
		}
		return ret;
	}
}
#endif // __GDAL__UTILITY__H__