Add local files to repository

app.py

@@ -0,0 +1,273 @@
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+import requests
+from copy import deepcopy
+import cv2
+from test_gradio import load_image, image_editing
+
+import options.options as option
+from utils.JPEG import DiffJPEG
+from scipy.io.wavfile import read as wav_read
+from scipy.io import wavfile
+
+import os
+import math
+import argparse
+import random
+import logging
+
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from data.data_sampler import DistIterSampler
+
+from utils import util
+from data.util import read_img 
+from models import create_model as create_model_editguard
+import base64
+import gradio as gr
+
+from scipy.ndimage import zoom
+
+import matplotlib.pyplot as plt
+
+def img_to_base64(filepath):
+    with open(filepath, "rb") as img_file:
+        return base64.b64encode(img_file.read()).decode()
+
+logo_base64 = img_to_base64("../logo.png")
+
+html_content = f"""
+<div style='display: flex; align-items: center; justify-content: center; padding: 20px;'>
+    <img src='data:image/png;base64,{logo_base64}' alt='Logo' style='height: 50px; margin-right: 20px;'>
+    <strong><font size='8'>EditGuard<font></strong>
+</div>
+"""
+
+# Examples
+examples = [
+    ["../dataset/examples/0011.png"],
+    ["../dataset/examples/0012.png"],
+    ["../dataset/examples/0003.png"],
+    ["../dataset/examples/0004.png"],
+    ["../dataset/examples/0005.png"],
+    ["../dataset/examples/0006.png"],
+    ["../dataset/examples/0007.png"],
+    ["../dataset/examples/0008.png"],
+    ["../dataset/examples/0009.png"],
+    ["../dataset/examples/0010.png"],
+    ["../dataset/examples/0002.png"],
+]
+
+default_example = examples[0]
+
+def hiding(image_input, bit_input, model):
+
+    message = np.array([int(bit_input[i:i+1]) for i in range(0, len(bit_input), 1)])
+    message = message - 0.5
+    val_data = load_image(image_input, message)
+    model.feed_data(val_data)
+    container = model.image_hiding()
+
+    from PIL import Image
+    image = Image.fromarray(container)
+    return container, container
+
+def rand(num_bits=64):
+    random_str = ''.join([str(random.randint(0, 1)) for _ in range(num_bits)])
+    return random_str
+
+def ImageEdit(img, prompt, model_index):
+    image, mask = img["image"], np.float32(img["mask"])
+
+    received_image = image_editing(image, mask, prompt)
+    return received_image, received_image, received_image
+
+
+def imgae_model_select(ckp_index=0):
+    # options
+    opt = option.parse("options/test_editguard.yml", is_train=True)
+    # distributed training settings
+    opt['dist'] = False
+    rank = -1
+    print('Disabled distributed training.')
+
+    # loading resume state if exists
+    if opt['path'].get('resume_state', None):
+        # distributed resuming: all load into default GPU
+        device_id = torch.cuda.current_device()
+        resume_state = torch.load(opt['path']['resume_state'],
+                                    map_location=lambda storage, loc: storage.cuda(device_id))
+        option.check_resume(opt, resume_state['iter'])  # check resume options
+    else:
+        resume_state = None
+
+    # convert to NoneDict, which returns None for missing keys
+    opt = option.dict_to_nonedict(opt)
+    torch.backends.cudnn.benchmark = True
+    # create model
+
+    model = create_model_editguard(opt)
+
+    if ckp_index == 0:
+        model_pth = '../checkpoints/clean.pth'
+    print(model_pth)
+    model.load_test(model_pth)
+    return model
+
+
+
+def Gaussian_image_degradation(image, NL):
+    image = torch.from_numpy(np.transpose(image, (2, 0, 1)))
+    image = image.unsqueeze(0)
+    NL = NL / 255.0
+    noise = np.random.normal(0, NL, image.shape)
+    torchnoise = torch.from_numpy(noise).float()
+    y_forw = image + torchnoise
+    y_forw = torch.clamp(y_forw, 0, 1)
+    y_forw = y_forw.permute(0, 2, 3, 1)
+    y_forw = y_forw.cpu().detach().numpy().squeeze()
+
+    y_forw = (y_forw * 255.0).astype(np.uint8)
+    return y_forw, y_forw
+
+
+
+def JPEG_image_degradation(image, NL):
+    image = image.astype(np.float32)
+    image = torch.from_numpy(np.transpose(image, (2, 0, 1)))
+    image = image.unsqueeze(0)
+    JPEG = DiffJPEG(differentiable=True, quality=int(NL))
+    y_forw = JPEG(image)
+    y_forw = y_forw.permute(0, 2, 3, 1)
+    y_forw = y_forw.cpu().detach().numpy().squeeze()
+    y_forw = (y_forw * 255.0).astype(np.uint8)
+
+    return y_forw, y_forw
+
+
+def revealing(image_edited, input_bit, model_list, model):
+
+    if model_list==0:
+        number = 0.2
+    else:
+        number = 0.2
+
+    container_data = load_image(image_edited) ## load tampered images
+    model.feed_data(container_data)
+    mask, remesg = model.image_recovery(number)
+    mask = Image.fromarray(mask.astype(np.uint8))
+    remesg = remesg.cpu().numpy()[0]
+    remesg = ''.join([str(int(x)) for x in remesg])
+    bit_acc = calculate_similarity_percentage(input_bit, remesg)
+    return mask, remesg, bit_acc
+
+
+
+def calculate_similarity_percentage(str1, str2):
+
+    if len(str1) == 0:
+        return "原始版权水印未知"
+    elif len(str1) != len(str2):
+        return "输入输出水印长度不同"
+    total_length = len(str1)
+    same_count = sum(1 for x, y in zip(str1, str2) if x == y)
+    similarity_percentage = (same_count / total_length) * 100
+    return f"{similarity_percentage}%"
+
+
+
+# Description
+title = "<center><strong><font size='8'>EditGuard<font></strong></center>"
+
+css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
+
+with gr.Blocks(css=css, title="EditGuard") as demo:
+    gr.HTML(html_content)
+    model = gr.State(value = None)
+    save_h = gr.State(value = None)
+    save_w = gr.State(value = None)
+    sam_global_points = gr.State([])
+    sam_global_point_label = gr.State([])
+    sam_original_image = gr.State(value=None)
+    sam_mask = gr.State(value=None)
+
+    with gr.Tabs():
+        with gr.TabItem('多功能取证水印'):
+
+            DESCRIPTION = """
+            ## 使用方法：
+            - 上传图像和版权水印（64位比特序列），点击"嵌入水印"按钮，生成带水印的图像。
+            - 涂抹要编辑的区域，并使用Inpainting算法编辑图像。
+            - 点击"提取"按钮检测篡改区域并输出版权水印。"""
+            
+            gr.Markdown(DESCRIPTION)
+            save_inpainted_image = gr.State(value=None)
+            with gr.Column():
+                with gr.Row():
+                    model_list = gr.Dropdown(label="选择模型", choices=["模型1"], type = 'index')
+                    clear_button = gr.Button("清除全部")
+                with gr.Box():
+                    gr.Markdown("# 1. 嵌入水印")
+                    with gr.Row():
+                        with gr.Column():
+                            image_input = gr.Image(source='upload', label="原始图片", interactive=True, type="numpy", value=default_example[0])
+                            with gr.Row():
+                                bit_input = gr.Textbox(label="输入版权水印（64位比特序列）", placeholder="在这里输入...")
+                                rand_bit = gr.Button("🎲 随机生成版权水印")
+                            hiding_button = gr.Button("嵌入水印")
+                        with gr.Column():
+                            image_watermark = gr.Image(source="upload", label="带有水印的图片", interactive=True, type="numpy")
+
+
+                with gr.Box():
+                    gr.Markdown("# 2. 篡改图片")
+                    with gr.Row():
+                        with gr.Column():
+                            image_edit = gr.Image(source='upload',tool="sketch", label="选取篡改区域", interactive=True, type="numpy")
+                            inpainting_model_list = gr.Dropdown(label="选择篡改模型", choices=["模型1：SD_inpainting"], type = 'index')
+                            text_prompt = gr.Textbox(label="篡改提示词")
+                            inpainting_button = gr.Button("篡改图片")
+                        with gr.Column():
+                            image_edited = gr.Image(source="upload", label="篡改结果", interactive=True, type="numpy")
+                
+
+                with gr.Box():
+                    gr.Markdown("# 3. 提取水印&篡改区域")
+                    with gr.Row():
+                        with gr.Column():
+                            image_edited_1 = gr.Image(source="upload", label="待提取图片", interactive=True, type="numpy")
+                            
+                            revealing_button = gr.Button("提取")
+                        with gr.Column():
+                            edit_mask = gr.Image(source='upload', label="编辑区域蒙版预测", interactive=True, type="numpy")
+                            bit_output = gr.Textbox(label="版权水印预测")
+                            acc_output = gr.Textbox(label="水印预测准确率")
+                
+                gr.Examples(
+                            examples=examples,
+                            inputs=[image_input],
+                        )
+
+
+                model_list.change(
+                    imgae_model_select, inputs = [model_list], outputs=[model]
+                    )
+                hiding_button.click(
+                    hiding, inputs=[image_input, bit_input, model], outputs=[image_watermark, image_edit]
+                    )
+                rand_bit.click(
+                    rand, inputs=[], outputs=[bit_input]
+                    )
+
+
+                inpainting_button.click(
+                    ImageEdit, inputs = [image_edit, text_prompt, inpainting_model_list], outputs=[image_edited, image_edited_1, save_inpainted_image]
+                    )
+
+                revealing_button.click(
+                    revealing, inputs=[image_edited_1, bit_input, model_list, model], outputs=[edit_mask, bit_output, acc_output]
+                    )
+
+demo.launch(server_name="0.0.0.0", server_port=2004, share=True, favicon_path='../logo.png')

maskextract.py

@@ -0,0 +1,46 @@
+import os
+from PIL import Image
+import numpy as np
+import argparse
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--threshold', default=0.2, type=float, help='Path to option YMAL file.')
+    args = parser.parse_args()
+
+    input_folder = 'results/test_age-set'
+    output_folder = 'results/mask'
+
+    for filename in os.listdir(input_folder):
+        if filename.endswith('_0_0_LRGT.png'):
+            digits = filename.split('_')[0]
+            if digits.isdigit():
+                digits = int(digits)
+                
+                if digits >= 0 and digits <= 1000:
+
+                    input_path_LRGT = os.path.join(input_folder, filename)
+                    input_path_SR_h = os.path.join(input_folder, filename).replace('LRGT', 'SR_h')
+
+                    image_LRGT = Image.open(input_path_LRGT).convert("RGB")
+                    image_SR_h = Image.open(input_path_SR_h).convert("RGB")
+
+                    w, h = image_SR_h.size
+                    image_LRGT = image_LRGT.resize((w, h))
+                    
+                    array_LRGT = np.array(image_LRGT) / 255.
+                    array_SR_h = np.array(image_SR_h) / 255.
+                    
+                    residual = np.abs(array_LRGT - array_SR_h)
+
+                    threshold = args.threshold
+                    mask = np.where(residual > threshold, 1, 0)
+                    
+                    os.makedirs(output_folder, exist_ok=True)
+                    
+                    output_path = os.path.join(output_folder, str(digits+1).zfill(4)+'.png')
+
+                    mask = np.sum(mask, axis=2)
+
+                    mask_image = Image.fromarray((mask * 255).astype(np.uint8))
+                    mask_image.save(output_path)

test.py

@@ -0,0 +1,188 @@
+import os
+import math
+import argparse
+import random
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from data.data_sampler import DistIterSampler
+
+import options.options as option
+from utils import util
+from data import create_dataloader, create_dataset
+from models import create_model
+import numpy as np
+
+
+def init_dist(backend='nccl', **kwargs):
+    ''' initialization for distributed training'''
+    # if mp.get_start_method(allow_none=True) is None:
+    if mp.get_start_method(allow_none=True) != 'spawn':
+        mp.set_start_method('spawn')
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+def cal_pnsr(sr_img, gt_img):
+    # calculate PSNR
+    gt_img = gt_img / 255.
+    sr_img = sr_img / 255.
+
+    psnr = util.calculate_psnr(sr_img * 255, gt_img * 255)
+
+    return psnr
+
+def get_min_avg_and_indices(nums):
+    # Get the indices of the smallest 1000 elements
+    indices = sorted(range(len(nums)), key=lambda i: nums[i])[:900]
+    
+    # Calculate the average of these elements
+    avg = sum(nums[i] for i in indices) / 900
+    
+    # Write the indices to a txt file
+    with open("indices.txt", "w") as file:
+        for index in indices:
+            file.write(str(index) + "\n")
+    
+    return avg
+
+
+def main():
+    # options
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, help='Path to option YMAL file.')
+    parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
+                        help='job launcher')
+    parser.add_argument('--ckpt', type=str, default='/userhome/NewIBSN/EditGuard_open/checkpoints/clean.pth', help='Path to pre-trained model.')
+    parser.add_argument('--local_rank', type=int, default=0)
+    args = parser.parse_args()
+    opt = option.parse(args.opt, is_train=True)
+
+    # distributed training settings
+    if args.launcher == 'none':  # disabled distributed training
+        opt['dist'] = False
+        rank = -1
+        print('Disabled distributed training.')
+    else:
+        opt['dist'] = True
+        init_dist()
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+
+    # loading resume state if exists
+    if opt['path'].get('resume_state', None):
+        # distributed resuming: all load into default GPU
+        device_id = torch.cuda.current_device()
+        resume_state = torch.load(opt['path']['resume_state'],
+                                  map_location=lambda storage, loc: storage.cuda(device_id))
+        option.check_resume(opt, resume_state['iter'])  # check resume options
+    else:
+        resume_state = None
+
+    # convert to NoneDict, which returns None for missing keys
+    opt = option.dict_to_nonedict(opt)
+
+    torch.backends.cudnn.benchmark = True
+    # torch.backends.cudnn.deterministic = True
+
+    #### create train and val dataloader
+    dataset_ratio = 200  # enlarge the size of each epoch
+    for phase, dataset_opt in opt['datasets'].items():
+        print("phase", phase)
+        if phase == 'TD':
+            val_set = create_dataset(dataset_opt)
+            val_loader = create_dataloader(val_set, dataset_opt, opt, None)
+        elif phase == 'val':
+            val_set = create_dataset(dataset_opt)
+            val_loader = create_dataloader(val_set, dataset_opt, opt, None)
+        else:
+            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
+
+    # create model
+    model = create_model(opt)
+    model.load_test(args.ckpt)
+            
+    # validation
+    avg_psnr = 0.0
+    avg_psnr_h = [0.0]*opt['num_image']
+    avg_psnr_lr = 0.0
+    biterr = []
+    idx = 0
+    for image_id, val_data in enumerate(val_loader):
+        img_dir = os.path.join('results',opt['name'])
+        util.mkdir(img_dir)
+
+        model.feed_data(val_data)
+        model.test(image_id)
+
+        visuals = model.get_current_visuals()
+
+        t_step = visuals['SR'].shape[0]
+        idx += t_step
+        n = len(visuals['SR_h'])
+
+        a = visuals['recmessage'][0]
+        b = visuals['message'][0]
+
+        bitrecord = util.decoded_message_error_rate_batch(a, b)
+        print(bitrecord)
+        biterr.append(bitrecord)
+
+        for i in range(t_step):
+
+            sr_img = util.tensor2img(visuals['SR'][i])  # uint8
+            sr_img_h = []
+            for j in range(n):
+                sr_img_h.append(util.tensor2img(visuals['SR_h'][j][i]))  # uint8
+            gt_img = util.tensor2img(visuals['GT'][i])  # uint8
+            lr_img = util.tensor2img(visuals['LR'][i])
+            lrgt_img = []
+            for j in range(n):
+                lrgt_img.append(util.tensor2img(visuals['LR_ref'][j][i]))
+
+            # Save SR images for reference
+            save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'SR'))
+            util.save_img(sr_img, save_img_path)
+
+            for j in range(n):
+                save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:d}_{:s}.png'.format(image_id, i, j, 'SR_h'))
+                util.save_img(sr_img_h[j], save_img_path)
+
+            save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'GT'))
+            util.save_img(gt_img, save_img_path)
+
+            save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'LR'))
+            util.save_img(lr_img, save_img_path)
+
+            for j in range(n):
+                save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:d}_{:s}.png'.format(image_id, i, j, 'LRGT'))
+                util.save_img(lrgt_img[j], save_img_path)
+
+            psnr = cal_pnsr(sr_img, gt_img)
+            psnr_h = []
+            for j in range(n):
+                psnr_h.append(cal_pnsr(sr_img_h[j], lrgt_img[j]))
+            psnr_lr = cal_pnsr(lr_img, gt_img)
+
+            avg_psnr += psnr
+            for j in range(n):
+                avg_psnr_h[j] += psnr_h[j]
+            avg_psnr_lr += psnr_lr
+
+    avg_psnr = avg_psnr / idx
+    avg_biterr = sum(biterr) / len(biterr)
+    print(get_min_avg_and_indices(biterr))
+
+    avg_psnr_h = [psnr / idx for psnr in avg_psnr_h]
+    avg_psnr_lr = avg_psnr_lr / idx
+    res_psnr_h = ''
+    for p in avg_psnr_h:
+        res_psnr_h+=('_{:.4e}'.format(p))
+    print('# Validation # PSNR_Cover: {:.4e}, PSNR_Secret: {:s}, PSNR_Stego: {:.4e},  Bit_Error: {:.4e}'.format(avg_psnr, res_psnr_h, avg_psnr_lr, avg_biterr))
+
+
+if __name__ == '__main__':
+    main()

test_gradio.py

@@ -0,0 +1,92 @@
+import sys
+
+import os
+import math
+import argparse
+import random
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from data.data_sampler import DistIterSampler
+
+import options.options as option
+from utils import util
+from data.util import read_img 
+from data import create_dataloader, create_dataset
+from models import create_model
+import numpy as np
+from PIL import Image
+from diffusers import StableDiffusionInpaintPipeline
+
+
+def init_dist(backend='nccl', **kwargs):
+    ''' initialization for distributed training'''
+    # if mp.get_start_method(allow_none=True) is None:
+    if mp.get_start_method(allow_none=True) != 'spawn':
+        mp.set_start_method('spawn')
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+
+def load_image(image, message = None):
+    # img_GT = read_img(None, image_path)
+    img_GT = image / 255
+    # print(img_GT)
+    img_GT = img_GT[:, :, [2, 1, 0]]
+    img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float().unsqueeze(0)
+    img_GT = torch.nn.functional.interpolate(img_GT, size=(512, 512), mode='nearest', align_corners=None)
+    img_GT = img_GT.unsqueeze(0)
+
+    _, T, C, W, H = img_GT.shape
+    list_h = []
+    R = 0
+    G = 0
+    B = 255
+    image = Image.new('RGB', (W, H), (R, G, B))
+    result = np.array(image) / 255.
+    expanded_matrix = np.expand_dims(result, axis=0) 
+    expanded_matrix = np.repeat(expanded_matrix, T, axis=0)
+    imgs_LQ = torch.from_numpy(np.ascontiguousarray(expanded_matrix)).float()
+    imgs_LQ = imgs_LQ.permute(0, 3, 1, 2)
+    imgs_LQ = torch.nn.functional.interpolate(imgs_LQ, size=(W, H), mode='nearest', align_corners=None)
+    imgs_LQ = imgs_LQ.unsqueeze(0)
+
+    list_h.append(imgs_LQ)
+
+    list_h = torch.stack(list_h, dim=0)
+
+    return {
+            'LQ': list_h,
+            'GT': img_GT,
+            'MES': message
+        }
+
+
+def image_editing(image_numpy, mask_image, prompt):
+
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-inpainting",
+            torch_dtype=torch.float16,
+        ).to("cuda")
+    
+    pil_image = Image.fromarray(image_numpy)
+    print(mask_image.shape)
+    print("maskmin", mask_image.min(), "maskmax", mask_image.max())
+    mask_image = Image.fromarray(mask_image.astype(np.uint8)).convert("L")
+    image_init = pil_image.convert("RGB").resize((512, 512))
+    
+    h, w = mask_image.size
+    
+    image_inpaint = pipe(prompt=prompt, image=image_init, mask_image=mask_image, height=w, width=h).images[0]
+    image_inpaint = np.array(image_inpaint) / 255.
+    image = np.array(image_init) / 255.
+    mask_image = np.array(mask_image)
+    mask_image = np.stack([mask_image] * 3, axis=-1) / 255.
+    mask_image = mask_image.astype(np.uint8)
+    image_fuse = image * (1 - mask_image) + image_inpaint * mask_image
+
+    return image_fuse

train.py

@@ -0,0 +1,280 @@
+import os
+import math
+import argparse
+import random
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from data.data_sampler import DistIterSampler
+
+import options.options as option
+from utils import util
+from data import create_dataloader, create_dataset
+from models import create_model
+
+
+def init_dist(backend='nccl', **kwargs):
+    ''' initialization for distributed training'''
+    # if mp.get_start_method(allow_none=True) is None:
+    if mp.get_start_method(allow_none=True) != 'spawn':
+        mp.set_start_method('spawn')
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+def cal_pnsr(sr_img, gt_img):
+    # calculate PSNR
+    gt_img = gt_img / 255.
+    sr_img = sr_img / 255.
+    psnr = util.calculate_psnr(sr_img * 255, gt_img * 255)
+
+    return psnr
+
+def main():
+    # options
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, help='Path to option YMAL file.')  # config 文件
+    parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
+                        help='job launcher')
+    parser.add_argument('--local_rank', type=int, default=0)
+    args = parser.parse_args()
+    opt = option.parse(args.opt, is_train=True)
+
+    # distributed training settings
+    if args.launcher == 'none':  # disabled distributed training
+        opt['dist'] = False
+        rank = -1
+        print('Disabled distributed training.')
+    else:
+        opt['dist'] = True
+        init_dist()
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+
+    # loading resume state if exists
+    if opt['path'].get('resume_state', None):
+        # distributed resuming: all load into default GPU
+        device_id = torch.cuda.current_device()
+        resume_state = torch.load(opt['path']['resume_state'],
+                                  map_location=lambda storage, loc: storage.cuda(device_id))
+        # resume_state = torch.load(opt['path']['resume_state'],
+        #                           map_location=lambda storage, loc: storage.cuda(device_id), strict=False)
+        option.check_resume(opt, resume_state['iter'])  # check resume options
+    else:
+        resume_state = None
+
+    # mkdir and loggers
+    if rank <= 0:  # normal training (rank -1) OR distributed training (rank 0)
+        if resume_state is None:
+            util.mkdir_and_rename(
+                opt['path']['experiments_root'])  # rename experiment folder if exists
+            util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
+                         and 'pretrain_model' not in key and 'resume' not in key))
+
+        # config loggers. Before it, the log will not work
+        util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO,
+                          screen=True, tofile=True)
+        util.setup_logger('val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO,
+                          screen=True, tofile=True)
+        logger = logging.getLogger('base')
+        logger.info(option.dict2str(opt))
+        # tensorboard logger
+        if opt['use_tb_logger'] and 'debug' not in opt['name']:
+            version = float(torch.__version__[0:3])
+            if version >= 1.1:  # PyTorch 1.1
+                from torch.utils.tensorboard import SummaryWriter
+            else:
+                logger.info(
+                    'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
+                from tensorboardX import SummaryWriter
+            tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
+    else:
+        util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
+        logger = logging.getLogger('base')
+
+    # convert to NoneDict, which returns None for missing keys
+    opt = option.dict_to_nonedict(opt)
+
+    # random seed
+    seed = opt['train']['manual_seed']
+    if seed is None:
+        seed = random.randint(1, 10000)
+    if rank <= 0:
+        logger.info('Random seed: {}'.format(seed))
+    util.set_random_seed(seed)
+
+    torch.backends.cudnn.benchmark = True
+    # torch.backends.cudnn.deterministic = True
+
+    #### create train and val dataloader
+    dataset_ratio = 200  # enlarge the size of each epoch
+    for phase, dataset_opt in opt['datasets'].items():
+        if phase == 'train':
+            train_set = create_dataset(dataset_opt)
+            train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
+            total_iters = int(opt['train']['niter'])
+            total_epochs = int(math.ceil(total_iters / train_size))
+            if opt['dist']:
+                train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio)
+                total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
+            else:
+                train_sampler = None
+            train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
+            if rank <= 0:
+                logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
+                    len(train_set), train_size))
+                logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
+                    total_epochs, total_iters))
+        elif phase == 'val':
+            val_set = create_dataset(dataset_opt)
+            val_loader = create_dataloader(val_set, dataset_opt, opt, None)
+            if rank <= 0:
+                logger.info('Number of val images in [{:s}]: {:d}'.format(
+                    dataset_opt['name'], len(val_set)))
+        else:
+            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
+    assert train_loader is not None
+
+    # create model
+    model = create_model(opt)
+    # resume training
+    if resume_state:
+        logger.info('Resuming training from epoch: {}, iter: {}.'.format(
+            resume_state['epoch'], resume_state['iter']))
+
+        start_epoch = resume_state['epoch']
+        current_step = resume_state['iter']
+        model.resume_training(resume_state)  # handle optimizers and schedulers
+    else:
+        current_step = 0
+        start_epoch = 0
+
+    # training
+    logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
+    for epoch in range(start_epoch, total_epochs + 1):
+        if opt['dist']:
+            train_sampler.set_epoch(epoch)
+        for _, train_data in enumerate(train_loader):
+            current_step += 1
+            if current_step > total_iters:
+                break
+            # training
+            model.feed_data(train_data)
+            model.optimize_parameters(current_step)
+
+            # update learning rate
+            model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])
+
+            # log
+            if current_step % opt['logger']['print_freq'] == 0:
+                logs = model.get_current_log()
+                message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
+                    epoch, current_step, model.get_current_learning_rate())
+                for k, v in logs.items():
+                    message += '{:s}: {:.4e} '.format(k, v)
+                    # tensorboard logger
+                    if opt['use_tb_logger'] and 'debug' not in opt['name']:
+                        if rank <= 0:
+                            tb_logger.add_scalar(k, v, current_step)
+                if rank <= 0:
+                    logger.info(message)
+
+            # validation
+            if current_step % opt['train']['val_freq'] == 0 and rank <= 0:
+                avg_psnr = 0.0
+                avg_psnr_h = [0.0]*opt['num_image']
+                avg_psnr_lr = 0.0
+                avg_biterr = 0.0
+                idx = 0
+                for image_id, val_data in enumerate(val_loader):
+                    img_dir = os.path.join(opt['path']['val_images'])
+                    util.mkdir(img_dir)
+
+                    model.feed_data(val_data)
+                    model.test(image_id)
+
+                    visuals = model.get_current_visuals()
+
+                    t_step = visuals['SR'].shape[0]
+                    idx += t_step
+                    n = len(visuals['SR_h'])
+
+                    avg_biterr += util.decoded_message_error_rate_batch(visuals['recmessage'][0], visuals['message'][0])
+
+                    for i in range(t_step):
+
+                        sr_img = util.tensor2img(visuals['SR'][i])  # uint8
+                        sr_img_h = []
+                        for j in range(n):
+                            sr_img_h.append(util.tensor2img(visuals['SR_h'][j][i]))  # uint8
+                        gt_img = util.tensor2img(visuals['GT'][i])  # uint8
+                        lr_img = util.tensor2img(visuals['LR'][i])
+                        lrgt_img = []
+                        for j in range(n):
+                            lrgt_img.append(util.tensor2img(visuals['LR_ref'][j][i]))
+
+                        # Save SR images for reference
+                        save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'SR'))
+                        util.save_img(sr_img, save_img_path)
+
+                        for j in range(n):
+                            save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:d}_{:s}.png'.format(image_id, i, j, 'SR_h'))
+                            util.save_img(sr_img_h[j], save_img_path)
+
+                        save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'GT'))
+                        util.save_img(gt_img, save_img_path)
+
+                        save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'LR'))
+                        util.save_img(lr_img, save_img_path)
+
+                        for j in range(n):
+                            save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:d}_{:s}.png'.format(image_id, i, j, 'LRGT'))
+                            util.save_img(lrgt_img[j], save_img_path)
+
+                        psnr = cal_pnsr(sr_img, gt_img)
+                        psnr_h = []
+                        for j in range(n):
+                            psnr_h.append(cal_pnsr(sr_img_h[j], lrgt_img[j]))
+                        psnr_lr = cal_pnsr(lr_img, gt_img)
+
+                        avg_psnr += psnr
+                        for j in range(n):
+                            avg_psnr_h[j] += psnr_h[j]
+                        avg_psnr_lr += psnr_lr
+
+                avg_psnr = avg_psnr / idx
+                avg_psnr_h = [psnr / idx for psnr in avg_psnr_h]
+                avg_psnr_lr = avg_psnr_lr / idx
+                avg_biterr = avg_biterr / idx
+
+                # log
+                res_psnr_h = ''
+                for p in avg_psnr_h:
+                    res_psnr_h+=('_{:.4e}'.format(p))
+
+                logger.info('# Validation # PSNR_Cover: {:.4e}, PSNR_Secret: {:s}, PSNR_Stego: {:.4e}, Bit_acc: {: .4e}'.format(avg_psnr, res_psnr_h, avg_psnr_lr, avg_biterr))
+                logger_val = logging.getLogger('val')  # validation logger
+                logger_val.info('<epoch:{:3d}, iter:{:8,d}> PSNR_Cover: {:.4e}, PSNR_Secret: {:s}, PSNR_Stego: {:.4e}, Bit_acc: {: .4e}'.format(
+                    epoch, current_step, avg_psnr, res_psnr_h, avg_psnr_lr, avg_biterr))
+                # tensorboard logger
+                if opt['use_tb_logger'] and 'debug' not in opt['name']:
+                    tb_logger.add_scalar('psnr', avg_psnr, current_step)
+
+            # save models and training states
+            if current_step % opt['logger']['save_checkpoint_freq'] == 0:
+                if rank <= 0:
+                    logger.info('Saving models and training states.')
+                    model.save(current_step)
+                    model.save_training_state(epoch, current_step)
+
+    if rank <= 0:
+        logger.info('Saving the final model.')
+        model.save('latest')
+        logger.info('End of training.')
+
+
+if __name__ == '__main__':
+    main()

train_bit.py

@@ -0,0 +1,246 @@
+import os
+import math
+import argparse
+import random
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from data.data_sampler import DistIterSampler
+
+import options.options as option
+from utils import util
+from data import create_dataloader, create_dataset
+from models import create_model
+
+
+def init_dist(backend='nccl', **kwargs):
+    ''' initialization for distributed training'''
+    # if mp.get_start_method(allow_none=True) is None:
+    if mp.get_start_method(allow_none=True) != 'spawn':
+        mp.set_start_method('spawn')
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+def cal_pnsr(sr_img, gt_img):
+    # calculate PSNR
+    gt_img = gt_img / 255.
+    sr_img = sr_img / 255.
+    psnr = util.calculate_psnr(sr_img * 255, gt_img * 255)
+
+    return psnr
+
+def main():
+    # options
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, help='Path to option YMAL file.')  # config 文件
+    parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
+                        help='job launcher')
+    parser.add_argument('--local_rank', type=int, default=0)
+    args = parser.parse_args()
+    opt = option.parse(args.opt, is_train=True)
+
+    # distributed training settings
+    if args.launcher == 'none':  # disabled distributed training
+        opt['dist'] = False
+        rank = -1
+        print('Disabled distributed training.')
+    else:
+        opt['dist'] = True
+        init_dist()
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+
+    # loading resume state if exists
+    if opt['path'].get('resume_state', None):
+        # distributed resuming: all load into default GPU
+        device_id = torch.cuda.current_device()
+        resume_state = torch.load(opt['path']['resume_state'],
+                                  map_location=lambda storage, loc: storage.cuda(device_id))
+        # resume_state = torch.load(opt['path']['resume_state'],
+        #                           map_location=lambda storage, loc: storage.cuda(device_id), strict=False)
+        option.check_resume(opt, resume_state['iter'])  # check resume options
+    else:
+        resume_state = None
+
+    # mkdir and loggers
+    if rank <= 0:  # normal training (rank -1) OR distributed training (rank 0)
+        if resume_state is None:
+            util.mkdir_and_rename(
+                opt['path']['experiments_root'])  # rename experiment folder if exists
+            util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
+                         and 'pretrain_model' not in key and 'resume' not in key))
+
+        # config loggers. Before it, the log will not work
+        util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO,
+                          screen=True, tofile=True)
+        util.setup_logger('val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO,
+                          screen=True, tofile=True)
+        logger = logging.getLogger('base')
+        logger.info(option.dict2str(opt))
+        # tensorboard logger
+        if opt['use_tb_logger'] and 'debug' not in opt['name']:
+            version = float(torch.__version__[0:3])
+            if version >= 1.1:  # PyTorch 1.1
+                from torch.utils.tensorboard import SummaryWriter
+            else:
+                logger.info(
+                    'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
+                from tensorboardX import SummaryWriter
+            tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
+    else:
+        util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
+        logger = logging.getLogger('base')
+
+    # convert to NoneDict, which returns None for missing keys
+    opt = option.dict_to_nonedict(opt)
+
+    # random seed
+    seed = opt['train']['manual_seed']
+    if seed is None:
+        seed = random.randint(1, 10000)
+    if rank <= 0:
+        logger.info('Random seed: {}'.format(seed))
+    util.set_random_seed(seed)
+
+    torch.backends.cudnn.benchmark = True
+    # torch.backends.cudnn.deterministic = True
+
+    #### create train and val dataloader
+    dataset_ratio = 200  # enlarge the size of each epoch
+    for phase, dataset_opt in opt['datasets'].items():
+        if phase == 'train':
+            train_set = create_dataset(dataset_opt)
+            train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
+            total_iters = int(opt['train']['niter'])
+            total_epochs = int(math.ceil(total_iters / train_size))
+            if opt['dist']:
+                train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio)
+                total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
+            else:
+                train_sampler = None
+            train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
+            if rank <= 0:
+                logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
+                    len(train_set), train_size))
+                logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
+                    total_epochs, total_iters))
+        elif phase == 'val':
+            val_set = create_dataset(dataset_opt)
+            val_loader = create_dataloader(val_set, dataset_opt, opt, None)
+            if rank <= 0:
+                logger.info('Number of val images in [{:s}]: {:d}'.format(
+                    dataset_opt['name'], len(val_set)))
+        else:
+            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
+    assert train_loader is not None
+
+    # create model
+    model = create_model(opt)
+    # resume training
+    if resume_state:
+        logger.info('Resuming training from epoch: {}, iter: {}.'.format(
+            resume_state['epoch'], resume_state['iter']))
+
+        start_epoch = resume_state['epoch']
+        current_step = resume_state['iter']
+        model.resume_training(resume_state)  # handle optimizers and schedulers
+    else:
+        current_step = 0
+        start_epoch = 0
+
+    # training
+    logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
+    for epoch in range(start_epoch, total_epochs + 1):
+        if opt['dist']:
+            train_sampler.set_epoch(epoch)
+        for _, train_data in enumerate(train_loader):
+            current_step += 1
+            if current_step > total_iters:
+                break
+            # training
+            model.feed_data(train_data)
+            model.optimize_parameters(current_step)
+
+            # update learning rate
+            model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])
+
+            # log
+            if current_step % opt['logger']['print_freq'] == 0:
+                logs = model.get_current_log()
+                message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
+                    epoch, current_step, model.get_current_learning_rate())
+                for k, v in logs.items():
+                    message += '{:s}: {:.4e} '.format(k, v)
+                    # tensorboard logger
+                    if opt['use_tb_logger'] and 'debug' not in opt['name']:
+                        if rank <= 0:
+                            tb_logger.add_scalar(k, v, current_step)
+                if rank <= 0:
+                    logger.info(message)
+
+            # validation
+            if current_step % opt['train']['val_freq'] == 0 and rank <= 0:
+                avg_psnr = 0.0
+                avg_psnr_h = [0.0]*opt['num_image']
+                avg_psnr_lr = 0.0
+                avg_biterr = 0.0
+                idx = 0
+                for image_id, val_data in enumerate(val_loader):
+                    img_dir = os.path.join(opt['path']['val_images'])
+                    util.mkdir(img_dir)
+
+                    model.feed_data(val_data)
+                    model.test(image_id)
+
+                    visuals = model.get_current_visuals()
+
+                    t_step = visuals['recmessage'].shape[0]
+                    idx += t_step
+                    n = 1
+                    # print(visuals['message'].shape)
+                    avg_biterr += util.decoded_message_error_rate_batch(visuals['recmessage'][0], visuals['message'][0])
+                    print(util.decoded_message_error_rate_batch(visuals['recmessage'][0], visuals['message'][0]))
+
+                    for i in range(t_step):
+
+                        gt_img = util.tensor2img(visuals['GT'][i])  # uint8
+                        lr_img = util.tensor2img(visuals['LR'][i])
+
+                        save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'GT'))
+                        util.save_img(gt_img, save_img_path)
+
+                        save_img_path = os.path.join(img_dir,'{:d}_{:d}_{:s}.png'.format(image_id, i, 'LR'))
+                        util.save_img(lr_img, save_img_path)
+                        psnr_lr = cal_pnsr(lr_img, gt_img)
+                        avg_psnr_lr += psnr_lr
+
+                avg_psnr_lr = avg_psnr_lr / idx
+                avg_biterr = avg_biterr / idx
+
+                logger.info('# Validation # PSNR_Stego: {:.4e}, Bit_acc: {: .4e}'.format(avg_psnr_lr, avg_biterr))
+                logger_val = logging.getLogger('val')  # validation logger
+                logger_val.info('<epoch:{:3d}, iter:{:8,d}> PSNR_Stego: {:.4e}, Bit_acc: {: .4e}'.format(
+                    epoch, current_step, avg_psnr_lr, avg_biterr))
+                # tensorboard logger
+                if opt['use_tb_logger'] and 'debug' not in opt['name']:
+                    tb_logger.add_scalar('psnr', avg_psnr, current_step)
+
+            # save models and training states
+            if current_step % opt['logger']['save_checkpoint_freq'] == 0:
+                if rank <= 0:
+                    logger.info('Saving models and training states.')
+                    model.save(current_step)
+                    model.save_training_state(epoch, current_step)
+
+    if rank <= 0:
+        logger.info('Saving the final model.')
+        model.save('latest')
+        logger.info('End of training.')
+
+
+if __name__ == '__main__':
+    main()