UUWR/test.py at main · yfjn/UUWR · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
from unet.unet_ddpm import *
import cv2
import torchvision
from get_config import parse_args_and_config
import torch.nn as nn
import os
import numpy as np
from tqdm import tqdm
import random

method = 'rivaGan'
length = 4
domain = 'rgb'
mode = 'val'
amount = 100
# modelG_path = 'saved_dwtDct/modelG_240803-23:24_5000_37.77.pth'
# modelG_path = 'saved_dwtDctSvd/modelG_240802-0027_17000_37.40.pth'
# modelG_path = 'saved_Digimarc/modelG_240904-2344_2000_39.45.pth'
# modelG_path = 'saved_rivaGan/modelG_240906-1440_3000_39.83.pth'
# modelG_path = 'saved_HiDDeN/modelG_240822-1112_5500_35.67.pth'  # 无失真训练的
# modelG_path = 'saved_HiDDeN/modelG_240908-1624_4000_33.02.pth'  # 实测图像质量不佳
# modelG_path = 'saved_HiDDeN/modelG_241004-0815_1000_32.01.pth'
# modelG_path = 'saved_StegaStamp/modelG_240809-1814_19000_28.17.pth'
modelG_path = 'saved_combine/modelG_241009-1112_10000_29.88_28.86_33.74.pth'  # 更佳
# modelG_path = 'saved_combine/modelG_241010-1830_31000_30.75_28.83_36.59.pth'  # 有些BER不高
if length == 4:
    input_folder = f'/media/dongli911/Documents/Datasets/watermark_dataset/{method}/wm_{mode}'
    output_folder = f'/media/dongli911/Documents/Datasets/watermark_dataset/{method}/rm_{mode}'
else:
    input_folder = f'/media/dongli911/Documents/Datasets/watermark_dataset/{method}/wm_{mode}_{length}'
    output_folder = f'/media/dongli911/Documents/Datasets/watermark_dataset/{method}/rm_{mode}_{length}'

transform = torchvision.transforms.ToTensor()
_, config = parse_args_and_config()
modelG = ModelG(config).cuda()
# modelD = torchvision.models.convnext_tiny(pretrained=True)
# modelD.classifier[2] = nn.Linear(768, 256)
# modelD.classifier.add_module("3", nn.Sequential(nn.GELU(), nn.Linear(256, 1)))
# modelD.cuda()
# modelD.load_state_dict(torch.load(rf"saved_{method}\modelD_20240609_35000_29.57.pth", map_location='cpu'), strict=True)
modelG.load_state_dict(torch.load(modelG_path, map_location='cpu'), strict=True)

if not os.path.exists(output_folder):
    os.makedirs(output_folder)
else:
    for filename in os.listdir(output_folder):
        file_path = os.path.join(output_folder, filename)
        os.remove(file_path)

image_files = [f for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
image_files = random.sample(image_files, min(amount, len(image_files)))
for filename in tqdm(image_files):
    img_path = os.path.join(input_folder, filename)
    img = cv2.imread(img_path)
    assert img.shape == (256, 256, 3)
    if domain == 'rgb':
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
    img = transform(img).unsqueeze(0).cuda()
    modelG.eval()
    img = torch.clamp(modelG(img), max=1, min=0)

    # 保存处理后的图片
    output_path = os.path.join(output_folder, filename)
    image_np = (img.squeeze(0).permute(1, 2, 0).detach().cpu().numpy() * 255).astype(np.uint8)
    if domain == 'rgb':
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    else:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_YUV2BGR)
    cv2.imwrite(output_path, image_np)