CE-VAE: Capsule Enhanced Variational AutoEncoderfor Underwater Image Enhancement

Rita Pucci, Niki Martinel; Proceedings of the Winter Conference on Applications of Computer Vision (WACV), 2025, pp. 2113-2123

Unmanned underwater image analysis for marine monitoring faces two key challenges: (i) degraded image quality due to light attenuation and (ii) hardware storage constraints limiting high-resolution image collection. Existing methods primarily address image enhancement with approaches that hinge on storing the full-size input. In contrast we introduce the Capsule Enhanced Variational AutoEncoder (CE-VAE) a novel architecture designed to efficiently compress and enhance degraded underwater images. Our attention-aware image encoder can project the input image onto a latent space representation while being able to run online on a remote device. The only information that needs to be stored on the device or sent to a beacon is a compressed representation. There is a dual-decoder module that performs offline full-size enhanced image generation. One branch reconstructs spatial details from the compressed latent space while the second branch utilizes a capsule-clustering layer to capture entity-level structures and complex spatial relationships. This parallel decoding strategy enables the model to balance fine-detail preservation with context-aware enhancements. CE-VAE achieves state-of-the-art performance in underwater image enhancement on six benchmark datasets providing up to 3x higher compression efficiency than existing approaches. Code available at https://github.com/iN1k1/ce-vae-underwater-image-enhancement.