Skip to main content

Hazy Image Restoration Using Keras

An introduction to building an Image Restoration model using Tensorflow, Keras, and Weights & Biases.
Created on October 15|Last edited on February 16

Who am I?



Introduction

The existence of haze in the atmosphere–due to the presence of aerosols such as dust, mist, and smoke–adds complicated noise to photographs. It dramatically degrades the visibility of outdoor images (and sometimes even indoor images), where contrasts are reduced and surface colors become faint.
From an ML perspective, a hazy image jeopardizes the effectiveness of many subsequent high-level computer vision tasks such as object detection, recognition, and segmentation.

Predictions
Number-of-Objects
Is-Hazy
2
4
Predictions
Number-of-Objects
Is-Hazy
1
3
The presence of atmospheric haze significatlt impacts the performance of State-of-the-Art Object Detection systems such as YOLOv6. This makes the restoration of images an extremely crucial aspect of any computer vision pipeline.
4

It follows then that problem of image-dehazing is a widely explored problem in computer vision. In this report, we'll be exploring a fast and simple end-to-end method for image dehazing proposed by the paper An All-in-One Network for Dehazing and Beyond.
The paper, albeit quite old (in terms of the blazing fast pace in which ML research moves), was the first paper to propose an end-to-end trainable image dehazing model, that can be trained to directly produce a clean image from a hazy image. The paper formulates a simple neural network architecture called AODNet (All-in-One De-hazing Network) which we train using Tensorflow and Keras. We'll also share and visualize the results of our experiments using Weights & Biases.
This article was written as a Weights & Biases Report which is a project management and collaboration tool for machine learning projects. Reports let you organize and embed visualizations, describe your findings, share updates with collaborators, and more. To know more about reports, check out Collaborative Reports.
💡

What We'll Be Covering



Looking at the Dataset

We'll train the AODNet model primarily on synthetic hazy images, and test on both synthetic and real natural images from multiple datasets. We list the datasets and Weave panels to visualize and explore these datasets below:

The Image-dehazing Dataset
298




The Data Input Pipeline

For training our model, we build a simple Tensorflow-based data input pipeline using the tf.data API. The tf.data API makes it possible to handle large amounts of data, read from different data formats, and perform complex transformations.

The Data Input Pipeline
0




The All-in-One Dehazing Network

Now that we have the data input pipeline in place, let's take a look at the concept of the AODNet architecture in detail.

The Mathematical Model for Hazy Images

The classical description for the hazy image generation (and conversely image-dehazing too) has been described by the atmospheric scattering model. According to this model, a hazy image is given by...
I(x)=J(x)t(x)+A(1t(x))\Large{I(x)=J(x) t(x)+A(1-t(x))}

...where:
  • I(x)I(x) is the observed hazy image
  • J(x)J(x) is the idea clean image
  • AA denotes the global atmospheric light
  • t(x)t(x) is the transmission matrix
The image dehazing problem can be formulated by expressing the atmospheric scattering model in terms of the hazy image:
J(x)=1t(x)I(x)A1t(x)+A\Large{J(x)=\frac{1}{t(x)} I(x)-A \frac{1}{t(x)}+A}

A lot of prior approaches in image dehazing such as DehazeNet and Multi-scale CNNs follow the identical three-step procedure:
  1. Estimating the transmission matrix t(x)t(x) from the hazy image I(x)I(x) using a sophisticated deep model.
  2. Estimating the global atmospheric light AA using some kind of empirical method.
  3. Estimating the clean image J(x)J(x) using the aforementioned formulation of the same.
However, such a procedure for image-dehazing leads to a sub-optimal solution that does not directly minimize the image reconstruction errors. The separate estimation of t(x)t(x) and AA will cause accumulated or even amplified errors, when combining them together to calculate the clean image J(x)J(x).

The Core Idea Behind AODNet

The core idea behind AODNet is to unify the transmission matrix t(x)t(x) and the global atmospheric light AA into a single expression given by K(x)K(x). Therefore, the formulation of the clean image can be transformed to...
J(x)=K(x)I(x)K(x)+1\Large{J(x)=K(x) I(x)-K(x)+1}

Since K(x)K(x) is dependent on I(x)I(x), we can then build a deep convolutional model with the hazy images as input, that minimizes the reconstruction error between the output J(x)J(x) and the ground-truth clean image.
The overall architecture of AODNet can be summarized by the following diagram
Source: Figure 4(a) from https://arxiv.org/pdf/1707.06543.pdf
Note, that K(x)K(x) is just t(x)t(x) and AA translated to the pixel domain. We can express it as K(x)=1t(x)(I(x)A)+(Ab)I(x)1K(x)=\frac{\frac{1}{t(x)}(I(x)-A)+(A-b)}{I(x)-1}. However, the exact expression of K(x)K(x) is not relevant for us, since we would train a deep convolutional neural network to approximate this expression.
💡

Implementing the AODNet

As per the architecture summarized in the aforementioned diagram, we can see that the AODNet Model essentially consists of two modules:
  1. A K-estimation module that estimates K(x)K(x) from the hazy image I(x)I(x). This module is responsible for estimating the depth and relative haze level.
  2. A clean image generation module that generated the clean image from K(x)K(x) using the transformed formulation of the atmospheric scattering model.

Implementation of the AODNet Architecture
1




Training AODNet

Now that we have both the data input pipeline and the model implementation in place, let's train the model. We'll also explore the Weights & Biases callbacks for Keras which enables us to log and track the results of our training experiments and the model checkpoint easily on Weights & Biases.

The Loss, Optimizer and the Metrics
0


Experiment Tracking using Weights & Biases

Weights & Biases provides us with three callbacks for experiment tracking:
  • WandbMetricsLogger: This callback automatically tracks the training and validation logs (such as losses and metrics) and logs them on our Weights & Biases dashboard. It also tracks our system metrics, such as CPU and GPU utilization.
  • WandbModelCheckpoint: This callback periodically saves a Keras model or model weights and uploads it to Weights & Biases as an Artifact. This callback provides the following features:
    • Save the model that has achieved the best performance based on the monitor.
    • Save the model at the end of every epoch regardless of the performance.
    • Save the model at the end of the epoch or after a fixed number of training batches.
    • Save only model weights, or save the whole model.
    • Save the model either in SavedModel format or in h5 format.
  • WandbEvalCallback: This is an abstract base class to build Keras callbacks for model prediction visualization. We can build callbacks for visualizing model predictions on_epoch_end that can be passed to model.fit() for different machine-learning tasks such as classification, object detection, segmentation, etc. To use this, we can create our own visualization callback that inherits from this base callback class and implement the add_ground_truth and add_model_prediction methods. The base WandbEvalCallback class will take care of the following:
    • Initialize data_table for logging the ground truth and pred_table for predictions.
    • The data uploaded to data_table is used as a reference for the pred_table. This is to reduce the memory footprint. The data_table_ref is a list that can be used to access the referenced data. We will demonstrate this with an example in the training notebook.
    • Log the tables to Weights & Biases as Artifacts.
    • Each new pred_table is logged as a new version with aliases.


Keras callbacks for logging into Weights & Biases
0

Now, let's take a look at the training logs:

Training Logs
2




Evaluation Results


Evaluation Results
2




Interactive Demo

We also created an interactive inference demo using Gradio that enables us to easily perform inference on hazy images of our choice and log them into our Weights & Biases dashboard for reproducibility and future reference.



Demo
25




Conclusion




Similar Reports


Iterate on AI agents and models faster. Try Weights & Biases today.
List<File<(table)>>
List<File<(table)>>