暗通道先验DCP代码

对代码进行修改可直接将文件夹内的图像统一进行去雾操作

源代码GitHub - He-Zhang/image_dehaze: Single Image Haze Removal Using Dark Channel Prior

import math
import numpy as np
import sys
import cv2
import os

def DarkChannel(im,sz):
    b,g,r = cv2.split(im)
    dc = cv2.min(cv2.min(r,g),b);
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(sz,sz))
    dark = cv2.erode(dc,kernel)
    return dark

def AtmLight(im,dark):
    [h,w] = im.shape[:2]
    imsz = h*w
    numpx = int(max(math.floor(imsz/1000),1))
    darkvec = dark.reshape(imsz);
    imvec = im.reshape(imsz,3);

    indices = darkvec.argsort();
    indices = indices[imsz-numpx::]

    atmsum = np.zeros([1,3])
    for ind in range(1,numpx):
       atmsum = atmsum + imvec[indices[ind]]

    A = atmsum / numpx;
    return A

def TransmissionEstimate(im,A,sz):
    omega = 0.95;
    im3 = np.empty(im.shape,im.dtype);

    for ind in range(0,3):
        im3[:,:,ind] = im[:,:,ind]/A[0,ind]

    transmission = 1 - omega*DarkChannel(im3,sz);
    return transmission

def Guidedfilter(im,p,r,eps):
    mean_I = cv2.boxFilter(im,cv2.CV_64F,(r,r));
    mean_p = cv2.boxFilter(p, cv2.CV_64F,(r,r));
    mean_Ip = cv2.boxFilter(im*p,cv2.CV_64F,(r,r));
    cov_Ip = mean_Ip - mean_I*mean_p;

    mean_II = cv2.boxFilter(im*im,cv2.CV_64F,(r,r));
    var_I   = mean_II - mean_I*mean_I;

    a = cov_Ip/(var_I + eps);
    b = mean_p - a*mean_I;

    mean_a = cv2.boxFilter(a,cv2.CV_64F,(r,r));
    mean_b = cv2.boxFilter(b,cv2.CV_64F,(r,r));

    q = mean_a*im + mean_b;
    return q;

def TransmissionRefine(im,et):
    gray = cv2.cvtColor(im,cv2.COLOR_BGR2GRAY);
    gray = np.float64(gray)/255;
    r = 60;
    eps = 0.0001;
    t = Guidedfilter(gray,et,r,eps);

    return t;

def Recover(im,t,A,tx = 0.1):
    res = np.empty(im.shape,im.dtype);
    t = cv2.max(t,tx);

    for ind in range(0,3):
        res[:,:,ind] = (im[:,:,ind]-A[0,ind])/t + A[0,ind]

    return res


if __name__ == '__main__':
    try:
        fn = sys.argv[1]
    except:
        fn = r' '  # 修改为文件夹路径


    def nothing(*argv):
        pass


    output_dir = './output'  # 输出文件夹路径，可以根据需要自行修改
    if not os.path.exists(output_dir):  # 如果输出文件夹不存在，创建它
        os.makedirs(output_dir)
    img_dir = os.listdir(fn)  # 获取文件夹下的所有文件名
    for img_file in img_dir:
        if img_file.endswith('.png') or img_file.endswith('.jpg'):  # 仅处理png和jpg格式的图片
            src = cv2.imread(os.path.join(fn, img_file))  # 使用os.path.join连接路径和文件名
            ...  # 后续处理代码不变

            I = src.astype('float64') / 255;
            dark = DarkChannel(I, 15);
            A = AtmLight(I, dark);
            te = TransmissionEstimate(I, A, 15);
            t = TransmissionRefine(src, te);
            J = Recover(I, t, A, 0.1)  # 恢复图片并乘以255
            output_path = os.path.join(output_dir, os.path.splitext(img_file)[0] + '.png')  # 构建输出文件路径，以原文件名（去掉扩展名）加上输出路径和新的扩展名.png
            cv2.imwrite(output_path, J * 255)  # 写入输出文件夹中