zero-shot-dehaze 0.1.0

A Zero-Shot Underwater Image Dehazing Library

An Enhanced Multi-Stage Approach for Dehazing Underwater Images

This repository contains a Python implementation of the image enhancement pipeline described in the 2024 paper, "An Enhanced Multi-Stage Approach for Dehazing Underwater Images" by Murugan et al. The project faithfully implements the three-stage process outlined in the paper to correct for color distortion, low contrast, and haze in underwater photographs.

---

Disclaimer: This project is a re-implementation of the research work described in the referenced papers. The author of this repository is not affiliated with the original authors of the papers. This implementation represents a best-effort attempt to accurately reproduce the described methodology, but it is not guaranteed to be an exact or fully accurate replication of the authors' original work or results.

Features

• Multi-Stage Enhancement Pipeline: Combines pre-processing, deep learning-based dehazing, and post-processing color correction for robust underwater image restoration.
• Zero-Shot Learning: Utilizes an unsupervised optimization approach that requires no training data, adapting to each specific image at runtime.
• Physical Model Inversion: Solves the underwater image formation model to estimate scene radiance, transmission maps, and atmospheric light.
• Advanced Network Architectures: Implements specialized U-Net and CNN architectures (J-Net, T-Net, A-Net) designed for underwater scene estimation.
• Comprehensive Post-Processing: Includes automated white balancing, adaptive saturation adjustment, and gamma correction to restore visual quality.
• Flexible CLI & API: Provides both a command-line interface for batch processing and a Python API for integration into other workflows.
• Quantitative Metrics: Built-in support for calculating PSNR and SSIM quality metrics against reference images.
• GPU Acceleration: Fully supports CUDA and MPS (Mac) acceleration for faster inference.

Methodology and Architecture

The enhancement process implements a three-stage workflow inspired by "An Enhanced Multi-Stage Approach for Dehazing Underwater Images" (Murugan et al., 2024). Network architectures are based on the Zero-UMSIE framework referenced in the original work.

Stage 1: Pre-processing

This initial stage performs ONLY contrast enhancement to prepare the image for dehazing:

• Contrast Limited Adaptive Histogram Equalization (CLAHE): Applied in LAB color space to enhance contrast while preserving color information

Important: Only one contrast enhancement method (histogram equalization or CLAHE) is applied. Other color corrections (white balance, gamma correction, saturation adjustment) are deferred to the post-processing stage, AFTER dehazing.

Stage 2: Zero-Shot Image Dehazing (ZID)

The core of the pipeline is the ZID model, which decomposes each underwater image into three fundamental components using the underwater imaging model:

I(x) = J(x)·t(x) + A·(1 - t(x))

Where:

• I(x) = Observed underwater image
• J(x) = Scene radiance (clean image)
• t(x) = Transmission map (3-channel, wavelength-dependent)
• A = Global background light (1-channel)

This is achieved by jointly optimizing three neural networks (J-Net, T-Net, and A-Net) using a composite loss function. The model is "zero-shot" because it optimizes on each input image individually without requiring pre-training on paired datasets.

Network Architectures (Based on Zero-UMSIE)

The network structures are detailed in the Zero-UMSIE framework:

• J-Net (Scene Radiance Estimation Network):

  • 5 convolutional layers (3×3 kernels, 64 channels each)
  • 4 instance normalization layers
  • 4 ReLU activation layers
  • Sigmoid output layer (normalizes to [0,1])
  • No downsampling to preserve fine details
  • Output: 3-channel RGB scene radiance

• T-Net (Transmission Map Estimation Network):

  • Similar architecture to J-Net
  • Output: 3-channel RGB transmission map
  • Accounts for wavelength-selective attenuation underwater

• A-Net (Global Background Light Estimation Network):

  • 6 convolutional layers
  • 5 LeakyReLU activation layers
  • Sigmoid output layer
  • Based on Retinex decomposition concept
  • Output: 1-channel grayscale global background light

Loss Functions

The loss function enforces physical and statistical priors about underwater image formation:

• Reconstruction Loss (L_rec): MSE loss ensuring the dehazing equation holds
• Atmospheric Light Loss (L_atm): Variational inference for global light estimation
• Dark Channel Prior Loss (L_dark): Statistics-based loss for scene radiance
• Laplacian Regularization Loss (L_reg): Smoothness prior for transmission and atmospheric light

Stage 3: Post-processing

After dehazing, this final stage applies color corrections to produce the final enhanced image. The techniques include:

• White Balance Correction: Neutralizes color casts by adjusting RGB channel means
• Color Enhancement: Saturation adjustments to restore vibrant underwater colors
• Gamma Correction: Final brightness/contrast adjustments
• CLAHE (Optional): Additional contrast enhancement if needed

Note: These color corrections are intentionally applied in this stage AFTER dehazing, not during preprocessing.

Installation

There are two primary ways to install and use this library:

1. Install from PyPI (Recommended for Users)

Once the package is published to PyPI, you can install it directly using pip:

pip install zero-shot-dehaze

2. Install from Source (For Developers)

If you intend to contribute to the project, make modifications, or run the examples from the repository, you can install from source:

1. Clone the repository:

   git clone https://github.com/jasmitte/Underwater_Image_Enhancement.git
   cd Underwater_Image_Enhancement
   

2. Create and activate a Python virtual environment (recommended):

   python3 -m venv .venv
   source .venv/bin/activate
   

3. Install the package in editable mode:

   pip install -e .
   

Usage

Command-Line Interface (CLI)

The easiest way to use the model is through the provided CLI.

To enhance a single image:

python -m zero_shot_dehaze.cli --input path/to/your/image.jpg --output outputs/

To batch-process a directory of images:

python -m zero_shot_dehaze.cli --input path/to/image/directory --output outputs/

To evaluate with reference images:

If you have ground-truth (clean) images, you can calculate PSNR and SSIM metrics by providing a reference path.

python -m zero_shot_dehaze.cli --input /path/to/raw_images --output outputs/ --reference /path/to/reference_images

Key CLI Arguments

• --input: Path to the input image or directory.
• --output: Path to the output directory.
• --reference: (Optional) Path to the ground-truth image or directory for metric calculation.
• --max-iterations: Optimization steps for the ZID solver (default: 600). Reduce for faster previews.
• --device: Force a specific device (e.g., cpu, cuda, cuda:0).
• --no-clahe: Disable CLAHE in both pre and post-processing stages for a different aesthetic.

Library Usage

You can also import and use the enhancement pipeline directly in your Python code.

from pathlib import Path
from zero_shot_dehaze.pipeline import MultistageUnderwaterEnhancer, EnhancerConfig
from zero_shot_dehaze.zid import ZIDConfig
from zero_shot_dehaze.utils import load_image, save_image

# Configure the pipeline
config = EnhancerConfig()
config.zid.max_iterations = 400
config.zid.device = "cuda"  # or "cpu"

# Initialize the enhancer
enhancer = MultistageUnderwaterEnhancer(config)

# Load an image and enhance it
input_image_path = Path("input.jpg")
image_array = load_image(input_image_path)
enhanced_array = enhancer.enhance_array(image_array)

# Save the result
output_path = Path("outputs/enhanced_image.png")
output_path.parent.mkdir(exist_ok=True)
save_image(enhanced_array, output_path)

print(f"Enhanced image saved to {output_path}")

Running Tests

A simple test suite is included to verify that the end-to-end pipeline executes correctly.

python -m pytest

Citation

This project is an implementation of the work described in the following papers. Please consider citing the original authors if you use this code in your research.

Primary Paper (Methodology):

T. K. Murugan, S. Sharma, A. Ganguly, A. Banerjee, and K. Kejriwal, "An Enhanced Multi-Stage Approach for Dehazing Underwater Images," IEEE Access, vol. 12, pp. 156803-156822, 2024. DOI: 10.1109/ACCESS.2024.3486456

Network Architectures (J-Net, T-Net, A-Net):

T. Liu, K. Zhu, W. Cao, B. Shan, and F. Guo, "Zero-UMSIE: a zero-shot underwater multi-scale image enhancement method based on isomorphic features," Optics Express, vol. 32, no. 23, pp. 40398-40415, 2024. DOI: 10.1364/OE.538120