【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計算功能：

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()