WMFormer++: Nested Transformer for Visible Watermark Removal
via Implict Joint Learning

In the realm of digital copyright protection, digital watermarks have assumed a critical role. Early approaches (Pei and Zeng 2006; Park, Tai, and Kweon 2012) for watermark removal heavily relied on manually crafted features and necessitated user intervention for watermark localization, leading to usability challenges. Some researchers (Dekel et al. 2017; Gandelsman, Shocher, and Irani 2019) explored methods based on multiple images, but these demanded extensive prior knowledge and were limited in their applicability to specific samples.

Recent advancements in deep learning-based techniques have shown significant promise in various computer vision tasks, including visible watermark removal.  Hertz et al. (2019) pioneered a groundbreaking approach by blindly removing visual motifs from images and introducing the single encoder with multi-decoder architecture for multi-task watermark removal, achieving impressive performance.  Cun and Pun (2021) further enhanced watermark removal performance by proposing two-stage networks for prediction and refinement, solidifying the multi-decoder multi-stage framework as a mainstream solution for watermark removal. Building on this backdrop,  Liang et al. (2021) introduced a set of intricate modules to enhance the quality of generated images, while  Sun, Su, and Wu (2023) employed a contrastive learning strategy to disentangle high-level embedding semantic information of images and watermarks.

While the multi-decoder multi-stage framework has been instrumental, it is crucial to reassess its rationale and adopt a more concise architecture. By seeking a more streamlined solution, we can further advance research in the field of watermark removal, providing more practical and sustainable approaches for real-world applications.

Watermark removal shares a resemblance to image dehazing (He, Sun, and Tang 2010), deraining (Ren et al. 2019), and shadow removal (Cun, Pun, and Shi 2020) tasks in image restoration. They all involve recovering the source image from a damaged version, but specific differences exist that cannot be ignored in practical applications.

In dehazing and deraining, the interference factors (haze and raindrops) permeate the entire image, with repeated patterns and semantics present within and across images. In contrast, watermarks are usually localized in specific areas of the image, with each watermark representing unique information independently. Moreover, shadow removal deals with meaningless grey areas, while watermarks symbolize media copyright and typically exhibit meaningful and colorful content. These distinctions make watermark removal a distinct and challenging research domain.

Recently, the adoption of a novel Transformer architecture (Zamir et al. 2022) in image restoration has shown remarkable results, inspiring us to explore the application of similar modules in the field of watermark removal.

UNet (Ronneberger, Fischer, and Brox 2015) is a classic neural network known for its effectiveness in image semantic segmentation. Its U-shaped architecture with symmetric encoding and decoding pathways enables it to capture both global and local features, achieving remarkable performance in image segmentation tasks.

UNet++ (Zhou et al. 2018) is an improved version of UNet. It builds upon the original UNet by introducing multi-level encoding-decoding paths. This design allows the network to progressively fuse feature information across layers, resulting in a better grasp of semantic information at different scales and further enhancing segmentation performance.

While UNet has been widely applied as a foundational framework in the field of watermark removal, the application of the UNet++ architecture in watermark removal remains relatively unexplored and lesser-known.

For a given watermarked image $J\in\mathbb{R}^{H\times W\times 3}$, our network produces a watermark-free restored background image $\hat{I}\in\mathbb{R}^{H\times W\times 3}$ as output, with the grayscale watermark mask $\hat{M}\in\mathbb{R}^{H\times W\times 1}$ being generated as a byproduct. The network architecture consists of a multi-stage encoder and a nested decoder, with both components composed of fundamental Transformer blocks.

A novel type of Transformer (Zamir et al. 2022) serves as the basic building block of the entire network, composed of two key components: Cross-channel Multi-head Attention and Gated-Dconv Feed-Forward Network.

(i) Cross-channel Multi-head Attention. This component serves to boost local restoration quality while enabling comprehensive restoration region anchoring. Initially, the module utilizes $1\times 1$ convolutions and $3\times 3$ depth-wise convolutions to gather spatial information from neighboring pixels. Subsequently, it calculates cross-channel attention maps to implicitly establish global context semantics across pixels.

Given the input $X$, the attention for the i-th head is computed using the following formula:

	$\displaystyle Q_{i},K_{i},V_{i}$	$\displaystyle=XW_{i}^{Q},XW_{i}^{K},XW_{i}^{V},$		(1)
	$\displaystyle head_{i}$	$\displaystyle=V_{i}\text{ softmax}(\sigma K_{i}^{T}Q_{i}),$		(2)

where $Q_{i},K_{i}$, and $V_{i}$ represent the query, key, and value for each head $i$, respectively. $W_{i}^{Q},W_{i}^{K},$ and $W_{i}^{V}$ are three projection matrices. $\sigma$ is a learnable parameter.

(ii) Gated-Dconv Feed-Forward Network. This module utilizes a gating mechanism to enhance information flow within the network, facilitating the collaborative optimization of the two implicit branches: watermark localization and background restoration. This synergistic approach enables them to collectively refine their respective tasks.

The entire process can be represented using the following formula:

	$\displaystyle X_{1},X_{2}$	$\displaystyle=split\_half\left(XW_{1},axis=channel\right),$		(3)
	$\displaystyle Y$	$\displaystyle=\left(gelu\left(X_{1}\right)\odot X_{2}\right)W_{2}+X,$		(4)

where $W_{1}$ and $W_{2}$ correspond to the projection matrices responsible for augmenting and reducing the channel dimensions of $X$. The output of the module is denoted by $Y$.

The Encoder module follows a similar structure as the UNet (Ronneberger, Fischer, and Brox 2015), comprising four hierarchical levels. As data progresses through each level, the feature map dimensions remain unchanged, but its channel count doubles and size is halved before transitioning to the subsequent level. This progressive encoding leads to the compression of the original input, resulting in a latent code enriched with semantic information.

In general, the latent code generated by the final level acts as input for the Decoder. Notably, in contrast to the conventional approach, we harness the latent codes generated at each level during the decoding process. This synergistic interaction among nested decoders empowers the creation of UNet models with varying depths, allowing for a more adaptable and flexible architecture.

The latent code encapsulates two crucial aspects of information: the overall structure of the repair regions, including the position of the watermark within the image, and various reconstruction details. The latent codes from different levels of the encoder harbor distinct information due to variations in encoding depth. To capitalize on this diversity, we decode latent codes from different levels to construct independent UNet pathways of varying depths. These pathways are then interlinked through dense connections, creating a nested UNet architecture that enhances feature fusion. The integration of UNet pathways at different depths enables the network to precisely locate watermark regions. Simultaneously, through skip connections at the same level, distinct UNet pathways can learn specific reconstruction details to aid in background restoration.

For a specific pathway of the decoder, the decoding process mirrors the reverse of the encoding procedure. For each level’s output, the channel count is reduced by half while the size doubles before moving to the previous level. This stepwise decoding gradually restores the latent code back to enriched feature maps containing discrete watermark and background information. Finally, through the last prediction head, both watermark-free images and watermark masks are generated.

At each decoder level, a decision point arises regarding the output strategy. The network can either upsample the feature by a factor of two while reducing the number of channels, thereby progressing to the previous decoder level, or it can directly transmit the feature to the decoder at the same level. The network benefits from these intra-level skip connections due to the distinct semantic information encoded within features originating from different depths of latent codes, even when they are at the same decoder level.

In our proposed framework, we leverage the nested architecture as the backbone, which can be seen as an integration of UNets with varying depths. To ensure that each UNet is capable of both image restoration and watermark extraction, during the training phase, we allow the refined features from each UNet’s output to pass through a shared prediction head. This step facilitates the generation of corresponding watermark mask $\hat{M}^{d}$ and background image $\hat{I}^{d}$, which are then supervised using the ground truth watermark mask $M$ and background image $I$.

Consistent with previous works (Hertz et al. 2019), we employ binary cross-entropy loss to ensure the close alignment between the predicted watermark mask $\hat{M}^{d}$ and the ground-truth watermark mask $M$.

	$\displaystyle\mathcal{L}_{\text{mask}}^{d}=-\sum_{i,j}($	$\displaystyle\left(1-M_{i,j}\right)\log(1-\hat{M}_{i,j}^{d})$
		$\displaystyle+M_{i,j}\log\hat{M}_{i,j}^{d}).$		(5)

As for the reconstruction task, we utilize the $L_{1}$ loss to ensure that the refined watermark-free image closely approximates the ground-truth image.

$$\mathcal{L}_{\text{image}}^{d}=\|I-\hat{I}^{d}\|_{1}.$$

(6)

Moreover, inspired by the works of  (Cun and Pun 2021; Liang et al. 2021), we incorporate an additional deep perceptual loss (Johnson, Alahi, and Fei-Fei 2016) to enhance the output’s overall quality.

$$\mathcal{L}_{\text{perc}}^{d}=\sum_{k\in 1,2,3}\left\|\Phi_{vgg}^{k}(\hat{I}^{d})-\Phi_{vgg}^{k}(I)\right\|_{1},$$

(7)

where $\Phi_{vgg}^{k}(\cdot)$ denotes the activation map of k-th layer in the pre-trained VGG16 (Simonyan and Zisserman 2014).

In conclusion, we aggregate all the aforementioned loss functions and introduce controllable hyper-parameters to form the final loss function. The losses from different depths of the UNet are combined harmoniously, playing a collective role in shaping the overall loss.

$$\mathcal{L}_{\text{all}}=\sum_{d\in 1,2,3}\lambda^{d}\left(\mathcal{L}_{\text{image}}^{d}+\lambda_{\text{mask}}\mathcal{L}_{\text{mask}}^{d}+\lambda_{\text{perc}}\mathcal{L}_{\text{perc}}^{d}\right).$$

(8)

This paper leverages various datasets, namely LOGO-L, LOGO-H, LOGO-Gray, and CLWD, for conducting the experiments.

The LOGO series dataset contains varying watermarks, such as LOGO-L with 12,151 training and 2,025 testing images, LOGO-H with similar quantities but larger watermark sizes and transparency, and LOGO-Gray with grayscale watermarks.

CLWD dataset contains 60,000 watermarked images, with 160 colored watermarks for training and 10,000 images with 40 colored watermarks for testing. The watermarks are collected from open-sourced logo images websites. The original images in the training and test sets are randomly chosen from the PASCAL VOC 2012 (Everingham et al. 2015) dataset. The transparency of the watermarks in CLWD is set in the range of 0.3 to 0.7, and the size, locations, rotation angle, and transparency of each watermark are randomly set in different images.

We implemented our method using Pytorch (Paszke et al. 2019) and conducted experiments on the aforementioned datasets. For training, we set the input image size as $256\times 256$. The AdamW (Loshchilov and Hutter 2017) optimizer was chosen with an initial learning rate of 3e-4, a batch size of 8, and momentum parameters $\beta_{1}=0.9$ and $\beta_{2}=0.999$. The hyper-parameters $\lambda^{d}$, $\lambda_{mask}$, and $\lambda_{vgg}$ in Eqn.(8) were empirically set as 1, 1, and 0.25, respectively, after several trials by observing the quality of predicted masks and reconstructed images.

To evaluate the effectiveness of our method, we employ widely recognized metrics in the field. Specifically, on the LOGO series datasets, we use Peak Signal-to-Noise Ratio (PSNR), Structural Similarity (SSIM) (Wang et al. 2004), and the deep perceptual similarity metric (LPIPS) (Zhang et al. 2018) for evaluation. On the CLWD dataset, we utilize PSNR, SSIM, Root Mean Square Error (RMSE), and weighted RMSE (RMSEw) as evaluation metrics.

Analysis of Individual Modules: As depicted in Tab. 1, we conducted a comprehensive evaluation of each module’s effectiveness in our framework through incremental removal and addition. Initially, we started with a basic UNet structure, as referenced  (Ronneberger, Fischer, and Brox 2015), with its performance displayed in the first row of the table. Subsequently, by replacing all convolutional modules in UNet’s encoder and decoder with Transformer modules, we observed a notable improvement in performance, as shown in the second row. This highlights the significant potential of applying Transformers in the field of watermark removal.

Furthermore, we devised a nested Transformer model as illustrated in Fig. 2, and its performance is documented in the third row. It is worth noting that, for a fair comparison, the Transformer-based UNet used in the second row employed a wider UNet to achieve a similar parameter count as the model experimented with in the third row. The comparison between the second and third rows provides compelling evidence for the efficacy of the nested architecture.

Finally, We applied the supervision signals to UNets with different depths, and the performance is presented in the fourth row. The comparison between the third and fourth rows clearly demonstrates that the model with deep supervision exhibits superior performance, achieving the best results across all experiments. This further validates the significance of incorporating deep supervision in our approach.

Visualization of Individual Modules: For a more intuitive presentation, we also generated corresponding visual results to complement the findings in Tab. 1. By zooming in on specific regions of the restored background images, we aimed to provide a clearer display. As shown in Fig. 3, with the gradual completion of the full model, the quality of the images significantly improves. The step-by-step enhancement of image quality through the addition of modules serves as concrete evidence of the effectiveness of each design in our approach. These visual results further reinforce the significance of our contributions.

Background Restoration. The quantitative and qualitative comparison of our proposed WMFormer++ with other existing watermark removal methods are summarized in Tab. 3, Tab. 4, and Fig. 4. Among them, BVMR (Hertz et al. 2019), SplitNet (Cun and Pun 2021), SLBR (Liang et al. 2021), and DENet (Sun, Su, and Wu 2023) are the latest technologies dedicated to watermark removal and generally adopt the multi-decoder multi-stage framework. On the other hand, Unet (Ronneberger, Fischer, and Brox 2015), SIRF (Zhang, Ng, and Chen 2018), BS²AM (Cun and Pun 2020), and DHAN (Cun, Pun, and Shi 2020) are migrated from related tasks, such as blind image harmonization and shadow removal, and typically utilize a single-decoder branch for image restoration. Overall, the multi-decoder multi-stage methods tend to exhibit better performance.

In contrast, our framework follows a single-decoder approach and outperforms all the other methods on the four datasets, achieving a new state-of-the-art performance. Our results significantly surpass previous non-multi-decoder multi-stage methods and even demonstrate unprecedented superiority over multi-decoder multi-stage methods by a large margin. The comprehensive experimental results provide strong evidence for the effectiveness of our approach.

It is worth noting that our design does not rely on complex modules like those in Liang et al. (Liang et al. 2021) and Cun et al. (Cun and Pun 2021). Instead, we use only two fundamental components, the Transformer block, and a small number of convolutional layers, for the entire model construction. This simplicity further highlights the efficiency and effectiveness of our approach in achieving outstanding watermark removal results.

In Tab. 2, the quantitative evaluation indicates that our method achieves significantly better results compared to previous approaches in terms of F1 and IoU. Moreover, Fig. 5 visually presents the watermark masks generated by our method, showcasing their high quality and precision.

These results serve as compelling evidence of the effectiveness and robustness of our approach, not only in removing watermarks from images but also in generating accurate and complete watermark masks as a valuable side benefit. This highlights the holistic superiority of our method in watermark removal and its related applications.