在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 py39torchconda 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 osfrom 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
