christophernavas/watermark-remover

Watermark Remover

Modèle de segmentation pour détecter et supprimer les watermarks dans les images.

Développé par DevynLabs dans le cadre du projet PixelForge.

Description

Watermark Remover utilise un pipeline en deux étapes :

1. Segmentation : UNet++ avec encodeur EfficientNet-B4 pour détecter les watermarks
2. Inpainting : LaMa pour reconstruire les zones masquées

Image → UNet++ (EfficientNet-B4) → Masque → LaMa → Image propre

Performance

Métrique	Score
IoU (validation)	97.48%
Taux de détection	95% (20/21 images)

Détails

• Entraîné sur des watermarks Banana (style texte semi-transparent)
• 1 faux négatif sur fond très lumineux
• Excellent sur watermarks similaires au dataset

Architecture

Composant	Valeur
Encoder	EfficientNet-B4 (pretrained ImageNet)
Decoder	UNet++
Input size	512x512
Output	Masque binaire (probabilité watermark)
Inpainting	LaMa (simple-lama-inpainting)

Usage

Installation

pip install segmentation-models-pytorch simple-lama-inpainting torch torchvision pillow

Détection seule

import torch
from PIL import Image
import segmentation_models_pytorch as smp
from torchvision import transforms

# Charger le modèle
model = smp.UnetPlusPlus(
    encoder_name="efficientnet-b4",
    encoder_weights=None,
    in_channels=3,
    classes=1,
)

# Charger les poids depuis HuggingFace
from huggingface_hub import hf_hub_download
weights_path = hf_hub_download("christophernavas/watermark-remover", "segmenter.pth")
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()

# Préparer l image
transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

image = Image.open("image_with_watermark.png").convert("RGB")
input_tensor = transform(image).unsqueeze(0)

# Prédiction
with torch.no_grad():
    mask = torch.sigmoid(model(input_tensor))
    mask = (mask > 0.5).float()

# mask contient le masque binaire des watermarks détectés

Pipeline complet (détection + suppression)

from simple_lama_inpainting import SimpleLama

# Après avoir obtenu le masque...
lama = SimpleLama()
result = lama(image, mask)
result.save("image_clean.png")

Training

• Framework: PyTorch + segmentation-models-pytorch
• Loss: BCE + Dice Loss
• Optimizer: Adam (lr=1e-4)
• Augmentations: Rotation, flip, color jitter, noise
• Platform: Modal (GPU T4)
• Epochs: 50

Dataset

Images synthétiques générées avec des watermarks Banana :

• Variations de position, taille, opacité
• Différents fonds (photos, illustrations)

Cas d usage

• Nettoyage d images pour e-commerce
• Préparation de datasets ML
• Restoration de photos

Limitations

• Optimisé pour watermarks textuels semi-transparents (style Banana)
• Peut avoir des difficultés avec :
  • Watermarks très opaques
  • Watermarks sur fonds très lumineux/blancs
  • Logos complexes (non-textuels)
• LaMa peut introduire des artefacts sur les textures complexes

License

MIT - Usage commercial autorisé.

Citation

@misc{watermarkremover2024,
  author = {Christopher Navas},
  title = {Watermark Remover: UNet++ Segmentation for Watermark Detection},
  year = {2024},
  publisher = {HuggingFace},
  url = {https://huggingface.co/christophernavas/watermark-remover}
}

Links

• DevynLabs
• PixelForge Documentation
• segmentation-models-pytorch
• LaMa Inpainting