Shallow water UAV imagery restored using diffusion and synthetic data

KAUSTics-Net: Diffusion-Based Through-Water UAV Restoration with a Semi-Synthetic Dataset

Computer Vision and Pattern Recognition

Summary

Taking pictures of shallow water from drones can be tricky because the water bends light and causes distortions that make the images unclear. The authors created a special dataset by combining real sharp images with computer-generated water effects to teach a new computer program how to fix these distortions. They designed a method called KAUSTics-Net that improves clarity by learning how water changes images and then reversing those changes. Their approach worked better than other methods on both their dataset and an independent one, while also helping keep important features useful for mapping underwater environments.

What this means in practice

  • For environmental monitoring teams: Improve clarity of drone images for mapping shallow-water habitats affected by water distortions using KAUSTics-Net restoration.
  • For marine survey operators: Enhance shallow water aerial imagery to better identify underwater features for resource assessment using the proposed diffusion-based restoration.

Authors

Dario Scilla, Omar A. Lopez, Kasper Johansen, Brian O. Nieuwenhuis, Matthew F. McCabe

Abstract

Accurate mapping of shallow-water environments using unmanned aerial vehicle (UAV) imagery is challenged by water-induced distortions, including refraction, caustics, and wavelength-dependent attenuation. The scarcity of paired distorted and reference imagery further limits the development of supervised restoration methods. To address this challenge, we introduce KAUSTics, a semi-synthetic dataset combining 1,280 reference patches derived from 32 high-resolution UAV images with 12,800 degraded observations generated through Blender-based water-surface rendering and empirical color and blur augmentations. We further propose KAUSTics-Net, a pixel-space conditional diffusion framework integrating a multi-scale conditional encoder, feature-wise linear modulation (FiLM), cross-attention, and a composite reconstruction objective to restore shallow-water imagery. We evaluated the framework against eight image-restoration methods, complemented by architectural ablations and cross-dataset experiments. On KAUSTics, KAUSTics-Net achieved a mean peak signal-to-noise ratio (PSNR) of 31.715 +/- 0.439 dB, structural similarity index (SSIM) of 0.849 +/- 0.002, and learned perceptual image patch similarity (LPIPS) of 0.145 +/- 0.007 across three training seeds. On the independent Sea-Undistort dataset, sampler-aware fine-tuning improved restoration fidelity, reaching 19.469 +/- 0.332 dB PSNR and 0.554 +/- 0.005 SSIM under combined water-surface distortions, although perceptual trade-offs remained. Additional experiments on real UAV imagery examined feature preservation and segmentation behavior. These findings demonstrate the potential of semi-synthetic training and conditional diffusion for shallow-water image restoration while highlighting the importance of evaluating both reconstruction fidelity and downstream remote-sensing utility.