在Python中使用SWCNN去除水印
在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