Python 批量图像去雾 - 暗通道先验算法实现

以下是一个使用暗通道先验算法批量去雾并保存的 Python 代码示例：

import os
import cv2
import numpy as np

def dark_channel_prior(image, patch_size):
    min_channel = np.min(image, axis=2)
    dark_channel = cv2.erode(min_channel, np.ones((patch_size, patch_size)))
    return dark_channel

def estimate_atmospheric_light(image, dark_channel, top_percentage):
    flat_dark_channel = dark_channel.flatten()
    num_pixels = int(flat_dark_channel.shape[0] * top_percentage / 100)
    sorted_indices = np.argsort(-flat_dark_channel)
    sorted_indices = sorted_indices[:num_pixels]
    atmospheric_light = np.max(image.reshape(-1, 3)[sorted_indices], axis=0)
    return atmospheric_light

def transmission_estimate(image, atmospheric_light, omega, patch_size):
    normalized_image = image.astype(np.float32) / atmospheric_light.astype(np.float32)
    transmission = 1 - omega * dark_channel_prior(normalized_image, patch_size)
    return transmission

def recover_scene(image, transmission, atmospheric_light, t0):
    transmission = np.maximum(transmission, t0)
    scene_radiance = np.zeros_like(image, dtype=np.float32)
    for i in range(3):
        scene_radiance[:,:,i] = (image[:,:,i] - atmospheric_light[i]) / transmission + atmospheric_light[i]
    scene_radiance = np.clip(scene_radiance, 0, 255).astype(np.uint8)
    return scene_radiance

def dehaze_images(input_folder, output_folder, omega=0.95, patch_size=15, top_percentage=0.1, t0=0.1):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for filename in os.listdir(input_folder):
        if filename.endswith('.jpg') or filename.endswith('.png'):
            image_path = os.path.join(input_folder, filename)
            image = cv2.imread(image_path)
            dark_channel = dark_channel_prior(image, patch_size)
            atmospheric_light = estimate_atmospheric_light(image, dark_channel, top_percentage)
            transmission = transmission_estimate(image, atmospheric_light, omega, patch_size)
            scene_radiance = recover_scene(image, transmission, atmospheric_light, t0)

            output_path = os.path.join(output_folder, filename)
            cv2.imwrite(output_path, scene_radiance)

# 设置输入文件夹和输出文件夹路径
input_folder = 'input_images'
output_folder = 'output_images'

# 调用批量去雾函数
dehaze_images(input_folder, output_folder)

请注意，这只是一个简单的示例代码，暗通道先验算法的效果可能会因为不同的参数设置和图像质量而有所不同。您可以根据需要自行调整参数。