分享一下常用的四种水下增强算法附源码

 一、直方图均衡化（Histogram Equalization，HE）：HE是一种用于增强图像对比度的经典技术。它通过重新分配图像像素的灰度级别，使得图像中的灰度值更加均匀分布，从而增强图像的整体对比度。

二、对比度有限自适应直方图均衡化（Contrast Limited Adaptive Histogram Equalization，CLAHE）：CLAHE是HE的改进版本，它在局部区域内应用直方图均衡化，以避免在图像中产生过度增强的噪声。

三、去雾算法（Dehazing，Dehaze）：水下图像常常受到水中悬浮颗粒的散射和吸收影响，导致图像中出现雾霾效果。去雾算法旨在通过估计和消除雾霾引起的图像退化，从而提高水下图像的清晰度和对比度。

四、多尺度反射消除算法（Multi-Scale Retinex with Color Restoration，MSRCR）：MSRCR算法基于多尺度的Retinex理论，通过估计图像的光照分量和反射分量来减少图像的光照不均匀性，从而改善图像的质量和视觉效果。

下图为浑浊度为10NTU左右水体30cm左右的成像

其中单张图像1080P实时性测试中：

HE:        5.95ms            

CLAHE:     2.99S     

DEHAZE:      83.77ms     

MSRCR:       2165.24ms

import cv2
import os
import numpy as np
import time

from tqdm import tqdm


"""HE方法"""
def he(image):
    B, G, R = cv2.split(image)
    B = cv2.equalizeHist(B)
    G = cv2.equalizeHist(G)
    R = cv2.equalizeHist(R)
    result = cv2.merge((B, G, R))
    return result




"""CLAHE方法"""
def clahe(image, clipLimit=2.0, tileGridSize=(8, 8)):
    B, G, R = cv2.split(image)
    clahe = cv2.createCLAHE(clipLimit, tileGridSize)
    clahe_B = clahe.apply(B)
    clahe_G = clahe.apply(G)
    clahe_R = clahe.apply(R)
    result = cv2.merge((clahe_B, clahe_G, clahe_R))
    return result



"""去雾方法"""
def zmMinFilterGray(src, r=7):
    return cv2.erode(src, np.ones((2 * r + 1, 2 * r + 1)))  


def guidedfilter(I, p, r, eps):
    height, width = I.shape
    m_I = cv2.boxFilter(I, -1, (r, r))
    m_p = cv2.boxFilter(p, -1, (r, r))
    m_Ip = cv2.boxFilter(I * p, -1, (r, r))
    cov_Ip = m_Ip - m_I * m_p
    m_II = cv2.boxFilter(I * I, -1, (r, r))
    var_I = m_II - m_I * m_I
    a = cov_Ip / (var_I + eps)
    b = m_p - a * m_I
    m_a = cv2.boxFilter(a, -1, (r, r))
    m_b = cv2.boxFilter(b, -1, (r, r))
    return m_a * I + m_b


def getV1(m, r, eps, w, maxV1):  # 输入rgb图像，值范围[0,1]
    '''计算大气遮罩图像V1和光照值A, V1 = 1-t/A'''
    V1 = np.min(m, 2)  # 得到暗通道图像
    V1 = guidedfilter(V1, zmMinFilterGray(V1, 7), r, eps)  # 使用引导滤波优化
    bins = 2000
    ht = np.histogram(V1, bins)  # 计算大气光照A
    d = np.cumsum(ht[0]) / float(V1.size)
    for lmax in range(bins - 1, 0, -1):
        if d[lmax] <= 0.999:
            break
    A = np.mean(m, 2)[V1 >= ht[1][lmax]].max()
    V1 = np.minimum(V1 * w, maxV1)  # 对值范围进行限制
    return V1, A


def defog(m, r=81, eps=0.001, w=0.95, maxV1=0.80, bGamma=False):
    m = m / 255.0
    Y = np.zeros(m.shape)
    V1, A = getV1(m, r, eps, w, maxV1)  # 得到遮罩图像和大气光照
    for k in range(3):
        Y[:, :, k] = (m[:, :, k] - V1) / (1 - V1 / A)  # 颜色校正
    Y = np.clip(Y, 0, 1)
    if bGamma:
        Y = Y ** (np.log(0.5) / np.log(Y.mean()))  # gamma校正,默认不进行该操作
    return Y * 255



"""MSRCR方法"""
def singleScaleRetinex(img, sigma):
    retinex = np.log10(img) - np.log10(cv2.GaussianBlur(img, (0, 0), sigma))
    return retinex


def multiScaleRetinex(img, sigma_list):
    retinex = np.zeros_like(img)
    for sigma in sigma_list:
        retinex += singleScaleRetinex(img, sigma)
    retinex = retinex / len(sigma_list)
    return retinex


def colorRestoration(img, alpha, beta):
    img_sum = np.sum(img, axis=2, keepdims=True)
    color_restoration = beta * (np.log10(alpha * img) - np.log10(img_sum))
    return color_restoration


def simplestColorBalance(img, low_clip, high_clip):
    total = img.shape[0] * img.shape[1]
    for i in range(img.shape[2]):
        unique, counts = np.unique(img[:, :, i], return_counts=True)
        current = 0
        for u, c in zip(unique, counts):
            if float(current) / total < low_clip:
                low_val = u
            if float(current) / total < high_clip:
                high_val = u
            current += c
        img[:, :, i] = np.maximum(np.minimum(img[:, :, i], high_val), low_val)
    return img


def MSRCR(img, sigma_list=[15, 80, 200], G=5.0, b=25.0, alpha=125.0, beta=46.0, low_clip=0.01, high_clip=0.99):
    img = np.float64(img) + 1.0
    img_retinex = multiScaleRetinex(img, sigma_list)
    img_color = colorRestoration(img, alpha, beta)
    img_msrcr = G * (img_retinex * img_color + b)
    for i in range(img_msrcr.shape[2]):
        img_msrcr[:, :, i] = (img_msrcr[:, :, i] - np.min(img_msrcr[:, :, i])) / \
                             (np.max(img_msrcr[:, :, i]) - np.min(img_msrcr[:, :, i])) * \
                             255
    img_msrcr = np.uint8(np.minimum(np.maximum(img_msrcr, 0), 255))
    img_msrcr = simplestColorBalance(img_msrcr, low_clip, high_clip)
    return img_msrcr
def process_image(image_path, output_folder):
    image = cv2.imread(image_path)

    
    he_image = he(image)
    clahe_image = clahe(image)
    defog_image = defog(image)
    msrcr_image = MSRCR(image)

    
    filename = os.path.splitext(os.path.basename(image_path))[0]
    cv2.imwrite(os.path.join(output_folder["he"], filename + ".jpg"), he_image)
    cv2.imwrite(os.path.join(output_folder["clahe"], filename + ".jpg"), clahe_image)
    cv2.imwrite(os.path.join(output_folder["defog"], filename + ".jpg"), defog_image)
    cv2.imwrite(os.path.join(output_folder["msrcr"], filename + ".jpg"), msrcr_image)

def batch_process_images(input_folder, output_folder):
    # 确保输出文件夹存在
    for folder in output_folder.values():
        if not os.path.exists(folder):
            os.makedirs(folder)


    image_files = [f for f in os.listdir(input_folder) if os.path.isfile(os.path.join(input_folder, f))]

    # 显示进度条
    for image_file in tqdm(image_files, desc="Processing Images"):
        image_path = os.path.join(input_folder, image_file)
        process_image(image_path, output_folder)


input_folder = r"C:\YOLO_Project\Convert\IMG_ENHANCEMENT-4\ImgEnhancement-4type\img"


output_folder = {
    "he": "he",
    "clahe": "clahe",
    "defog": "defog",
    "msrcr": "msrcr"
}

batch_process_images(input_folder, output_folder)

input_folder处输入相处理的图像文件夹路径

在图像路径img处新建四个生成文件夹路径，具体放置方法可以如下