Views
No views yet
| 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) |
| Metric | Value |
|---|---|
| PSNR | 34.04 dB |
| SSIM | 0.9670 |
1import torch
2from huggingface_hub import hf_hub_download
3
4# Load model
5ckpt = torch.load(hf_hub_download("GeraldNdawula/Watermark_Removal_UNet", "unet_final.pth"))
6model = UNet()
7model.load_state_dict(ckpt["model"])
8model.eval()
9
10# Inference
11import torchvision.transforms as transforms
12transform = transforms.Compose([
13 transforms.Resize((128, 128)),
14 transforms.ToTensor(),
15 transforms.Normalize([0.5]*3, [0.5]*3)
16])
17
18watermarked_tensor = transform(your_pil_image).unsqueeze(0)
19with torch.no_grad():
20 clean_tensor = model(watermarked_tensor)
21
22# Back to PIL
23denorm = transforms.Normalize([-1,-1,-1], [2,2,2])
24clean_img = transforms.ToPILImage()(denorm(clean_tensor.squeeze()).clamp(0,1))