在Python中使用SWCNN去除水印

shizidushu發表於2024-07-17
PythonCNN

在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

相關文章