Opencv实现暗通道先验去雾算法

今天读了何凯明博士的《Single Image Haze Removal Using Dark Channel Prior》，用opencv实现了一遍。

其中，暗通道及T(x)采用腐蚀erode算法。

具体代码如下：

class hazeOutProcesser
{
public:

hazeOutProcesser(const Mat& srcImg)
{
this->srcImg = srcImg.clone();
}

hazeOutProcesser(){}

void getDarkChannel(const Mat& src, Mat& dst)
{
Mat element = getStructuringElement(MORPH_RECT, Size(15, 15));
// erode(src, dst, element);
morphologyEx(src, dst, MORPH_ERODE, element);
// Mat temp = dst.clone();
// bilateralFilter(temp, dst, 25, 25*2, 25/2);
cvtColor(dst, dst, COLOR_BGR2GRAY);
}

std::pair<Mat, Mat> getAtmosphere(const Mat& sortedIndex,
const Mat& srcImg,
double rate = 0.001)
{
long n_pixel = srcImg.rows * srcImg.cols;
long n_search = n_pixel * rate;
Mat srcImgVector = srcImg.reshape(1, n_pixel).clone(); //1channel, 3rows(b, g, r)
//(n_pixel, 3)
cout << srcImgVector.rows << ", " << srcImgVector.cols << endl;
Mat accumulator = Mat::zeros(1, 3, CV_32FC1); //[(32f)0, (32f)0, (32f)0]
for(int i = 0; i < n_search; ++i) {
int index = sortedIndex.at<int>(i);
Mat temp = srcImgVector.rowRange(index, index + 1).clone();   //row:index
accumulator += temp;
}
Mat atmosphere = accumulator / n_search;
for(int i = 0; i < n_pixel; ++i) {
srcImgVector.rowRange(i, i + 1) = srcImgVector.rowRange(i, i + 1).mul(1 / atmosphere);
}
srcImgVector = srcImgVector.reshape(3).reshape(3, srcImg.rows);
// imshow("srcImgVector", srcImgVector);
// cout << srcImgVector.rows << ", " << srcImgVector.cols << endl;

4000
return std::make_pair(atmosphere, srcImgVector);
}

Mat getTransmission(const Mat& srcImg_divide_atmosphere)
{
float omega = 0.95;
Mat darkchannel_SDA;
getDarkChannel(srcImg_divide_atmosphere, darkchannel_SDA);
darkchannel_SDA *= omega; darkchannel_SDA = 1 - darkchannel_SDA;
// imshow("sda", darkchannel_SDA);
return darkchannel_SDA;
}

Mat hazeOut(const Mat& srcImg, const Mat& T, const Mat& A)
{
const float T0 = 0.1;
// long n_pixel = srcImg.rows * srcImg.cols;
// Mat srcImgVector = srcImg.reshape(1, n_pixel);

// for(int i = 0; i < n_pixel; ++i) {
//  srcImgVector.rowRange(i, i + 1) -= A;
// }
/*-------------------need improve-------------------
I use a loop to calculate the result matrix
but using maxtrix directly might speed up the procession
*/

Mat hazeout(srcImg.size(), CV_32FC3);
for(size_t r = 0; r < hazeout.rows; ++r) {
for(size_t c = 0; c < hazeout.cols; ++c) {
hazeout.at<Vec3f>(r, c) =
Vec3f((srcImg.at<Vec3f>(r, c)[0] - A.at<float_t>(0, 0)) / max(T.at<float>(r, c), T0) + A.at<float_t>(0, 0),
(srcImg.at<Vec3f>(r, c)[1] - A.at<float_t>(0, 1)) / max(T.at<float>(r, c), T0) + A.at<float_t>(0, 1),
(srcImg.at<Vec3f>(r, c)[2] - A.at<float_t>(0, 2)) / max(T.at<float>(r, c), T0) + A.at<float_t>(0, 2));
}
}

return hazeout;
}

void process()
{
getDarkChannel(srcImg, darkchannel);

Mat matVector = darkchannel.reshape(1, 1);                       //(1, 116749)
// cout << matVector.rows << ", " << matVector.cols << endl;
Mat_<int> sortIndex;
sortIdx(darkchannel,
sortIndex,
CV_SORT_EVERY_ROW | SORT_DESCENDING);
std::pair<Mat, Mat> p = getAtmosphere(sortIndex, srcImg);
/*-------check minimum-------------
double min = 0;
minMaxLoc(transmission, &min);
cout << "min:" << min << endl;
*/
Mat transmission = getTransmission(p.second);
hazeOut(srcImg, transmission, p.first).convertTo(dstImg, CV_8UC3, 255, 0);
}

void setInput(const Mat& input)
{
srcImg = input.clone();
}

Mat getDarkChannelOutput()
{
return darkchannel;
}

Mat getOutput()
{
return dstImg;
}

private:

Mat srcImg;
Mat dstImg;
Mat darkchannel;
};

#include <opencv2/opencv.hpp>
using namespace cv;
using namespace std;
int main(int argc, char const *argv[])
{
Mat srcImg = imread("../../../data/sea1.jpg");
assert(srcImg.data && srcImg.channels() == 3);
srcImg.convertTo(srcImg, CV_32FC3, 1.0 / 255, 0);
hazeOutProcesser processer;
processer.setInput(srcImg);
processer.process();
imshow("darkchannel1", processer.getDarkChannelOutput());
imshow("src1", srcImg);
imshow("hazeout1", processer.getOutput());

srcImg = imread("../../../data/sea2.png");
assert(srcImg.data && srcImg.channels() == 3);
srcImg.convertTo(srcImg, CV_32FC3, 1.0 / 255, 0);
processer.setInput(srcImg);
processer.process();
// imshow("darkchannel2", processer.getDarkChannelOutput());
imshow("src2", srcImg);
imshow("hazeout2", processer.getOutput());

srcImg = imread("../../../data/sea4.jpg");
assert(srcImg.data && srcImg.channels() == 3);
resize(srcImg, srcImg, Size(), 0.5, 0.5);
srcImg.convertTo(srcImg, CV_32FC3, 1.0 / 255, 0);
processer.setInput(srcImg);
processer.process();
// imshow("darkchannel2", processer.getDarkChannelOutput());
imshow("src3", srcImg);
imshow("hazeout3", processer.getOutput());
waitKey();
return 0;
}

后记

为了提高运算速度，尽量使用矩阵运算，当然以上还有需要改善的地方。

当然，本人水平有限，若有地方存在问题，还欢迎留言一起探讨！