DCPNet复现：批量处理图片

import os

import cv2
import numpy as np
import matplotlib.pyplot as plt
import math
from skimage.metrics import structural_similarity as sk_cpt_ssim

from skimage.metrics import structural_similarity as compare_ssim
from skimage.metrics import peak_signal_noise_ratio as compare_psnr
from skimage.metrics import mean_squared_error as compare_mse

def guided_filter(I,p,win_size,eps):

    mean_I = cv2.blur(I,(win_size,win_size))
    mean_p = cv2.blur(p,(win_size,win_size))

    corr_I = cv2.blur(I*I,(win_size,win_size))
    corr_Ip = cv2.blur(I*p,(win_size,win_size))

    var_I = corr_I-mean_I*mean_I
    cov_Ip = corr_Ip - mean_I*mean_p

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

    mean_a = cv2.blur(a,(win_size,win_size))
    mean_b = cv2.blur(b,(win_size,win_size))

    q = mean_a*I + mean_b
    return q
def get_min_channel(img):
    return np.min(img,axis=2)

def min_filter(img,r):
    kernel = np.ones((2*r-1,2*r-1))
    return cv2.erode(img,kernel, 2)#最小值滤波器，可用腐蚀替代

def get_A(img_haze,dark_channel,bins_l):  #在get_A中已经实现了从三通道计算A，即设置axis=0

    hist,bins = np.histogram(dark_channel,bins=bins_l) # 得到直方图 hist ：直方图的值  bin_edges： dtype 浮点数组 bins 中的每个元素表示对应箱子的右边缘的值
    d = np.cumsum(hist)/float(dark_channel.size) # 累加 通过 np.cumsum 函数对直方图进行累加，计算每个阈值下的像素分布比例。在这里，变量d存储了每个阈值下的像素分布比例。len(d) = 2000

    threshold=0
    for i in range(bins_l-1,0,-1):
        # print('testdiishere', d[i])
        if d[i]<=0.999:
            threshold=i
            break

    A = np.max(img_haze[dark_channel>=bins[threshold]], axis=0) #axis=0意思是沿着第 0 维（即行）求最大值
    A = A.max()
    #候选区域可视化
    show  = np.copy(img_haze)
    show[dark_channel>=bins[threshold]] = 0,0,255
    return A

def optimize_lambda(img_haze, dark_channel, mean_intensity, bins, threshold, init_lambda=0.6, lr=0.1, max_iter=1000, eps=1e-4):

    lambda1 = init_lambda
    for i in range(max_iter):
        grad = (objective_function(lambda1 + eps, mean_intensity,img_haze, dark_channel, bins, threshold) -
                objective_function(lambda1 - eps, mean_intensity,img_haze, dark_channel, bins, threshold)) / (2 * eps)
        lambda1 -= lr * grad
        if abs(grad) < eps:
            break
    return lambda1

def objective_function(lambda1, mean_intensity, img_haze, dark_channel, bins, threshold):
    I = cv2.imread('testdcp5.jpg') / 255.0  # 归一化
    C = cv2.imread('testdcp5.jpg')/255.0
    A = lambda1 * mean_intensity + (1 - lambda1) * img_haze[dark_channel >= bins[threshold]].max()

    t = get_t(I, A)
    t = guided_filter(dark_channel, t, 81, 0.001)
    t = t[:, :, np.newaxis].repeat(3, axis=2)  # 升维至(r,w,3)

    J = (I-A)/t +A
    J = np.clip(J,0,1)
    J = J*255
    J =np.uint8(J)
    ssim2 = sk_cpt_ssim(J, C * 255, win_size=11, data_range=255, channel_axis=2)
    print(f'geta3-ssim2 = %s, lambda1 = %s' %(ssim2, lambda1))
    return -ssim2


# 得到三个通道的大气光
def get_atomospheric_light(I, dark_channel_1):
    """
    calculate the atomospheric light A
    :param I: the observed intensity image I
    :return: A
    """
    # step1: get the dark_channel of image I
    I_dark = dark_channel_1

    # step2: get the top 0.1% brightest pixels index in the dark channel
    h, w = I_dark.shape
    top_pixels = h * w // 1000

    # define a function to find the N maximum indices of a NumPy array
    def largest_indices(ary, n):
        """Returns the n largest indices from a numpy array."""
        flat = ary.flatten()
        indices = np.argpartition(flat, -n)[-n:]
        indices = indices[np.argsort(-flat[indices])]
        return np.unravel_index(indices, ary.shape)

    bright_index = largest_indices(I_dark, top_pixels)
    A = np.zeros(shape=(3,))
    A[0] = np.max(I[bright_index][0])
    A[1] = np.max(I[bright_index][1])
    A[2] = np.max(I[bright_index][2])
    return A

def estimation_A(image, dark):
    """
    get A by haze image and dark channel
    Args:
        image: haze image with numpy type
        dark: dark channel of haze image with gray type

    Returns:
        air light A, scalar
    """
    if not isinstance(image, np.ndarray):
        raise ValueError("input image is not numpy type")
    if not len(dark.shape) == 2:
        raise ValueError("the dark channel should be gray scale")
    # calculate the number of pixels of the top 0.1%
    top_pixels_number = int(image.shape[0]*image.shape[1]*0.001)

    # reshape image and dark into shape (-1, 3) and (-1)
    # copy the data, in case of it changed
    image_copy = image.copy()
    dark_copy = dark.copy()
    image_copy = image_copy.reshape(-1, 3)
    dark_copy = dark_copy.reshape(-1)


    index_of_sort = dark_copy.argsort()[::-1]

    index_of_sort = index_of_sort[0:top_pixels_number]


    A = image_copy[index_of_sort]

    A = np.mean(A, axis=0)
    return A

def get_t(img_haze,A,t0=0.1,w=0.95):
    out = get_min_channel(img_haze)
    out = min_filter(out,r=7)

    t = 1-w*out/A #需要乘上一系数w，为远处的物体保留少量的雾

    t = np.clip(t,t0,1)#论文4.4所提到t(x)趋于0容易产生噪声，所以设置一最小值0.1 t的值和暗通道一样，每一个像素点有一个值

    return t

def PSNR(target,ref):
    #必须归一化
    target=target/255.0
    ref=ref/255.0
    MSE = np.mean((target-ref)**2)
    if MSE<1e-10:
        return 100
    MAXI=1
    PSNR = 20*math.log10(MAXI/math.sqrt(MSE))
    return PSNR

def handledcp(image_addr, name):

    I = cv2.imread(image_addr) / 255.0  # 归一化
    dark_channel = get_min_channel(I)
    dark_channel_1 = min_filter(dark_channel, r=7)
    A = get_A(I, dark_channel_1, bins_l=2000)
    t = get_t(I, A)
    t = guided_filter(dark_channel, t, 81, 0.001)
    t = t[:, :, np.newaxis].repeat(3, axis=2)  # 升维至(r,w,3)
    J = (I - A) / t + A
    J = np.clip(J, 0, 1)
    J = J * 255
    J = np.uint8(J)
    addr_save = "D:/xxxxx/pythonProject/dcpnet/gt/" + name + '.png'
    cv2.imwrite(addr_save, J)


def dehazeFile(img_Dir):
    img_pathDir = os.listdir(img_Dir)  # 提取所有文件名，并存在列表中
    print("img_Dir:", img_Dir)  # 输出文件路径
    print("img_pathDir:", img_pathDir)  # 输出文件名列表
    print(len(img_pathDir))  # 输出文件数
    num = 0
    for i in img_pathDir:
        addr = 'D:/xxxxx/pythonProject/dcpnet/hazy/' + i
        # print(addr)
        name = i.split("_")[0]
        handledcp(addr, name)
        num = num + 1
        print('handling:{}/{}'.format(num, len(img_pathDir)))
    return


if __name__ == '__main__':
    img_Dir = r"D:\xxxx\pythonProject\dcpnet\hazy" #本地文件路径
    dehazeFile(img_Dir)