在Python中使用SWCNN去除水印

说明

• 首次发表日期：2024-07-17
• SWCNN Github官方仓库： https://github.com/hellloxiaotian/SWCNN
• SWCNN 论文链接： https://arxiv.org/abs/2403.05807

准备

运行环境

首先创建一个conda环境，安装SWCNN官方建议的库：

conda create -n py39torch python=3.9
conda activate py39torch

conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/

pip install tensorboard==2.9.1 scikit-image==0.19.3 pyyaml==5.1 h5py==3.7.0 opencv-python==4.6.0.66 matplotlib==3.5.2 -i https://pypi.tuna.tsinghua.edu.cn/simple

Clone仓库、下载权重，准备图片

Clone下Github仓库：

git clone https://github.com/hellloxiaotian/SWCNN

然后从百度网盘 https://pan.baidu.com/s/1UbnK2F1FuIMKvqESNk5SvA (passwd: jlbv) 或者 Google Drive 中下载权重文件，只需要下载其中的models_SWCNN文件夹即可

新建data文件夹，并将下载的models_SWCNN文件夹放入其中。

mkdir -p data
# download models_SWCNN from https://drive.google.com/drive/folders/1-f7tVsytSXNjeqFkpEUMPbzkb4ZsDngQ or https://pan.baidu.com/s/1UbnK2F1FuIMKvqESNk5SvA (passwd: jlbv) 
# unzip and move folder into data

准备想要去除水印的图片：我使用PixWeaver随意生成了一张图片，通过一个网站随便加了点水印；然后在data文件夹下新建了一个名为my_images的目录，将水印图片放入其中。

运行推理

作者有提供一个test.py文件，见 https://github.com/hellloxiaotian/SWCNN/blob/main/test.py 。 这个文件说明了如何对模型进行测试，以下的代码基本上是从里面提取的。

首先，导入依赖：

import torch.nn as nn
import torch
import matplotlib.image as matImage
import os

from models import HN
from utils import *

然后，加载模型文件：

net = HN()
device_ids = [0]
model = nn.DataParallel(net, device_ids=device_ids).cuda()
model.load_state_dict(torch.load(os.path.join('data/models_SWCNN', "HNperL1n2nalpha1.0.pth")))
model.eval()

此处加载的是HNperL1n2nalpha1.0.pth这个模型文件，其中1.0代表的是alpha。

读取图片并进行预处理：

Img = cv2.imread("data/my_images/穿越到一个玄幻世界的文科生-watermarked.png")
Img = normalize(np.float32(Img[:, :, :]))
Img = np.expand_dims(Img, 0)
Img = np.transpose(Img, (0, 3, 1, 2))
_, _, w, h = Img.shape
w = int(int(w / 32) * 32)
h = int(int(h / 32) * 32)
Img = Img[:, :, 0:w, 0:h]
ISource = torch.Tensor(Img)

运行推理并保存图片：

with torch.no_grad():
    Out = torch.clamp(model(ISource), 0., 1.)
Out_np = Out.cpu().numpy()
pic = Out_np[0]

r, g, b = pic[0], pic[1], pic[2]
b = b[None, :, :]
r = r[None, :, :]
g = g[None, :, :]
pic = np.concatenate((b, g, r), axis=0)
pic = np.transpose(pic, (1, 2, 0))
matImage.imsave("data/my_images/out.jpg", pic)

效果

去除水印前：

去除水印后：

可以看出，大部分水印都去除了，但是还残留了一些。

另外，我有测试一张非AI生成的图片，效果相对比较好。

在实际应用中，我们可以准备一些素材进行训练（作者已开源训练脚本），相信会有不错的效果。

幕后花絮

运行作者提供的测试脚本时，发现了一些小问题，修改之后可以正常运行了，见我克隆的仓库： https://github.com/shizidushu/SWCNN