Utilizing Multi-step Loss for Single Image Reflection Removal

TL;DR

This work addresses single-image reflection removal by introducing a generalizable multi-step loss mechanism for image-to-image translation tasks, augmented by a RefGAN-synthesized dataset and a Ranged Depth Map to focus on scene content. The approach combines an Ranged Depth Map-guided Reflection Removal Module with a two-stage UNet-based architecture and a loss set consisting of Pixel, Feature, and Gradient components accumulated over multiple steps, formalized as $L^t$ and $L_{total}$. RefGAN, built on Pix2Pix with a UNet generator and PatchGAN discriminator, generates $7115$ ambient-transmission pairs to boost training diversity. Empirically, the method achieves state-of-the-art performance on the $SIR^2$ benchmark and several real-world datasets, demonstrating strong generalization and practical impact for improving image quality in single-image scenarios.

Abstract

Image reflection removal is crucial for restoring image quality. Distorted images can negatively impact tasks like object detection and image segmentation. In this paper, we present a novel approach for image reflection removal using a single image. Instead of focusing on model architecture, we introduce a new training technique that can be generalized to image-to-image problems, with input and output being similar in nature. This technique is embodied in our multi-step loss mechanism, which has proven effective in the reflection removal task. Additionally, we address the scarcity of reflection removal training data by synthesizing a high-quality, non-linear synthetic dataset called RefGAN using Pix2Pix GAN. This dataset significantly enhances the model's ability to learn better patterns for reflection removal. We also utilize a ranged depth map, extracted from the depth estimation of the ambient image, as an auxiliary feature, leveraging its property of lacking depth estimations for reflections. Our approach demonstrates superior performance on the SIR^2 benchmark and other real-world datasets, proving its effectiveness by outperforming other state-of-the-art models.

Figures (6)

• Figure 1: (c) is the estimated depth map for the ambient image --(a)--, (d) is the ranged depth map with k=4.
• Figure 2: The architecture of our method. The Depth module estimates the depth of the ambient image. The RDepth module extracts the ranged depth map from the generated depth map. The R-CNN network predicts the reflection. The T-CNN predicts the target. The blue arrows resembles the multi-step loss mechanism during training.
• Figure 3: The model could learn the relations between nearby pixels that fall in the same range in Ranged Depth Map.
• Figure 4: RefGAN samples generated by pix2pix GAN model.
• Figure 5: Contrasting our method against state-of-the-art models on nature dataset (Rows 1-2) and the $SIR^2$ benchmark (Row 3)