GeraldNdawula/Watermark_Removal_UNet

Watermark Removal U-Net

A U-Net model trained to remove synthetic proof-style watermarks from portrait images.

Training

Setting	Value
Architecture	U-Net (64→128→256→512 features)
Dataset	GeraldNdawula/Watermark_Dataset
Image size	128×128
Epochs	30
Optimizer	Adam (lr=0.0001)
Loss	L1 + Perceptual (VGG16, λ=0.1)

Final Metrics

Metric	Value
PSNR	34.04 dB
SSIM	0.9670

Usage

import torch
from huggingface_hub import hf_hub_download

# Load model
ckpt = torch.load(hf_hub_download("GeraldNdawula/Watermark_Removal_UNet", "unet_final.pth"))
model = UNet()
model.load_state_dict(ckpt["model"])
model.eval()

# Inference
import torchvision.transforms as transforms
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize([0.5]*3, [0.5]*3)
])

watermarked_tensor = transform(your_pil_image).unsqueeze(0)
with torch.no_grad():
    clean_tensor = model(watermarked_tensor)

# Back to PIL
denorm = transforms.Normalize([-1,-1,-1], [2,2,2])
clean_img = transforms.ToPILImage()(denorm(clean_tensor.squeeze()).clamp(0,1))