This repository by Diffusion Dynamics, showcases the core technology behind the watermark segmentation capabilities of our first product, clear.photo. This work leverages insights from research on diffusion models for image restoration tasks.

Effective watermark removal hinges on accurately identifying the watermark’s precise location and shape within the image. This code tackles the first crucial step: watermark segmentation.

We present a deep learning approach trained to generate masks highlighting watermark regions. This repository focuses on segmenting logo-based watermarks, demonstrating a robust technique adaptable to various watermark types. The methodologies employed draw inspiration from advancements in image segmentation.

This repository aims to consolidate key ideas from recent research in visible watermark removal and segmentation, including techniques presented in:

• arXiv:2108.03581 “Visible Watermark Removal via Self-calibrated Localization and Background Refinement”
• arXiv:2012.07616 “WDNet: Watermark-Decomposition Network for Visible Watermark Removal”
• arXiv:2312.14383 “Removing Interference and Recovering Content Imaginatively for Visible Watermark Removal”
• arXiv:2502.02676 “Blind Visible Watermark Removal with Morphological Dilation”

It distills these concepts into a minimal, functional codebase focused purely on the segmentation task. The goal is to provide a clear, understandable baseline that is easy to modify and build upon, even allowing for fine-tuning on consumer hardware like laptops with Apple M-series chips. It serves as a foundational example demonstrating the core techniques applicable to building more complex tools like clear.photo.

A typical advanced watermark removal pipeline involves:

1. Segmentation: Generating a precise mask that isolates the watermark pixels from the background image content.
2. Inpainting/Restoration: Using the mask to guide an algorithm (often generative, like diffusion models) to intelligently ‘fill in’ the region previously occupied by the watermark, seamlessly reconstructing the underlying image.

This project provides the necessary tools to train a watermark segmentation model and use it for inference. Key components include:

• watermark-segmentation.ipynb: A Jupyter notebook containing the end-to-end workflow:
  • Model definition (utilizing the segmentation_models.pytorch library).
  • Dataloader setup using dataset.py.
  • Training loop implementation using pytorch-lightning.
  • Inference logic to apply the trained model to new images.
• dataset.py: A Python script defining the Dataset class responsible for generating training data. It dynamically applies logo watermarks (from the logos/ directory) onto background images with randomized properties (scale, rotation, opacity, position, blend mode) to create diverse training samples and their corresponding ground truth masks. This data augmentation strategy is key to the model’s robustness.
• requirements.txt: Lists all necessary Python dependencies.
• *.pth: Model weights from different training epochs.
• logos/: A directory for sample watermark logos. Populate this with logos relevant to your use case.
• lightning_logs/: Default directory where PyTorch Lightning saves training logs and checkpoints.

Follow these steps to set up and run the project:

1. Prerequisites:

• Python 3.10 or later.
• wget and unzip (or equivalent tools for downloading and extracting datasets).

python --version
# Ensure it shows 3.10.x or newer

2. Clone the Repository:

git clone https://github.com/Diffusion-Dynamics/watermark-segmentation
cd watermark-segmentation

3. Install Dependencies:

It’s recommended to use a virtual environment:

python -m venv venv
source venv/bin/activate # On Windows use `virtualScriptsactivate`
pip install -r requirements.txt