Universal and Quality-Preserving Watermark Removal based on Unpaired Learning

We propose a universal and quality-preserving watermark removal method based on unpaired learning. The targeted invisible watermarks include Digimarc, HiDDeN, StegaStamp, DwtDct, DwtDctSvd and RivaGAN. The remover is based on DDPM-UNet, and the discriminator utilizes a ConvNeXt model. Other models like simple UNet are also provided.

Dataset

The project uses the ILSVRC2012 dataset, which is a subset of the ImageNet dataset. The dataset can be accessed via the following link: ImageNet

Project Structure

UUMR
├── configs                  # Configuration files
├── Digimarc                 # Code related to Digimarc
├── HiDDeN_pytorch           # PyTorch implementation of HiDDeN
├── invisible_watermark      # Project code for DwtDct, DwtDctSvd and RivaGAN
├── StegaStamp_tensorflow    # StegaStamp project code
├── tool                     # Utility scripts
├── unet                     # UNet models, including DDPM-UNet and standard UNet
├── dataset.py               # Dataset class
├── Discriminator.py         # Custom discriminator (less effective)
├── get_config.py            # Configuration class
├── pretrain_D.py            # Pretraining script for discriminator D
├── pretrain_G.py            # Pretraining script for generator G
├── README.md                # Project documentation
├── requirements.txt         # Project dependencies
├── slide_test_clean.py      # Sliding window test for clean images
├── slide_test.py            # Sliding window-based removal of StegaStamp watermark
├── test_clean.py            # Testing clean images
├── test.py                  # Removing other types of watermarks
├── train_pair.py            # Paired training script
├── train.py                 # Training generator G and discriminator D
└── util.py                  # Utility functions

Environment Setup

Follow the steps below to configure and run the project:

1. Create and activate a conda virtual environment:

   conda create -n uuwr python=3.8
   conda activate uuwr

2. Install project dependencies: Use the following command to install the required packages:

   pip install -r requirements.txt

Project Workflow

Please ensure that all file paths are modified as needed before running the steps below.

Step 1: Pretrain Generator G

Run the pretrain_G.py script to obtain a pretrained generator G. This step trains the model on the specified dataset and saves the pretrained generator model.

Step 2: Pretrain Discriminator D and Joint Training

Run the train.py script to train the watermark remover G. This involves training the discriminator D with the generator G fixed, followed by training the generator G with the discriminator D fixed. My training records can be found in tool/record.txt. The ones marked with ★ are the best models for the watermark type, while those marked with ☆ are the second-best models for the watermark type.

Step 3: Testing

Run test.py or slide_test.py to test the effectiveness of watermark removal.

• test.py is used to remove other types of watermarks.
• slide_test.py applies a sliding window test on clean images.

References:

1. invisible-watermark
2. HiDDeN
3. StegaStamp
4. Pytorch-UNet