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

后记

为了提高运算速度，尽量使用矩阵运算，当然以上还有需要改善的地方。
当然，本人水平有限，若有地方存在问题，还欢迎留言一起探讨！