Physical-Guided Lightweight Underwater Image Enhancement

这是一个面向个人水下图像数据的 PyTorch 工程第一版，重点是“可运行、可训练、可推理、可解释”。项目不封装第三方增强模型，而是实现了两个可训练网络：

• baseline: 纯深度学习轻量 encoder-decoder，用作对照组。
• physics_guided: 估计背景光 A 和传输图 t(x)，先按物理模型粗恢复，再用轻量网络 refinement。

当前版本默认按无配对真实水下图像训练，也保留了配对监督训练接口。

方法思路

水下退化可以近似写成：

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

主模型先学习：

• A_estimator: 估计全局/低频背景光 A
• transmission_estimator: 估计传输图 t(x)
• coarse_restore: J_coarse = (I - A) / max(t, eps) + A
• refinement_net: 输入原图、粗恢复图和传输图，输出最终增强图

总损失为：

L = w_color L_color
  + w_exp L_exposure
  + w_struct L_gradient
  + w_edge L_edge
  + w_tv L_tv
  + w_phy L_physics
  + w_sup L_supervised

其中监督项只有在 data.target_dir 提供配对图像，并在 YAML 中打开对应权重时生效。

目录结构

underwater_enhancement/
├─ configs/
│  ├─ baseline.yaml
│  └─ physics_guided.yaml
├─ data/
│  └─ raw/images/
├─ datasets/
│  └─ underwater_dataset.py
├─ losses/
│  ├─ color_loss.py
│  ├─ structure_loss.py
│  ├─ physics_loss.py
│  └─ total_loss.py
├─ models/
│  ├─ baseline_net.py
│  ├─ physics_guided_net.py
│  └─ modules.py
├─ scripts/
│  ├─ scan_dataset.py
│  ├─ plot_rgb_hist.py
│  ├─ train.py
│  ├─ infer.py
│  └─ compare_methods.py
├─ utils/
│  ├─ metrics.py
│  ├─ image_ops.py
│  ├─ logger.py
│  └─ vis.py
├─ results/
├─ checkpoints/
├─ README.md
├─ requirements.txt
└─ main.py

旧版 src/ 和 run_baselines.py 保留，用于传统算法对照，不影响新 PyTorch 工程。

安装依赖

建议使用虚拟环境：

cd D:\study\02_science_product\underwater_enhancement
pip install -r requirements.txt

如果你使用 CUDA，请按本机 CUDA 版本从 PyTorch 官网安装对应的 torch。

数据准备

默认配置读取：

data/raw/images

可以直接把 JPG/PNG/JPEG 图像放进去，例如：

data/raw/images/G0273464.JPG

如果你的数据在其他目录，修改：

data:
  input_dir: D:/your/path/images

若后续有配对清晰图，把目标图放入单独目录，并保证文件名一致，然后设置：

data:
  target_dir: D:/your/path/targets
loss:
  weights:
    supervised_l1: 1.0
    supervised_charbonnier: 0.5
    supervised_ssim: 0.2

扫描数据集

python scripts/scan_dataset.py --input_dir data/raw/images

或：

python main.py scan --input_dir data/raw/images

输出图像数量、文件名示例和尺寸分布，适合先检查高分辨率数据是否被正确读取。

绘制 RGB 直方图

python scripts/plot_rgb_hist.py --input_dir data/raw/images -n 30

结果保存到：

results/histograms/rgb_histogram.png

可用于观察数据是否整体偏蓝、偏绿或红通道衰减明显。

训练 baseline

python scripts/train.py --config configs/baseline.yaml

输出：

• checkpoints/baseline/latest.pth
• checkpoints/baseline/best.pth
• results/baseline/loss_curve.png
• results/baseline/samples/epoch_xxx_comparison.jpg

断点续训：

python scripts/train.py --config configs/baseline.yaml --resume checkpoints/baseline/latest.pth

训练 physics-guided 模型

python scripts/train.py --config configs/physics_guided.yaml

输出除增强样例外，还会保存：

• epoch_xxx_coarse.jpg
• epoch_xxx_transmission.png
• epoch_xxx_A.png

这些图可以展示物理先验分支的可解释性。

推理单张图

python scripts/infer.py ^
  --config configs/physics_guided.yaml ^
  --checkpoint checkpoints/physics_guided/best.pth ^
  --input data/raw/images/G0273464.JPG ^
  --output_dir results/inference_physics ^
  --tile_size 1024 ^
  --overlap 96 ^
  --save_metrics

输出：

• *_enhanced.jpg
• *_compare.jpg
• *_coarse.jpg
• *_transmission.png
• *_A.png
• metrics.csv，若开启 --save_metrics

tile_size 用于大图分块推理，可避免一次性把 5568×4872 图像放进显存。

推理整个文件夹

python scripts/infer.py ^
  --config configs/physics_guided.yaml ^
  --checkpoint checkpoints/physics_guided/best.pth ^
  --input data/raw/images ^
  --output_dir results/inference_physics

baseline 推理只需换配置和权重：

python scripts/infer.py ^
  --config configs/baseline.yaml ^
  --checkpoint checkpoints/baseline/best.pth ^
  --input data/raw/images ^
  --output_dir results/inference_baseline

对比 baseline 与 physics-guided

python scripts/compare_methods.py ^
  --input_dir data/raw/images ^
  --baseline_dir results/inference_baseline ^
  --physics_dir results/inference_physics ^
  --output_dir results/comparisons

当前版本局限性

• 无配对训练的目标是稳妥增强，不保证达到公开数据集 SOTA。
• UIQM 当前提供轻量 proxy，UCIQE、信息熵、RMS 对比度、颜色均衡误差已可用。
• 背景光 A 目前是全局 RGB 估计，后续可扩展为空间低频图。
• 感知损失接口预留在 supervised loss 位置，第一版没有强依赖 VGG，以保持安装和训练简单。

后续扩展方向

• 加入全量 UIQM 和更多无参考水下质量指标。
• 引入更稳健的红通道补偿先验，约束 refinement 不过度增红。
• 加入 paired/unpaired 混合训练策略。
• 支持 ONNX 导出和简历展示用 Web demo。
• 引入 EMA、学习率调度和更完整的实验记录。