Rich Feature Distillation with Feature Affinity Module for Efficient Image Dehazing

TL;DR

This paper presents a lightweight, efficient image dehazing framework that leverages rich feature distillation from a pre-trained super-resolution model, significantly improving performance while reducing model size.

Contribution

Introduces a novel feature affinity module for heterogeneous knowledge distillation, enhancing dehazing performance with a lightweight model.

Findings

Achieves up to 15% PSNR improvement on RESIDE-Standard dataset.

Reduces model size by approximately 20 times.

Demonstrates robustness on synthetic and real-world haze images.

Abstract

Single-image haze removal is a long-standing hurdle for computer vision applications. Several works have been focused on transferring advances from image classification, detection, and segmentation to the niche of image dehazing, primarily focusing on contrastive learning and knowledge distillation. However, these approaches prove computationally expensive, raising concern regarding their applicability to on-the-edge use-cases. This work introduces a simple, lightweight, and efficient framework for single-image haze removal, exploiting rich "dark-knowledge" information from a lightweight pre-trained super-resolution model via the notion of heterogeneous knowledge distillation. We designed a feature affinity module to maximize the flow of rich feature semantics from the super-resolution teacher to the student dehazing network. In order to evaluate the efficacy of our proposed framework, its performance as a plug-and-play setup to a baseline model is examined. Our experiments are carried out on the RESIDE-Standard dataset to demonstrate the robustness of our framework to the synthetic and real-world domains. The extensive qualitative and quantitative results provided establish the effectiveness of the framework, achieving gains of upto 15\% (PSNR) while reducing the model size by $\sim$20 times.