【ISP算法精粹】什么是global tone mapping和local tone mapping?
1. 简介
全局色调映射(Global Tone Mapping)和局部色调映射(Local Tone Mapping)是高动态范围(HDR)图像处理中的两种关键技术,用于将高动态范围图像的亮度值映射到标准动态范围(LDR)内,同时保留图像的细节和视觉质量。
全局色调映射(Global Tone Mapping)
全局色调映射对图像中的所有像素应用相同的映射函数,常见方法包括:
- 线性映射:将图像的亮度范围直接线性缩放到显示设备的动态范围内。
- 对数映射:使用对数函数压缩高亮度区域,扩展低亮度区域。
- 伽马校正:调整图像的亮度分布,常用于显示设备的非线性响应补偿。
局部色调映射(Local Tone Mapping)
局部色调映射考虑图像的局部特征,对不同区域应用不同的映射函数,常见方法包括:
- 直方图均衡化:通过调整图像的直方图来增强对比度。
- 拉普拉斯滤波:基于图像的局部梯度信息进行处理。
- 多尺度分解:将图像分解为不同尺度的分量,分别处理后再合并。
下面是使用Python实现这两种色调映射方法的代码:
import numpy as np
import cv2
from matplotlib import pyplot as pltdef global_tone_mapping(image, gamma=2.2, exposure=1.0):"""全局色调映射实现参数:image: numpy数组,输入的HDR图像gamma: 伽马值,用于调整亮度分布exposure: 曝光值,用于调整整体亮度返回:numpy数组,处理后的LDR图像"""# 应用曝光调整exposed = np.clip(image * exposure, 0, None)# 应用伽马校正ldr_image = np.power(exposed, 1.0 / gamma)# 将值归一化到[0, 1]范围ldr_image = np.clip(ldr_image, 0, 1)return ldr_imagedef local_tone_mapping(image, sigma=15, contrast=1.0):"""局部色调映射实现(基于拉普拉斯滤波)参数:image: numpy数组,输入的HDR图像sigma: 高斯核标准差,控制局部区域大小contrast: 对比度增强因子返回:numpy数组,处理后的LDR图像"""# 将图像转换为对数空间log_image = np.log1p(image)# 计算全局平均亮度global_mean = np.mean(log_image)# 计算局部亮度(使用高斯滤波)local_mean = cv2.GaussianBlur(log_image, (0, 0), sigma)# 计算局部对比度(拉普拉斯算子)laplacian = cv2.Laplacian(log_image, cv2.CV_64F)# 应用局部色调映射tonemapped = global_mean + contrast * (log_image - local_mean) + 0.5 * laplacian# 将图像转回线性空间ldr_image = np.expm1(tonemapped)# 将值归一化到[0, 1]范围ldr_image = np.clip(ldr_image / np.max(ldr_image), 0, 1)return ldr_imagedef main():# 创建示例HDR图像(这里使用合成图像)# 在实际应用中,你可能需要使用OpenCV或其他库读取真实的HDR图像h, w = 200, 200x = np.linspace(0, 1, w)y = np.linspace(0, 1, h)xx, yy = np.meshgrid(x, y)# 创建具有高动态范围的图像hdr_image = np.zeros((h, w, 3))hdr_image[:, :, 0] = 5.0 * np.exp(-10 * (xx - 0.3)**2 - 10 * (yy - 0.7)**2) # 红色高光hdr_image[:, :, 1] = 2.0 * np.exp(-5 * (xx - 0.7)**2 - 5 * (yy - 0.3)**2) # 绿色高光hdr_image[:, :, 2] = 1.0 * (xx + yy) # 蓝色渐变# 应用全局色调映射global_ldr = global_tone_mapping(hdr_image, gamma=2.2, exposure=0.5)# 应用局部色调映射local_ldr = local_tone_mapping(hdr_image, sigma=15, contrast=1.5)# 显示结果plt.figure(figsize=(15, 5))plt.subplot(131)plt.title('HDR Image (Scaled)')plt.imshow(np.clip(hdr_image / 5.0, 0, 1)) # 缩放以显示HDR图像plt.axis('off')plt.subplot(132)plt.title('Global Tone Mapping')plt.imshow(global_ldr)plt.axis('off')plt.subplot(133)plt.title('Local Tone Mapping')plt.imshow(local_ldr)plt.axis('off')plt.tight_layout()plt.show()if __name__ == "__main__":main()
这段代码实现了两种色调映射方法:
global_tone_mapping
函数通过曝光调整和伽马校正实现全局色调映射。local_tone_mapping
函数通过对数变换、高斯滤波和拉普拉斯算子实现局部色调映射。
主函数创建了一个合成的HDR图像,并对其应用两种色调映射方法,最后显示原始HDR图像和处理后的LDR图像进行对比。在实际应用中,你可能需要使用OpenCV或其他库读取真实的HDR图像文件。
2. 实现全局色调映射的代码中,gamma值如何选择?
在全局色调映射中,gamma值的选择直接影响图像的亮度分布和视觉效果。以下是关于gamma值选择的详细说明:
gamma值的基本原理
- gamma>1:降低图像亮度,使暗部更暗,适合处理过亮的图像。
- gamma=1:不改变图像亮度(线性映射)。
- gamma<1:提高图像亮度,使亮部更亮,适合处理过暗的图像。
如何选择合适的gamma值
-
基于显示设备特性
- 标准sRGB显示设备通常使用gamma=2.2。
- 若显示设备支持线性色彩空间,则使用gamma=1.0。
-
基于图像内容特性
- 高对比度图像:使用较大的gamma值(如2.2-2.6)压缩动态范围。
- 低对比度图像:使用较小的gamma值(如1.8-2.0)增强对比度。
-
经验公式
- 自动计算gamma值:
gamma = log(0.5) / log(mean_intensity)
- 其中
mean_intensity
是图像的平均亮度(范围0-1)。
- 自动计算gamma值:
下面是一个改进的全局色调映射实现,增加了自动gamma计算功能:
import numpy as np
import cv2
from matplotlib import pyplot as pltdef global_tone_mapping(image, gamma=None, exposure=1.0):"""全局色调映射实现参数:image: numpy数组,输入的HDR图像gamma: 伽马值,用于调整亮度分布。若为None,则自动计算exposure: 曝光值,用于调整整体亮度返回:numpy数组,处理后的LDR图像"""# 应用曝光调整exposed = np.clip(image * exposure, 0, None)# 自动计算gamma值(如果未指定)if gamma is None:# 计算图像的平均亮度gray = cv2.cvtColor(exposed, cv2.COLOR_RGB2GRAY) if exposed.ndim == 3 else exposedmean_intensity = np.mean(gray)# 避免除以零if mean_intensity > 0:gamma = np.log(0.5) / np.log(mean_intensity)else:gamma = 2.2 # 默认值# 应用伽马校正ldr_image = np.power(exposed, 1.0 / gamma)# 将值归一化到[0, 1]范围ldr_image = np.clip(ldr_image, 0, 1)return ldr_imagedef local_tone_mapping(image, sigma=15, contrast=1.0):"""局部色调映射实现(基于拉普拉斯滤波)参数:image: numpy数组,输入的HDR图像sigma: 高斯核标准差,控制局部区域大小contrast: 对比度增强因子返回:numpy数组,处理后的LDR图像"""# 将图像转换为对数空间log_image = np.log1p(image)# 计算全局平均亮度global_mean = np.mean(log_image)# 计算局部亮度(使用高斯滤波)local_mean = cv2.GaussianBlur(log_image, (0, 0), sigma)# 计算局部对比度(拉普拉斯算子)laplacian = cv2.Laplacian(log_image, cv2.CV_64F)# 应用局部色调映射tonemapped = global_mean + contrast * (log_image - local_mean) + 0.5 * laplacian# 将图像转回线性空间ldr_image = np.expm1(tonemapped)# 将值归一化到[0, 1]范围ldr_image = np.clip(ldr_image / np.max(ldr_image), 0, 1)return ldr_imagedef main():# 创建示例HDR图像(这里使用合成图像)# 在实际应用中,你可能需要使用OpenCV或其他库读取真实的HDR图像h, w = 200, 200x = np.linspace(0, 1, w)y = np.linspace(0, 1, h)xx, yy = np.meshgrid(x, y)# 创建具有高动态范围的图像hdr_image = np.zeros((h, w, 3))hdr_image[:, :, 0] = 5.0 * np.exp(-10 * (xx - 0.3)**2 - 10 * (yy - 0.7)**2) # 红色高光hdr_image[:, :, 1] = 2.0 * np.exp(-5 * (xx - 0.7)**2 - 5 * (yy - 0.3)**2) # 绿色高光hdr_image[:, :, 2] = 1.0 * (xx + yy) # 蓝色渐变# 应用全局色调映射(自动gamma)global_ldr_auto = global_tone_mapping(hdr_image, exposure=0.5)# 应用全局色调映射(手动gamma)global_ldr_manual = global_tone_mapping(hdr_image, gamma=2.2, exposure=0.5)# 应用局部色调映射local_ldr = local_tone_mapping(hdr_image, sigma=15, contrast=1.5)# 显示结果plt.figure(figsize=(15, 10))plt.subplot(221)plt.title('HDR Image (Scaled)')plt.imshow(np.clip(hdr_image / 5.0, 0, 1)) # 缩放以显示HDR图像plt.axis('off')plt.subplot(222)plt.title('Global Tone Mapping (Auto Gamma)')plt.imshow(global_ldr_auto)plt.axis('off')plt.subplot(223)plt.title('Global Tone Mapping (Manual Gamma=2.2)')plt.imshow(global_ldr_manual)plt.axis('off')plt.subplot(224)plt.title('Local Tone Mapping')plt.imshow(local_ldr)plt.axis('off')plt.tight_layout()plt.show()if __name__ == "__main__":main()
这个改进版本的全局色调映射函数允许:
- 当gamma参数为None时,自动根据图像平均亮度计算gamma值。
- 保留手动指定gamma值的选项,以满足特定需求。
在实际应用中,你可以根据图像内容和显示设备特性调整gamma值,或者使用自动计算功能获得较好的默认效果。
3. 局部色调映射怎么优化?
局部色调映射的优化可以从多个维度进行,包括算法改进、计算效率提升和视觉效果增强等方面。以下是几种常见的优化方法及其Python实现:
1. 多尺度分解优化
将图像分解为不同尺度的细节层和基础层,对每层应用不同的处理参数,保留更多细节的同时避免光晕现象。
2. 自适应滤波
使用导向滤波或双边滤波替代简单的高斯滤波,更好地保留边缘信息。
3. 对比度限制
引入对比度限制机制,防止局部对比度过度增强导致的伪影。
下面是优化后的局部色调映射实现:
import numpy as np
import cv2
from matplotlib import pyplot as pltdef global_tone_mapping(image, gamma=None, exposure=1.0):"""全局色调映射实现参数:image: numpy数组,输入的HDR图像gamma: 伽马值,用于调整亮度分布。若为None,则自动计算exposure: 曝光值,用于调整整体亮度返回:numpy数组,处理后的LDR图像"""# 应用曝光调整exposed = np.clip(image * exposure, 0, None)# 自动计算gamma值(如果未指定)if gamma is None:# 计算图像的平均亮度gray = cv2.cvtColor(exposed, cv2.COLOR_RGB2GRAY) if exposed.ndim == 3 else exposedmean_intensity = np.mean(gray)# 避免除以零if mean_intensity > 0:gamma = np.log(0.5) / np.log(mean_intensity)else:gamma = 2.2 # 默认值# 应用伽马校正ldr_image = np.power(exposed, 1.0 / gamma)# 将值归一化到[0, 1]范围ldr_image = np.clip(ldr_image, 0, 1)return ldr_imagedef guided_filter(I, p, r, eps):"""导向滤波实现参数:I: 导向图像p: 输入图像r: 滤波半径eps: 正则化参数返回:numpy数组,滤波结果"""# 计算导向图像的均值mean_I = cv2.boxFilter(I, -1, (r, r))# 计算输入图像的均值mean_p = cv2.boxFilter(p, -1, (r, r))# 计算I和p的协方差corr_I = cv2.boxFilter(I * I, -1, (r, r))corr_Ip = cv2.boxFilter(I * p, -1, (r, r))# 计算方差和协方差var_I = corr_I - mean_I * mean_Icov_Ip = corr_Ip - mean_I * mean_p# 计算线性系数a = cov_Ip / (var_I + eps)b = mean_p - a * mean_I# 计算系数的均值mean_a = cv2.boxFilter(a, -1, (r, r))mean_b = cv2.boxFilter(b, -1, (r, r))# 输出结果q = mean_a * I + mean_breturn qdef optimized_local_tone_mapping(image, scales=3, sigma_base=15, sigma_detail=5, contrast=1.0, clip_limit=1.5):"""优化的局部色调映射实现(基于多尺度分解和导向滤波)参数:image: numpy数组,输入的HDR图像scales: 分解的尺度数sigma_base: 基础层滤波的标准差sigma_detail: 细节层滤波的标准差contrast: 对比度增强因子clip_limit: 局部对比度限制因子返回:numpy数组,处理后的LDR图像"""# 将图像转换为对数空间log_image = np.log1p(image)# 计算全局平均亮度global_mean = np.mean(log_image)# 初始化输出图像tonemapped = np.zeros_like(log_image)# 多尺度分解for s in range(scales):# 计算当前尺度的权重weight = 1.0 / (2 ** s)# 计算当前尺度的滤波参数sigma = sigma_base * (2 ** s)# 使用导向滤波计算基础层if s == 0:# 第一层使用原图作为导向if log_image.ndim == 3:# 对于彩色图像,使用亮度通道作为导向gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)base = np.stack([guided_filter(gray, log_image[:,:,c], int(sigma), 0.01) for c in range(3)], axis=2)else:base = guided_filter(log_image, log_image, int(sigma), 0.01)else:# 后续层使用上一层的基础层作为导向if base.ndim == 3:gray_base = cv2.cvtColor(np.expm1(base), cv2.COLOR_RGB2GRAY)base = np.stack([guided_filter(gray_base, log_image[:,:,c], int(sigma), 0.01) for c in range(3)], axis=2)else:base = guided_filter(base, log_image, int(sigma), 0.01)# 计算细节层detail = log_image - base# 对比度限制detail_std = np.std(detail)detail = np.clip(detail, -clip_limit * detail_std, clip_limit * detail_std)# 应用细节增强enhanced_detail = detail * contrast# 累加结果tonemapped += weight * (global_mean + enhanced_detail)# 将图像转回线性空间ldr_image = np.expm1(tonemapped)# 将值归一化到[0, 1]范围ldr_image = np.clip(ldr_image / np.max(ldr_image), 0, 1)return ldr_imagedef main():# 创建示例HDR图像(这里使用合成图像)# 在实际应用中,你可能需要使用OpenCV或其他库读取真实的HDR图像h, w = 200, 200x = np.linspace(0, 1, w)y = np.linspace(0, 1, h)xx, yy = np.meshgrid(x, y)# 创建具有高动态范围的图像hdr_image = np.zeros((h, w, 3))hdr_image[:, :, 0] = 5.0 * np.exp(-10 * (xx - 0.3)**2 - 10 * (yy - 0.7)**2) # 红色高光hdr_image[:, :, 1] = 2.0 * np.exp(-5 * (xx - 0.7)**2 - 5 * (yy - 0.3)**2) # 绿色高光hdr_image[:, :, 2] = 1.0 * (xx + yy) # 蓝色渐变# 应用全局色调映射(自动gamma)global_ldr_auto = global_tone_mapping(hdr_image, exposure=0.5)# 应用原始局部色调映射local_ldr_original = local_tone_mapping(hdr_image, sigma=15, contrast=1.5)# 应用优化的局部色调映射local_ldr_optimized = optimized_local_tone_mapping(hdr_image, scales=3, sigma_base=15, sigma_detail=5, contrast=1.5, clip_limit=1.5)# 显示结果plt.figure(figsize=(15, 10))plt.subplot(221)plt.title('HDR Image (Scaled)')plt.imshow(np.clip(hdr_image / 5.0, 0, 1)) # 缩放以显示HDR图像plt.axis('off')plt.subplot(222)plt.title('Global Tone Mapping')plt.imshow(global_ldr_auto)plt.axis('off')plt.subplot(223)plt.title('Original Local Tone Mapping')plt.imshow(local_ldr_original)plt.axis('off')plt.subplot(224)plt.title('Optimized Local Tone Mapping')plt.imshow(local_ldr_optimized)plt.axis('off')plt.tight_layout()plt.show()if __name__ == "__main__":main()
这个优化版本的局部色调映射具有以下特点:
-
多尺度分解:将图像分解为多个尺度的基础层和细节层,每层使用不同的处理参数。
-
导向滤波:替代简单的高斯滤波,更好地保留边缘信息,减少光晕现象。
-
对比度限制:对细节层进行对比度限制,防止局部对比度过度增强导致的伪影。
-
自适应权重:不同尺度层使用不同的权重,平衡全局和局部效果。
你可以通过调整scales
、sigma_base
、sigma_detail
、contrast
和clip_limit
等参数来获得不同的视觉效果。例如:
- 增加
scales
可以保留更多细节,但会增加计算复杂度。 - 调整
clip_limit
可以控制局部对比度的增强程度。 - 修改
sigma_base
和sigma_detail
可以调整滤波的精细程度。
这些优化方法能够显著提升局部色调映射的效果,特别是在保留细节和减少伪影方面。
4. 还有哪些优化局部色调映射的思路?
除了前面提到的多尺度分解、导向滤波和对比度限制等方法,局部色调映射还有许多其他优化思路。以下是一些进阶的优化方向及其实现原理:
1. 基于视网膜感知模型的优化
模仿人眼视网膜的处理机制,将图像分为亮度通道和颜色通道分别处理:
- 亮度通道:使用对数变换压缩动态范围,保留细节。
- 颜色通道:通过色度保持算法避免颜色失真。
def retina_tone_mapping(image, alpha=0.5, beta=0.5):"""基于视网膜感知模型的色调映射"""# 分离亮度和色度lab = cv2.cvtColor((image * 255).astype(np.uint8), cv2.COLOR_RGB2LAB)L, a, b = cv2.split(lab)# 对数压缩亮度通道L_log = np.log1p(L.astype(np.float32) / 100.0)L_mapped = (L_log / np.max(L_log) * 100).astype(np.uint8)# 重构LAB图像并转回RGBlab_mapped = cv2.merge([L_mapped, a, b])rgb_mapped = cv2.cvtColor(lab_mapped, cv2.COLOR_LAB2RGB)return rgb_mapped.astype(np.float32) / 255.0
2. 直方图优化技术
通过改进的直方图处理增强局部对比度:
- 自适应直方图均衡化(CLAHE):将图像分块处理,避免全局直方图均衡化的过度增强问题。
- 双直方图均衡化:分别处理亮度的上下部分,保留更多细节。
def clahe_tone_mapping(image, clip_limit=2.0, tile_grid_size=(8, 8)):"""使用CLAHE进行局部色调映射"""# 转换为LAB色彩空间lab = cv2.cvtColor((image * 255).astype(np.uint8), cv2.COLOR_RGB2LAB)L, a, b = cv2.split(lab)# 应用CLAHE到亮度通道clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)L_clahe = clahe.apply(L)# 重构LAB图像并转回RGBlab_clahe = cv2.merge([L_clahe, a, b])rgb_clahe = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB)return rgb_clahe.astype(np.float32) / 255.0
3. 梯度域优化
直接在梯度域操作,保留重要的梯度信息:
- 梯度域滤波:对图像梯度进行滤波,避免传统方法中的光晕效应。
- 梯度保留插值:在压缩动态范围的同时保留关键梯度。
def gradient_domain_tone_mapping(image, sigma=0.5, epsilon=0.01):"""梯度域色调映射"""# 计算亮度通道gray = cv2.cvtColor((image * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0# 计算梯度sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)gradient_mag = np.sqrt(sobelx**2 + sobely**2)# 梯度压缩compressed_gradient = gradient_mag / (1 + sigma * gradient_mag)# 重建图像(简化版,实际需通过泊松方程求解)# 这里使用简化的方法:基于梯度比值调整亮度scale = compressed_gradient / (gradient_mag + epsilon)scaled_gray = gray * scale# 重构彩色图像scaled_rgb = image * np.stack([scaled_gray / (gray + epsilon)]*3, axis=2)scaled_rgb = np.clip(scaled_rgb, 0, 1)return scaled_rgb
4. 基于深度学习的方法
利用神经网络学习最优的色调映射函数:
- CNN模型:端到端学习从HDR到LDR的映射。
- 生成对抗网络(GAN):生成更真实自然的结果。
# 简化的深度学习色调映射模型(实际需训练)
import tensorflow as tfdef build_tone_mapping_model(input_shape):"""构建简单的CNN色调映射模型"""model = tf.keras.Sequential([tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding='same', input_shape=input_shape),tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding='same'),tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding='same'),tf.keras.layers.Conv2D(3, (3,3), activation='sigmoid', padding='same')])return model
5. 混合方法
结合多种技术的优势:
- 全局+局部混合:先用全局方法处理基础亮度,再用局部方法增强细节。
- 自适应参数调整:根据图像内容自动调整局部处理参数。
def hybrid_tone_mapping(image, global_weight=0.6, local_weight=0.4):"""混合全局和局部方法"""# 全局处理global_result = global_tone_mapping(image)# 局部处理local_result = optimized_local_tone_mapping(image)# 加权融合hybrid_result = global_weight * global_result + local_weight * local_resultreturn np.clip(hybrid_result, 0, 1)
6. GPU加速优化
对于实时应用,使用CUDA或OpenCL加速计算密集型操作:
- 并行计算:将滤波、梯度计算等操作并行化。
- 内存优化:减少中间结果的内存占用。
# 使用numba进行GPU加速示例
from numba import cuda@cuda.jit
def parallel_log_transform(src, dst):"""并行对数变换"""i, j = cuda.grid(2)if i < src.shape[0] and j < src.shape[1]:dst[i, j] = np.log1p(src[i, j])
优化方向总结
方法 | 核心优势 | 适用场景 |
---|---|---|
视网膜感知模型 | 符合人眼视觉特性,保色性好 | 人像、风景照片 |
直方图优化 | 增强局部对比度 | 医学影像、监控视频 |
梯度域处理 | 精确控制细节保留 | 高动态范围场景 |
深度学习 | 自适应学习最优映射 | 复杂场景、高质量需求 |
混合方法 | 平衡效率与效果 | 通用场景 |
GPU加速 | 实时处理 | 视频、游戏等实时应用 |
这些方法可以根据具体需求组合使用,例如将深度学习模型与传统算法结合,或在GPU上加速混合方法,以达到最佳的视觉效果和处理效率。
5. 用 PyTorch 实现基于深度学习的色调映射方法
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import torchvision.transforms.functional as F
import os
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt# 定义色调映射模型
class ToneMappingNet(nn.Module):def __init__(self):super(ToneMappingNet, self).__init__()# 编码器部分 - 提取特征self.encoder = nn.Sequential(nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),nn.ReLU(inplace=True),nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),nn.ReLU(inplace=True),nn.MaxPool2d(kernel_size=2, stride=2))# 解码器部分 - 生成LDR图像self.decoder = nn.Sequential(nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),nn.ReLU(inplace=True),nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1),nn.ReLU(inplace=True),nn.Conv2d(32, 3, kernel_size=3, stride=1, padding=1),nn.Sigmoid() # 将输出限制在[0,1]范围内)def forward(self, x):x = self.encoder(x)x = self.decoder(x)return x# 自定义数据集类
class HDRLDataset(Dataset):def __init__(self, hdr_dir, ldr_dir, transform=None):self.hdr_dir = hdr_dirself.ldr_dir = ldr_dirself.transform = transformself.hdr_files = sorted(os.listdir(hdr_dir))self.ldr_files = sorted(os.listdir(ldr_dir))# 确保HDR和LDR图像数量匹配assert len(self.hdr_files) == len(self.ldr_files), "HDR和LDR图像数量不匹配"def __len__(self):return len(self.hdr_files)def __getitem__(self, idx):hdr_path = os.path.join(self.hdr_dir, self.hdr_files[idx])ldr_path = os.path.join(self.ldr_dir, self.ldr_files[idx])# 加载HDR图像 (使用OpenCV或其他库读取HDR格式)hdr_img = cv2.imread(hdr_path, cv2.IMREAD_ANYDEPTH)hdr_img = cv2.cvtColor(hdr_img, cv2.COLOR_BGR2RGB)# 加载LDR图像ldr_img = Image.open(ldr_path).convert('RGB')# 应用变换if self.transform:hdr_img = self.transform(hdr_img)ldr_img = self.transform(ldr_img)# 将HDR图像转换为Tensor并归一化hdr_tensor = torch.tensor(hdr_img, dtype=torch.float32).permute(2, 0, 1)ldr_tensor = transforms.ToTensor()(ldr_img)return hdr_tensor, ldr_tensor# 训练函数
def train_model(model, train_loader, criterion, optimizer, device, epochs=100):model.train()for epoch in range(epochs):running_loss = 0.0for hdr_images, ldr_images in train_loader:hdr_images = hdr_images.to(device)ldr_images = ldr_images.to(device)# 前向传播outputs = model(hdr_images)loss = criterion(outputs, ldr_images)# 反向传播和优化optimizer.zero_grad()loss.backward()optimizer.step()running_loss += loss.item()# 打印训练信息avg_loss = running_loss / len(train_loader)print(f'Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.6f}')return model# 推理函数
def predict(model, hdr_image, device):model.eval()with torch.no_grad():hdr_tensor = torch.tensor(hdr_image, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device)ldr_pred = model(hdr_tensor)ldr_pred = ldr_pred.squeeze(0).cpu().permute(1, 2, 0).numpy()return ldr_pred# 可视化结果
def visualize_results(hdr_image, ldr_pred, ldr_gt=None):plt.figure(figsize=(15, 5))plt.subplot(131)plt.title('HDR Image (Scaled)')# 缩放HDR图像以便显示hdr_scaled = np.clip(hdr_image / np.max(hdr_image), 0, 1)plt.imshow(hdr_scaled)plt.axis('off')plt.subplot(132)plt.title('Predicted LDR')plt.imshow(ldr_pred)plt.axis('off')if ldr_gt is not None:plt.subplot(133)plt.title('Ground Truth LDR')plt.imshow(ldr_gt)plt.axis('off')plt.tight_layout()plt.show()# 主函数
def main():# 设置设备device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# 初始化模型model = ToneMappingNet().to(device)# 定义损失函数和优化器criterion = nn.MSELoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 创建数据集和数据加载器# 注意:需要准备HDR-LDR图像对数据集transform = transforms.Compose([transforms.Resize((256, 256)),transforms.RandomCrop(224),transforms.RandomHorizontalFlip()])# 实际使用时替换为你的数据集路径train_dataset = HDRLDataset(hdr_dir='./data/hdr_train', ldr_dir='./data/ldr_train', transform=transform)train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)# 训练模型model = train_model(model, train_loader, criterion, optimizer, device, epochs=50)# 保存模型torch.save(model.state_dict(), 'tone_mapping_model.pth')# 加载模型进行推理model.load_state_dict(torch.load('tone_mapping_model.pth'))# 示例:对单张图像进行推理# 注意:需要提供HDR图像进行测试test_hdr_image = cv2.imread('test.hdr', cv2.IMREAD_ANYDEPTH)test_hdr_image = cv2.cvtColor(test_hdr_image, cv2.COLOR_BGR2RGB)# 预处理测试图像test_hdr_image = cv2.resize(test_hdr_image, (224, 224))test_hdr_image = test_hdr_image / np.max(test_hdr_image) # 归一化# 预测ldr_prediction = predict(model, test_hdr_image, device)# 可视化结果visualize_results(test_hdr_image, ldr_prediction)if __name__ == "__main__":main()
当然可以!以下是使用PyTorch实现基于深度学习的色调映射方法,包括模型定义、训练流程和推理代码:
这个实现包含以下核心组件:
-
ToneMappingNet模型:
- 使用编码器-解码器架构
- 包含卷积层和上采样层
- 输出范围通过Sigmoid限制在[0,1]
-
HDRLDataset数据集类:
- 处理HDR-LDR图像对
- 支持图像变换和增强
-
训练流程:
- 使用MSE损失函数
- Adam优化器
- 支持多轮训练
-
推理功能:
- 对单张HDR图像进行处理
- 返回预测的LDR图像
-
可视化工具:
- 对比显示HDR原图和生成的LDR图像
使用时需要准备HDR-LDR图像对数据集,其中HDR图像作为输入,对应的LDR图像作为目标输出。模型会学习它们之间的映射关系。
对于更复杂的场景,可以考虑扩展网络结构,例如加入残差块、注意力机制或使用更先进的U-Net架构。此外,还可以尝试使用感知损失(如VGG特征损失)来提高视觉质量。