AquaBalanceLoss: Underwater Image Enhancement

• AquaBalanceLoss is a multi-attribute loss function that integrates color, sharpness, and contrast constraints to achieve balanced underwater image enhancement.
• It computes a joint attribute vector and applies a squared difference between input and output, ensuring Pareto-efficient optimization across multiple quality indices.
• Empirical studies demonstrate that incorporating AquaBalanceLoss in frameworks like JDPNet improves PSNR, SSIM, UIQM, and other metrics without overemphasizing any single attribute.

AquaBalanceLoss is a multi-attribute loss function tailored for underwater image enhancement, designed to enable deep networks to achieve a visually balanced trade-off among color fidelity, image sharpness, and local contrast. It operates by aggregating three orthogonal image quality constraints into a single loss space, penalizing the change between input and output in this domain. Implemented as a squared difference between the joint attribute vector (color, sharpness, contrast) before and after enhancement, AquaBalanceLoss encourages Pareto-efficient optimization—improving a given attribute is only permitted unless it degrades others. It forms a core component in the JDPNet framework for joint degradation processing (Ye et al., 23 Dec 2025).

1. Mathematical Formulation

Let $$I_{in}, I_{out} \in \mathbb{R}^{H \times W \times 3}$$ denote the degraded input and enhanced output images respectively. AquaBalanceLoss decomposes image quality into three atomic indices:

• Color Index ($L_{coi}$): Utilizes oppositional channels—$RG(x,y) = R(x,y) - G(x,y)$ and $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$—and computes $\alpha$-trimmed means $(\mu_a)$ and variances $(s_a)$. These are combined as $l = \sqrt{\mu_a(RG)^2 + \mu_a(YB)^2}$, $r = \sqrt{s_a(RG)^2 + s_a(YB)^2}$, giving $L_{coi} = -0.027 \cdot l + 0.159 \cdot r$.
• Sharpness Index ($L_{coi}$0): Edge maps $L_{coi}$1 are obtained via Sobel operators on each channel, then weighted and summed: $L_{coi}$2.
• Contrast Index ($L_{coi}$3): The image is partitioned into non-overlapping blocks of $L_{coi}$4 pixels. For each block, $L_{coi}$5, $L_{coi}$6, contributing $L_{coi}$7 with $L_{coi}$8, $L_{coi}$9.

The joint attribute loss (AbL) aggregates these metrics:

$RG(x,y) = R(x,y) - G(x,y)$0

where $RG(x,y) = R(x,y) - G(x,y)$1.

AquaBalanceLoss itself is given as:

$RG(x,y) = R(x,y) - G(x,y)$2

with bias term $RG(x,y) = R(x,y) - G(x,y)$3 typically zero.

Composite training objectives in JDPNet integrate AquaBalanceLoss as follows:

$RG(x,y) = R(x,y) - G(x,y)$4

Default weights are $RG(x,y) = R(x,y) - G(x,y)$5, $RG(x,y) = R(x,y) - G(x,y)$6, $RG(x,y) = R(x,y) - G(x,y)$7, $RG(x,y) = R(x,y) - G(x,y)$8.

2. Theoretical Motivation and Principles

AquaBalanceLoss is structured to tackle the compounded nonlinear degradations found in underwater imagery, where color casts, blur, and low contrast often coexist and interact. The color component leverages opponent channel statistics with outlier-resistant $RG(x,y) = R(x,y) - G(x,y)$9-trimming, sharpening the response to persistent color shifts without sensitivity to specular outliers. The sharpness measure is Sobel-based, focusing gradients mainly on edges to maintain detail. The local contrast term is blockwise, entropy-sensitive, and parameterized for entropy scaling, ensuring no region is underrepresented in contrast optimization.

Linear combination of color, clarity, and contrast indices ensures that optimization proceeds along nearly orthogonal gradients in parameter space. This property facilitates learning trajectories that maintain balance among these attributes. The loss function’s construction enforces a Pareto front: improvement in one visual attribute is not at the expense of the others, preventing the network from overfitting color, over-sharpening, or over-contrasting single aspects.

3. Weight Selection and Hyperparameter Strategy

Parameter selection follows coarse grid search and cross-validation:

• Attribute weights $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$0 are fixed globally at $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$1 to balance color, sharpness, and contrast across validation sets.
• Loss term weights $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$2 are static, with best practical results at $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$3. Sweeps of $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$4 in the range $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$5–$YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$6 indicate performance peaks for PSNR/SSIM/UIQM/UCIQE around $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$7.
• No dynamic adjustment or learning of weights; all hyperparameters are chosen by cross-validation and held constant.
• $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$8-trim for color statistics is set to 5–10% to remove specular outliers, and contrast block size $YB(x,y) = \frac{1}{2}[R(x,y) + G(x,y)] - B(x,y)$9 and exponent $\alpha$0 (default $\alpha$1, $\alpha$2) are balanced to avoid artifact introduction.

If bespoke application demands weighted emphasis—such as increased sharpness—trade-offs along the Pareto boundary can be achieved by adjusting $\alpha$3 relative to $\alpha$4 and $\alpha$5.

4. Implementation Protocols

Integration involves appending AquaBalanceLoss as the fourth term in the overall training loss, jointly backpropagated with reconstruction ($\alpha$6), perceptual (VGG16 feature), and KL divergence terms. Training details include:

• Optimizer: Adam, initial learning rate $\alpha$7, reduced to $\alpha$8 after 50 epochs.
• Epochs: 300 (overtraining beyond $\alpha$9400–500 epochs may cause color/contrast overfitting).
• Batch size: 1
• Data augmentation: $(\mu_a)$0 crops, random horizontal flips, and random rotations (90°, 180°, 270°).
• Regularization: Dropout in Joint Feature Mining, no explicit weight decay.
• Normalization: images scaled to [0, 1] or channel-wise mean–std normalization.
• AquaBalanceLoss is evaluated on each batch, with AbL computed for both input and output, and the squared difference accumulated with weight $(\mu_a)$1.

5. Empirical Results, Ablations, and Generalization

Ablation studies demonstrate AquaBalanceLoss’s centrality: its removal (“w/o AquaBalanceLoss”) causes measurable declines in PSNR (23.007→22.66), SSIM (0.915→0.898), UIQM (3.045→2.958), PCQI (0.979→0.940), and UCIQE/CCF metrics. Visual outcomes are marked by desaturated color, diminished local contrast, and blurred edges.

Weight sweeps confirm that performance peaks at $(\mu_a)$2, with degradation for larger values due to over-contrasting. Addition of AquaBalanceLoss to different architectures (PUIE-MC, Spectroformer, WaterMamba) consistently improves color, sharpness, and contrast metrics, underscoring architectural agnosticism of its supervision signal.

Comparisons with other losses (reconstruction, perceptual, KL) indicate that AquaBalanceLoss’s removal most severely impacts the combined color–contrast–clarity index (CCF) and perceptual uniformity, highlighting its unique regulative role.

6. Limitations and Trade-Offs

AquaBalanceLoss exhibits several intrinsic trade-offs:

• High $(\mu_a)$3 can overemphasize contrast and color, leading to unnatural artifacts and reduced structural fidelity (lower PSNR/SSIM).
• Prolonged training drives bias toward attribute loss, risking overfit and visual instability.
• Choices of block granularity and entropy scaling (contrast term) raise risk of halo artifacts with inappropriate settings.
• Static weights may be suboptimal for specific domains or adaptation scenarios; no dynamic or learned weighting is employed.
• Optimization is subject to the Pareto front enforced by AbL aggregation: only multi-objective improvements are permitted, so attribute-specific preferences should be reflected in weight adjustment.

7. Practical Significance and Extension Guidance

AquaBalanceLoss provides a rigorous framework for guiding underwater image enhancement networks toward visually satisfying outputs that are perceptually balanced, not skewed by overcorrection of a single attribute. Its methodology is fully specified and directly implementable. The attribute space construction and aggregation principles may plausibly generalize to other scenarios with nonlinearly coupled degradations, provided suitable choices for atomic indices and weight calibration. Extensions should consider the interaction of attribute indices, potential for weight learning, and the effect of loss-space Pareto optimization on subjective image quality and task performance (Ye et al., 23 Dec 2025).