CE-VAE: Capsule Enhanced Variational AutoEncoder for Underwater Image Enhancement

TL;DR

CE-VAE is a novel architecture that efficiently compresses and enhances underwater images using a capsule-enhanced variational autoencoder, enabling high-quality image reconstruction with reduced storage requirements.

Contribution

Introduces CE-VAE, a new model combining capsule networks and variational autoencoders for efficient underwater image compression and enhancement, suitable for remote device deployment.

Findings

Achieves state-of-the-art enhancement performance on six datasets.

Provides up to 3x higher compression efficiency than existing methods.

Enables online processing on remote devices with minimal storage.

Abstract

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 \url{https://github.com/iN1k1/ce-vae-underwater-image-enhancement}.