使用 webgpu 实现图像暗通道去雾

因为自己是做设计的经常要处理各种图片,最经常搞的就是给渲染图跑一遍去雾来提升一下效果之类.

泼辣修图的效果其实挺好的而且快但是经常有需要批量处理一下或者想要直接在笔记里面处理去雾的情况,所以还是自己弄一下算了.

效果

个人喜欢暴露更多参数,所以界面长这个样子:

在思源里面的界面害没有写完,因为不能把上面那个直接塞到思源里面所以还在想来着:

暗通道先验理论基础

暗通道先验是基于对大量户外无雾图像的观察得出的统计规律：在大多数非天空区域的局部像素块中，至少有一个颜色通道(RGB)的值非常低。这个观察可以表示为：

J^dark(x) = min_{c∈{r,g,b}} (min_{y∈Ω(x)} J^c(y)) → 0

其中 J^c 表示图像的某个颜色通道，Ω(x)表示以像素 x 为中心的局部区域。

对于有雾图像 I，其大气散射模型为：

I(x) = J(x)t(x) + A(1-t(x))

其中：

• J(x)是无雾图像
• t(x)是透射率
• A 是大气光值

我数学比较烂就不说了,反正就是根据雾天图像的共同特征将去雾转化为一个蒙版 + 处理的过程.

WebGPU 架构设计

核心计算流程

整个去雾算法的核心流程在dehazing-webgpu.js中实现：

1. 暗通道计算 - 使用 GPU 并行计算每个像素邻域内的最小值
2. 大气光估计 - 基于暗通道结果估计全局大气光
3. 透射率估计 - 根据暗通道和大气光计算透射率图
4. 图像恢复 - 利用透射率图和大气光恢复无雾图像

关键 WebGPU 函数解析

1. 暗通道计算着色器

暗通道计算是去雾算法的第一步，在shaders.js中定义了对应的计算着色器：

@compute @workgroup_size(16, 8)
fn computeDarkChannel(@builtin(global_invocation_id) globalId: vec3<u32>) {
  let x = globalId.x;
  let y = globalId.y;
  
  if (x >= u32(uniforms.width) || y >= u32(uniforms.height)) {
    return;
  }
  
  let halfWindow = u32(uniforms.windowSize / 2.0);
  var minValue: f32 = 1.0;
  
  // 计算实际有效的窗口范围，避免采样到图像边界之外
  let startX = max(halfWindow, x);
  let endX = min(u32(uniforms.width) - 1 - halfWindow, x);
  let startY = max(halfWindow, y);
  let endY = min(u32(uniforms.height) - 1 - halfWindow, y);
  
  // 边缘区域使用较小的窗口
  var actualHalfWindow = halfWindow;
  if (x < halfWindow || x >= u32(uniforms.width) - halfWindow ||
      y < halfWindow || y >= u32(uniforms.height) - halfWindow) {
    actualHalfWindow = min(halfWindow, min(x, min(u32(uniforms.width) - 1 - x,
                                           min(y, u32(uniforms.height) - 1 - y))));
  }
  
  // 在窗口内寻找最小值
  for (var wy = y - actualHalfWindow; wy <= y + actualHalfWindow; wy++) {
    for (var wx = x - actualHalfWindow; wx <= x + actualHalfWindow; wx++) {
      if (wx < u32(uniforms.width) && wy < u32(uniforms.height)) {
        let pixel = textureLoad(inputTexture, vec2<i32>(i32(wx), i32(wy)), 0);
        
        // 使用人眼亮度权重计算暗通道
        let luminance = pixel.r * 0.299 + pixel.g * 0.587 + pixel.b * 0.114;
        let minChannel = min(min(pixel.r, pixel.g), pixel.b);
        
        // 结合亮度信息和最小通道值
        let darkChannel = min(luminance, minChannel);
        minValue = min(minValue, darkChannel);
      }
    }
  }
  
  let outputIndex = y * u32(uniforms.width) + x;
  outputBuffer[outputIndex] = minValue;
}

关键点：

• 图片边界的地方要记得特殊处理，避免采样到图片外面
• 根据人眼的视觉特性修正估算反正我自己是觉得效果是要好一点的

2. 透射率估计着色器

透射率估计基于暗通道先验，在shaders.js中实现了基础版本和空间自适应版本：

@compute @workgroup_size(16, 8)
fn estimateTransmission(@builtin(global_invocation_id) globalId: vec3<u32>) {
  let x = globalId.x;
  let y = globalId.y;
  
  if (x >= u32(uniforms.width) || y >= u32(uniforms.height)) {
    return;
  }
  
  let pixelIndex = y * u32(uniforms.width) + x;
  let darkValue = darkChannelBuffer[pixelIndex];
  let transmission = 1.0 - uniforms.omega * (darkValue / uniforms.atmosphericLight);
  transmissionBuffer[pixelIndex] = transmission;
}

空间自适应版本则根据每个位置的雾浓度动态调整参数：

fn computeSpatialAdaptiveParameters(darkValue: f32) -> vec2<f32> {
  // 基于暗通道值计算雾强度因子
  let hazeFactor = min(1.0, darkValue * 1.5);
  
  // 基于大气光调整因子
  let atmosphericFactor = min(1.0, uniforms.atmosphericLight * 1.2);
  
  // 基于暗通道值计算对比度因子
  let contrastFactor = min(1.0, darkValue * 1.2);
  
  // 综合计算自适应因子
  let adaptiveFactor = 
    uniforms.hazeWeight * hazeFactor +
    uniforms.atmosphericWeight * atmosphericFactor +
    (1.0 - uniforms.hazeWeight - uniforms.atmosphericWeight) * contrastFactor;
  
  // 基于用户设置的omega进行自适应调整
  let omegaAdjustment = uniforms.omegaAdjustRange * adjustedAdaptiveFactor;
  let adaptiveOmega = uniforms.userOmega + omegaAdjustment;
  
  // 基于用户设置的t0进行自适应调整
  let t0Adjustment = uniforms.t0AdjustRange * adjustedAdaptiveFactor;
  let adaptiveT0 = uniforms.userT0 + t0Adjustment;
  
  return vec2<f32>(
    clamp(adaptiveOmega, 0.1, 0.99),
    clamp(adaptiveT0, 0.01, 0.3)
  );
}

3. 图像恢复着色器

图像恢复是最后一步，在shaders.js中实现了基础版本和增强版本：

@compute @workgroup_size(16, 8)
fn recoverImage(@builtin(global_invocation_id) globalId: vec3<u32>) {
  let x = globalId.x;
  let y = globalId.y;
  
  if (x >= u32(uniforms.width) || y >= u32(uniforms.height)) {
    return;
  }
  
  let pos = vec2<i32>(i32(x), i32(y));
  let pixel = textureLoad(inputTexture, pos, 0);
  
  let r = pixel.r;
  let g = pixel.g;
  let b = pixel.b;
  let a = pixel.a;
  
  let transmission = max(transmissionBuffer[y * u32(uniforms.width) + x], uniforms.t0);
  
  // 计算原始亮度
  let originalLuminance = r * 0.299 + g * 0.587 + b * 0.114;
  
  // 基于亮度进行恢复
  let recoveredLuminance = (originalLuminance - uniforms.atmosphericLightLuminance) / transmission + uniforms.atmosphericLightLuminance;
  
  // 计算亮度变化比例
  let luminanceRatio = recoveredLuminance / max(originalLuminance, 1e-6);
  
  // 分别恢复RGB通道，保持色彩比例
  let recoveredR = (r - uniforms.atmosphericLightR) / transmission + uniforms.atmosphericLightR;
  let recoveredG = (g - uniforms.atmosphericLightG) / transmission + uniforms.atmosphericLightG;
  let recoveredB = (b - uniforms.atmosphericLightB) / transmission + uniforms.atmosphericLightB;
  
  // 结合亮度恢复和色彩恢复，确保自然效果
  let finalR = mix(recoveredR, r * luminanceRatio, 0.7);
  let finalG = mix(recoveredG, g * luminanceRatio, 0.7);
  let finalB = mix(recoveredB, b * luminanceRatio, 0.7);
  
  let outputIndex = (y * u32(uniforms.width) + x) * 4u;
  outputBuffer[outputIndex] = clamp(finalR, 0.0, 1.0);
  outputBuffer[outputIndex + 1u] = clamp(finalG, 0.0, 1.0);
  outputBuffer[outputIndex + 2u] = clamp(finalB, 0.0, 1.0);
  outputBuffer[outputIndex + 3u] = a;
}

增强版本还支持饱和度、对比度和明度调整：

// 饱和度增强函数
fn enhanceSaturation(r: f32, g: f32, b: f32, factor: f32) -> vec3<f32> {
  let hsv = rgbToHsv(r, g, b);
  let enhancedS = clamp(hsv.y * factor, 0.0, 1.0);
  return hsvToRgb(hsv.x, enhancedS, hsv.z);
}

// 对比度增强函数
fn enhanceContrast(r: f32, g: f32, b: f32, factor: f32) -> vec3<f32> {
  return vec3<f32>(
    clamp((r - 0.5) * factor + 0.5, 0.0, 1.0),
    clamp((g - 0.5) * factor + 0.5, 0.0, 1.0),
    clamp((b - 0.5) * factor + 0.5, 0.0, 1.0)
  );
}

这个跟直接调节对比那些的区别就是多了暗通道信息,可以针对性的做一下恢复.

性能优化策略

其实主要就是做了一下缓存.

代码

这个主要是用来开发我自己用的分支版本思源,单独仓库位于 https://github.com/leolee9086/dehaze
之后的更新主要会在 https://github.com/leolee9086/siyuan/tree/multipleAI/packages 进行