Underwater Image Enhancement Based on UNDERWATER-CUT Model
-
摘要: 针对水下图像的颜色失真和对比度失衡问题, 提出了一种基于UNDERWATER-CUT模型的弱监督水下图像增强算法。该算法网络训练时无需成对训练集, 通过图像分块的方式构建对比学习正负样本, 约束了图像生成的内容。使用了结构相似性损失函数对水下图像增强进行约束, 确保CUT模型在对水下图像域到脱水图像域进行转换过程中的物体结构不变。同时还使用alignment和uniformity 2个简单函数的线性组合来近似逼近CUT模型的InfoNCE损失函数, 使得改进后的模型训练更容易收敛到最优值。实验结果证明经过文中算法增强后的图像, 颜色失真得到极大的修正, 图像中的物体结构和修复前的图像基本一致。
-
关键词:
- 水下图像增强 /
- 结构相似性 /
- InfoNCE损失函数
Abstract: A weakly supervised underwater image enhancement algorithm based on the UNDERWATER-CUT model was proposed to address color distortion and contrast imbalance in underwater images. The algorithm network was trained without paired training sets, and the contrast learning positive and negative samples were constructed by image chunking, which constrains the image generation content. The structural similarity(SSIM) loss function was used to constrain the image enhancement to ensure that the structure of the objects remains unchanged during the transformation of the CUT model from the underwater image domain to the dewatered image domain. A linear combination of two simple functions was used to approximate the InfoNCE loss of the CUT model, converging the improved model training to the optimal value easily. The experimental results demonstrate that color distortion is significantly corrected in the enhanced image using this algorithm, and the object structure in the restored image is the same as that in the original image. -
表 1 各方法增强后图像指标对比
Table 1. Comparison of the image indicators enhanced by each method
模型 FID UCIQE PSNR CycleGAN 129.63 0.551 9 14.99 SSIM-CycleGAN 107.77 0.557 8 14.85 CUT 182.38 0.551 6 13.48 SSIM-CUT 113.66 0.625 2 18.62 UNDERWATER-CUT 104.53 0.634 2 18.18 -
[1] 杨爱萍, 曲畅, 王建, 等. 基于水下成像模型的图像清晰化算法[J]. 电子与信息学报, 2018, 40(2): 298-305.Yang Ai-ping, Qu Chang, Wang Jian, et al. Underwater Image Visibility Restoration Based on Underwater Imaging Model[J]. Journal of Electronics & Information Technology, 2018, 40(2): 298-305. [2] He K, Sun J, Tang X. Single Image Haze Removal Using Dark Channel Prior[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2010, 33(12): 2341-2353. [3] Wen H, Tian Y, Huang T, et al. Single Underwater Image Enhancement with a New Optical Model[C]//2013 IEEE International Symposium on Circuits and Systems(ISCAS). Beijing, China: IEEE, 2013: 753-756. [4] Li J, Skinner K A, Eustice R M, et al. WaterGAN: Unsupervised Generative Network to Enable Real-time Color Correction of Monocular Underwater Images[J]. IEEE Robotics and Automation Letters, 2017, 3(1): 387-394. [5] Li C, Guo J, Guo C. Emerging from Water: Underwater Image Color Correction Based on Weakly Supervised Color Transfer[J]. IEEE Signal Processing Letters, Institute of Electrical and Electronics Engineers Inc., 2018, 25(3): 323-327. [6] Zhu J Y, Park T, Isola P, et al. Unpaired Image-to-Image Translation Using Cycle-Consistent Adversarial Networks[C]//2017 IEEE International Conference on Computer Vision (ICCV). Venice, Italy: IEEE, 2017. [7] Wang Z, Bovik A C, Sheikh H R, et al. Image Quality Assessment: from Error Visibility to Structural Similarity[J]. IEEE Transactions on Image Processing, 2004, 13(4): 600-612. doi: 10.1109/TIP.2003.819861 [8] 李庆忠, 白文秀, 牛炯. 基于改进CycleGAN的水下图像颜色校正与增强[J/OL]. 自动化学报, 2020, 46(x): 1-11.[2021-10-10]. http://www.aas.net.cn/cn/article/doi/10.16383/j.aas.c200510.Li Qing-Zhong, Bai Wen-Xiu, Niu Jiong. Underwater Image Color Correction and Enhancement Based on Improved Cycle-consistent Generative Adversarial Networks[J/OL]. Acta Automatica Sinica, 2020, 46(x): 1−11. [2021-10-10]. http://www.aas.net.cn/cn/article/doi/10.16383/j.aas.c200510. [9] Park T, Efros A A, Zhang R, et al. Contrastive Learning for Unpaired Image-to-Image Translation[EB/OL]. (2020-08-20)[2021-11-10]. https://arxiv.org/abs/2007.15651. [10] Van den Oord A, Li Y, Vinyals O. Representation Learning with Contrastive Predictive Coding[EB/OL]. (2019-01-22)[2021-11-10]. http://export.arxiv.org/abs/1807.03748. [11] Goodfellow I J, Pouget-Abadie J, Mirza M, et al. Generative Adversarial Networks[EB/OL]. (2017-04-03)[2021-11-10]. https://doi.org/10.48550/arXiv.1701.00160. [12] He K, Zhang X, Ren S, et al. Deep Residual Learning for Image Recognition[C]//2016 IEEE Conference on Computer Vision and Pattern Recognition(CVPR). Las Vegas, NV, USA: IEEE, 2016: 770-778. [13] Isola P, Zhu J-Y, Zhou T, et al. Image-to-Image Translation with Conditional Adversarial Networks[C]//2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). Honolulu, HI, USA: IEEE, 2018. [14] Wang T, Isola P. Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere[EB/OL]. (2020-11-10)[2021-12-17]. https://doi.org/10.48550/arXiv.2005.10242. [15] Li C, Guo C, Ren W, et al. An Underwater Image Enhancement Benchmark Dataset and Beyond[J]. IEEE Transactions on Image Processing, 2020, 29: 4376-4389. doi: 10.1109/TIP.2019.2955241 [16] Paszke A, Gross S, Massa F, et al. PyTorch: An Imperative Style, High-Performance Deep Learning Library[EB/OL]. (2019-12-03)[2021-11-07]. https://doi.org/10.48550/arXiv.1912.01703. [17] Huynh-Thu Q, Ghanbari M. Scope of Validity of PSNR in Image/Video Quality Assessment[J]. Electronics Letters, IET, 2008, 44(13): 800-801. doi: 10.1049/el:20080522 [18] Heusel M, Ramsauer H, Unterthiner T, et al. GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium[EB/OL]. (2018-01-12)[2021-11-07]. https://doi.org/10.48550/arXiv.1706.08500. [19] Yang M, Sowmya A. An Underwater Color Image Quality Evaluation Metric[J]. IEEE Transactions on Image Processing, 2015, 24(12): 6062-6071. doi: 10.1109/TIP.2015.2491020