Git Product home page Git Product logo

slbr-visible-watermark-removal's Introduction

Visible Watermark Removal via Self-calibrated Localization and Background Refinement


Introduction

This is the official code of the following paper:

Visible Watermark Removal via Self-calibrated Localization and Background Refinement[1]
Jing Liang1, Li Niu1, Fengjun Guo2, Teng Long2 and Liqing Zhang1
1MoE Key Lab of Artificial Intelligence, Shanghai Jiao Tong University
2INTSIG
(ACM MM 2021 | Bibtex)

SLBR Network

Here is our proposed SLBR(Self-calibrated Localization and Background Refinement). Top row depicts the whole framework of SLBR and bottom row elaborates the details of our proposed three modules.

Some examples of inharmonious region
Some examples of inharmonious region

Quick Start

Install

Data Preparation

In this paper, we conduct all of the experiments on the latest released dataset CLWD[2] and LVW[3]. You can contact the authors of LVW to obtain the dataset.

Train and Test

  • How to train and test my model?

    We provide an example of training and a test bash respectively:scripts/train.sh, scripts/test.sh

    Please specify the checkpoint save path in --checkpoint and dataset path in--dataset_dir.

  • How to test on my data?

    We also provide an example of a custom data test bash: scripts/test_custom.sh And you can further tailor test_custom.py to meet your demands. For the best performance, it is better to finetune on your dataset since our training data size is set as 256x256.

Pretrained Model

Here is the model trained on CLWD dataset:

Visualization Results

We also show some qualitative comparision with state-of-art methods:

Some examples of inharmonious region

Acknowledgements

Part of the code is based upon the previous work SplitNet[4].

Citation

If you find this work or code is helpful in your research, please cite:

@inproceedings{liang2021visible,
  title={Visible Watermark Removal via Self-calibrated Localization and Background Refinement},
  author={Liang, Jing and Niu, Li and Guo, Fengjun and Long, Teng and Zhang, Liqing},
  booktitle={Proceedings of the 29th ACM International Conference on Multimedia},
  pages={4426--4434},
  year={2021}
}

Resources

We have summarized the existing papers, codes, and datasets on visible watermark removal in the following repository: https://github.com/bcmi/Awesome-Visible-Watermark-Removal

Reference

[1] Jing Liang, Li Niu, Fengjun Guo, Teng Long and Liqing Zhang. 2021. Visible Watermark Removal via Self-calibrated Localization and Background Refinement. In Proceedings of the 29th ACM International Conference on Multimedia. download

[2] Liu, Yang and Zhu, Zhen and Bai, Xiang. 2021. WDNet: Watermark-Decomposition Network for Visible Watermark Removal. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision.

[3] Danni Cheng, Xiang Li, Wei-Hong Li, Chan Lu, Fake Li, Hua Zhao, and WeiShi Zheng. 2018. Large-scale visible watermark detection and removal with deep convolutional networks. In Chinese Conference on Pattern Recognition and Computer Vision. 27–40.

[4] Xiaodong Cun and Chi-Man Pun. 2020. Split then Refine: Stacked Attentionguided ResUNets for Blind Single Image Visible Watermark Removal. arXiv preprint arXiv:2012.07007 (2020).

slbr-visible-watermark-removal's People

Contributors

jimleungjing avatar ustcnewly avatar

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

slbr-visible-watermark-removal's Issues

关于训练效果以及调参方面的问题

您好!非常感谢您优秀的工作!这里我有几个问题想要请教您:
1.我在训练模型的时候发现水印预测效果很差,不知道从哪里入手去改善模型效果;
2.我发现我训练模型的时候40个epoch的效果最好,甚至比90个epoch的效果还要好,应当是过拟合了,正通过调参等操作来实现改变,请问您有比较好的建议嘛?
3.测试的时候我发现,psnr值一直没有显著提升,仅仅达到了19左右
4.我看到您在回答其他人提出的问题时涉及到了beta参数,请问在代码中具体指的是哪个参数呢?
Let's suppose that watermarked image is J, original image is I, wateramrk alpha map is A, watermark is W and controllable parameter β∈[0,1].
我在代码中有看到存在两个β参数分别是:
parser.add_argument('--beta1', default=0.5, type=float, help='initial learning rate')
parser.add_argument('--beta2', default=0.999, type=float, help='initial learning rate')
以上是两个默认参数,其中的β1我改成0.9,做了实验,但是很遗憾效果并不好,所以想请问您是我理解错了嘛还是这个β参数另有?
非常期待您的回复!感谢您的帮助!

运行报错KeyError: 'dhn'

运行的时候报错,是原因呢。
Traceback (most recent call last):
File "D:/kuaipan/caiji/DL_github/SLBR-Visible-Watermark-Removal-master/test.py", line 180, in
main(parser.parse_args())
File "D:/kuaipan/caiji/DL_github/SLBR-Visible-Watermark-Removal-master/test.py", line 85, in main
Machine = models.dict[args.models](datasets=data_loaders, args=args)
File "D:\kuaipan\caiji\DL_github\SLBR-Visible-Watermark-Removal-master\src\models_init_.py", line 7, in basic
return BasicModel(**kwargs)
File "D:\kuaipan\caiji\DL_github\SLBR-Visible-Watermark-Removal-master\src\models\BasicModel.py", line 35, in init
self.model = nets.dictself.args.nets
KeyError: 'dhn'
==> creating model

进程已结束,退出代码1

Problem with learning rate equal to 0.1

Hello! Thank you for your work, it's a great contribution.

I am experiencing an unexpected behaviour, which according to my experiments seems to be caused by setting the learning rate on 0.1. I am training on a subset of 200 examples from the CLWD dataset.

I have a training loop copied from your train.py, I set the lr as 0.1 on the args. If I execute the training, I see the loss L1 going down but when I retrieved the model and do inference on an example from the training set, I get an output that matches exactly the input. How can that be? It’s as if the model wasn’t being applied.

CODE:

  • ARGS
CROP_SIZE = 384
DATASET="CLWD"
DATASET_DIR = "media/mini-CLWD-train-equal-test"
# DATASET_DIR = "media/10CLWD-train-equal-test"
TEST_DIR = DATASET_DIR
CHECKPOINT_DIR = "checkpoint"   
RESUME_DIR = ""

ARGS = dict(checkpoint=CHECKPOINT_DIR, 
            crop_size=CROP_SIZE, 
            dataset='clwd', 
            dataset_dir=DATASET_DIR, debug=False, 
            epochs=21,
            evaluate=False,  
            freq=-1,
            lr=0.1,
            schedule=[100,200,300], 
            gamma=0.1,        
            sigma_decay=0,   
            resume=RESUME_DIR,
            start_epoch=0,
            nets='slbr',
            test_dir=TEST_DIR,
            train_batch=6,
            dlr=0.001,
            data='', data_augumentation=False,
            finetune='', flip=False,
            alpha=0.5, beta1=0.9, beta2=0.999, bg_mode='res_mask',
            gan_norm=False, gpu=True, gpu_id='0', hl=False,
            input_size=256, k_center=2, k_refine=3, k_skip_stage=3,
            lambda_content=0, lambda_iou=0, lambda_l1=4, lambda_mask=1,
            lambda_primary=0.01, lambda_style=0, loss_type='l2', 
            mask_mode='res', masked=False, momentum=0,
            name='slbr_v1', no_flip=True, normalized_input=False,
            preprocess='resize', project_mode='simple', requires_grad=False,
            res=False,
            sim_metric='cos', sltype='vggx', test_batch=1,  
            use_refine=True, weight_decay=0, workers=1)

class DictArgs(object):
    def __init__(self, d):
        for k, v in d.items():
            setattr(self, k, v)

ARGS = DictArgs(ARGS)
  • TRAINING LOOP
from __future__ import print_function, absolute_import

import argparse
import torch,time,os

torch.backends.cudnn.benchmark = True

from src.utils.misc import save_checkpoint, adjust_learning_rate
import src.models as models

import datasets as datasets
from options import Options
import numpy as np

def train(args):
    args.seed = 1
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    MODEL_NAME = "SLBR"
    
    args.dataset = args.dataset.lower()
    if args.dataset == 'clwd':
        dataset_func = datasets.CLWDDataset
    elif args.dataset == 'lvw':
        dataset_func = datasets.LVWDataset
    else:
        raise ValueError("Not known dataset:\t{}".format(args.dataset))

    train_loader = torch.utils.data.DataLoader(dataset_func('train',args),batch_size=args.train_batch, shuffle=True,
        num_workers=args.workers, pin_memory=True)
    
    val_loader = torch.utils.data.DataLoader(dataset_func('val',args),batch_size=args.test_batch, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    lr = args.lr
    data_loaders = (train_loader,val_loader)

#     model = models.__dict__[args.models](datasets=data_loaders, args=args)
    model = models.__dict__[MODEL_NAME](datasets=data_loaders, args=ARGS)
    print('============================ Initization Finish && Training Start =============================================')
    print(f"It will start in epoch: {model.args.start_epoch}\nIt will end in epoch: {model.args.epochs}")
    for epoch in range(model.args.start_epoch, model.args.epochs):
        lr = adjust_learning_rate(data_loaders, model, epoch, lr, args)
            
        print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))

        model.record('lr',lr, epoch)        
        model.train(epoch)
        
        
        # save model
        save_epochs = {1,10,20,50,100,200,300,400}
        if epoch in save_epochs:
            model.validate(epoch)
            model.flush() #algo que tiene que ver con memoria
            
            print(f"Saving checkpoint of epoch {epoch}")
            model.save_checkpoint(filename=f"checkpoint{epoch}.pth.tar")
            
        
        # model.validate(epoch)
#         if args.freq < 0:
#             model.validate(epoch)
#             model.flush()
#             model.save_checkpoint()
            
    return model

final_model = train(ARGS)

I know the lr is multiplied by the gamma factor (so scaled down x10) on the epochs the schedule indicates. I've tried it with and without adjusting the learning rate during training. In the code I give the lr is not adjusted since the first adjust would come on the 100th epoch and I run it just 21 epochs.

  • INFERENCE CODE
    For inference I based on what's avaible on test_custom.py, using the helper functions defined there
# helper functions
### setup 
import os
import time
from pathlib import Path

import torch
import torch.nn.functional as F
import cv2
import numpy as np

import datasets as datasets
import src.models as models

torch.backends.cudnn.benchmark = True

# TEST_DIR = "test_imgs/"
# MODEL = "SLBR"
# CROP_SIZE = 384
# ONLY_PRED = True


#helper functions
def tensor2np(x, isMask=False):
    if isMask:
        if x.shape[1] == 1:
            x = x.repeat(1,3,1,1)
        x = ((x.cpu().detach()))*255
    else:
        x = x.cpu().detach()
        mean = 0
        std = 1
        x = (x * std + mean)*255
    return x.numpy().transpose(0,2,3,1).astype(np.uint8)


def save_output(inputs, preds, save_dir, img_fn, extra_infos=None,  verbose=False, alpha=0.5, only_pred=True):
    outs = []
    image = inputs['I'] #, inputs['bg'], inputs['mask']
    image = cv2.cvtColor(tensor2np(image)[0], cv2.COLOR_RGB2BGR)
    bg_pred,mask_preds = preds['bg'], preds['mask']
    bg_pred = cv2.cvtColor(tensor2np(bg_pred)[0], cv2.COLOR_RGB2BGR)
    mask_pred = tensor2np(mask_preds, isMask=True)[0]
    outs = [image, bg_pred,  mask_pred]
    if only_pred:
        outimg = bg_pred
    else:
        outimg = np.concatenate(outs, axis=1)
    if verbose==True:
        cv2.imshow("out",outimg)
        cv2.waitKey(0)
    else:
        img_fn = os.path.split(img_fn)[-1]
        out_fn = os.path.join(save_dir, "{}{}".format(os.path.splitext(img_fn)[0], os.path.splitext(img_fn)[1]))
        cv2.imwrite(out_fn, outimg)


def preprocess(file_path, img_size=512):
    img_J = cv2.imread(file_path)
    assert img_J is not None, "NoneType"
    h,w,_ = img_J.shape
    img_J = cv2.cvtColor(img_J, cv2.COLOR_BGR2RGB).astype(float)/255.
    img_J = torch.from_numpy(img_J.transpose(2,0,1)[np.newaxis,...]) #[1,C,H,W]
    img_J = F.interpolate(img_J, size=(img_size, img_size), mode='bilinear')
    return img_J, (h, w)


def test_dataloader(folder, crop_size):
    files = [x.as_posix() for pat in ["*.jpg", "*.jpeg", "*.png"] for x in Path(folder).glob(pat) ]
    loaders = []
    save_fns = []
    orig_dims = []
    for file in files:
        processed, orig_size = preprocess(file, img_size=crop_size)
        loaders.append(processed)
        save_fns.append(file)
        orig_dims.append(orig_size)
    return loaders, save_fns, orig_dims

import matplotlib.pyplot as plt
def sanity_check(checkpoint_path: str, image_index: int=0):
    MODEL="SLBR"

    CROP_SIZE = 384
    DATASET="CLWD"
    DATASET_DIR = "media/mini-CLWD-train-equal-test"
    TEST_DIR = DATASET_DIR
    CHECKPOINT_DIR = "checkpoint"   
    RESUME_DIR = checkpoint_path

    ARGS = dict(checkpoint=CHECKPOINT_DIR, 
                crop_size=CROP_SIZE, 
                dataset='clwd', 
                dataset_dir=DATASET_DIR, debug=False, 
                dlr=0.001,
                epochs=100,
                evaluate=True, 
                finetune='', flip=False, 
                freq=-1,
                lr=0.01,
                resume=RESUME_DIR,
                start_epoch=0,
                nets='slbr',
                test_dir=TEST_DIR,
                train_batch=6,
                data='', data_augumentation=False, 
                alpha=0.5, beta1=0.9, beta2=0.999, bg_mode='res_mask',
                gamma=0.1, gan_norm=False, gpu=True, gpu_id='0', hl=False,
                input_size=256, k_center=2, k_refine=3, k_skip_stage=3,
                lambda_content=0, lambda_iou=0, lambda_l1=4, lambda_mask=1,
                lambda_primary=0.01, lambda_style=0, loss_type='l2', 
                mask_mode='res', masked=False, models='slbr', momentum=0,
                name='slbr_v1', no_flip=True, normalized_input=False,
                preprocess='resize', project_mode='simple', requires_grad=False,
                res=False,
                schedule=[5, 10], sigma_decay=0, sim_metric='cos', sltype='vggx', test_batch=1,  
                use_refine=True, weight_decay=0, workers=2)

    class DictArgs(object):
        def __init__(self, d):
            for k, v in d.items():
                setattr(self, k, v)

    ARGS = DictArgs(ARGS)

    model = models.__dict__[MODEL](datasets=(None, None), args=ARGS)
    
    
    ### inference and comparison
    TEST_DIR = "media/mini-CLWD-train-equal-test/train/Watermarked_image"
    doc_loader, fns, orig_dims = test_dataloader(TEST_DIR, ARGS.crop_size)
    
    input_example = doc_loader[image_index]
    
    input_example = input_example.to(model.device).float()
    imoutput, immask_all, imwatermark = model.model(input_example)
    imoutput = imoutput[0]  
    immask = immask_all[0]
    imfinal = imoutput * immask + model.norm(input_example) * (1-immask)
    
    print(f"Boolean comparison of tensors (torch.equal): {torch.equal(input_example, imfinal)}")

    input_example_p = tensor2np(F.interpolate(input_example, size=orig_dims[10], mode="bilinear"))[0]
    imfinal_p = tensor2np(F.interpolate(imfinal, size=orig_dims[10], mode="bilinear"))[0]
    
    print(f"Boolean comparison of processed numpy objects showed (np.array_equal): {np.array_equal(imfinal_p,input_example_p)}")
    fig, axes = plt.subplots(1,2)
    axes[0].imshow(input_example_p)
    axes[1].imshow(imfinal_p)

sanity_check("checkpoint/slbr_v1/checkpoint20.pth.tar", image_index=0)

Notice that I pass the checkpoint of the model trained on the resume_path variable. I do inference on a given example of the train set, I check if the matrix are equal and also check the image.

  • OUTPUT

Screen Shot 2023-04-14 at 12 34 06

Images are the same and they are equal, entry by entry. I've tried other saved checkpoints as well. What could be going on?

Thanks for your support.

CLWD测试集中没有Watermark文件夹

作者您好,我在运行测试脚本script/test.sh时,使用CLWD测试集,遇到以下错误:
/tmp/pip-build-c8_36uuo/opencv-python-headless/opencv/modules/imgcodecs/src/loadsave.cpp (239) findDecoder imread_('/home/ubuntu/sda/datasets/Watermark/CLWD/test/Watermark/2235.png'): can't open/read file: check file path/integrity
查看后发现确实缺少sda/datasets/Watermark/CLWD/test/Watermark文件夹,只有
Mask
Watermarked_image
Watermark_free_image
这三个文件夹,麻烦问下是需要生成Watermark文件夹吗?有具体的脚本吗?

About BVMR

您好,请问您能提供一下关于BVMR方法在LVW和CLWD数据集上的测试结果吗?

请教一个环境问题

初入这个领域,不懂的地方是在太多,现在环境倒是安好了,测试了好几个python 版本,目前来看3.7好像比较适合。现在卡在我们依赖的库太旧了。torch==1.6.0 我看官方已经不维护了,改为1.7 ,又出现
ERROR: torch has an invalid wheel, torch has an invalid wheel, .dist-info directory not found 这样的问题。
不只是否有时间否更新一下依赖库?

AUTOMATIC MIXED PRECISION - BCELOSS

Hey!
Since the model is heavy, I was trying to apply Mixed Precision to reduce GPU usage.
But after applying it I get the following error:

RuntimeError: torch.nn.functional.binary_cross_entropy and torch.nn.BCELoss are unsafe to autocast.
Many models use a sigmoid layer right before the binary cross entropy layer.
In this case, combine the two layers using torch.nn.functional.binary_cross_entropy_with_logits
or torch.nn.BCEWithLogitsLoss.  binary_cross_entropy_with_logits and BCEWithLogits are
safe to autocast.

This error happens due to this :

File "/home/ubuntu/SLBR-Visible-Watermark-Removal/src/models/SLBR.py", line 136, in train
    coarse_loss, refine_loss, style_loss, mask_loss = self.loss(
  File "/opt/conda/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/SLBR-Visible-Watermark-Removal/src/models/SLBR.py", line 74, in forward
    final_mask_loss += self.mask_loss(pred_ms[0], mask)

BCELoss is used here:

self.masked_l1_loss, self.mask_loss = l1_relative, nn.BCELoss()

Since this can be fixed by using torch.nn.BCEWithLogitsLoss. Can you please guide me from where should I remove the sigmoid so that I can use torch.nn.BCEWithLogitsLoss here

Question about the cost of training

Thanks for your impressive work. I have some question about the cost of training.
I used 4 Titan Xp, batchsize is set as 8, train on CLWD dataset. It seems to take 150 hours to complete the training, namely 600 GPU hour. What is your hardware environment during training?
Or maybe I got something wrong.
Looking forward to your reply.

Can't find a pretrained model

Hi Li. Thanks for your work, but I can't find a pretrained model in the code. Can you provide one? Besides, is there a demo script for custom images?

Dataset Source

Hey,

I have contacted the authors of LVW for the dataset, but he has not responded till now.
Can you please provide me with the dataset? My task is pending and it would be really great if you can provide me with the dataset.

Thank you!

Error when training with just 10 examples

Hello! Thank you for your work. It's a great contribution.

I’m training the model based on what you got on train.py. I’m doing it on a small subset of the CLWD dataset, because I’m trying to overfit the model (it’s just a check of consistency, I want to see that the model learns).

If I use a small sample, 10 examples (so 10 images on each of the dedicated folder: Watermarked_images, Watermark_free_images, etc) and run the training loop I get an error when the batching occurs. It’s an OpenCV error which seems to be related to the train_loader generator.

If I use 200 examples (as I'm using in #31) , I don’t get the problem.

CODE

  • ARGS
CROP_SIZE = 384
DATASET="CLWD"
# DATASET_DIR = "media/mini-CLWD-train-equal-test"
DATASET_DIR = "media/10CLWD-train-equal-test"
TEST_DIR = DATASET_DIR
CHECKPOINT_DIR = "checkpoint"   
RESUME_DIR = ""

ARGS = dict(checkpoint=CHECKPOINT_DIR, 
            crop_size=CROP_SIZE, 
            dataset='clwd', 
            dataset_dir=DATASET_DIR, debug=False, 
            epochs=21,
            evaluate=False,  
            freq=-1,
            lr=0.01,
            schedule=[100,200,300], 
            gamma=0.1,        
            sigma_decay=0,    
            resume=RESUME_DIR,
            start_epoch=0,
            nets='slbr',
            test_dir=TEST_DIR,
            train_batch=2,
            dlr=0.001,
            data='', data_augumentation=False,
            finetune='', flip=False,
            alpha=0.5, beta1=0.9, beta2=0.999, bg_mode='res_mask',
            gan_norm=False, gpu=True, gpu_id='0', hl=False,
            input_size=256, k_center=2, k_refine=3, k_skip_stage=3,
            lambda_content=0, lambda_iou=0, lambda_l1=4, lambda_mask=1,
            lambda_primary=0.01, lambda_style=0, loss_type='l2', 
            mask_mode='res', masked=False, momentum=0,
            name='slbr_v1', no_flip=True, normalized_input=False,
            preprocess='resize', project_mode='simple', requires_grad=False,
            res=False,
            sim_metric='cos', sltype='vggx', test_batch=1,  
            use_refine=True, weight_decay=0, workers=1)

class DictArgs(object):
    def __init__(self, d):
        for k, v in d.items():
            setattr(self, k, v)

ARGS = DictArgs(ARGS)

In media/10CLWD-train-equal-test I got a copy of the CLWD folder structure with just 10 examples, also with the test and train folders being equal. I'm not interested in testing performance on unseen data at this point, the test folder is there just so the code in the repository doesn't break.

  • TRAINING LOOP
from __future__ import print_function, absolute_import

import argparse
import torch,time,os

torch.backends.cudnn.benchmark = True

from src.utils.misc import save_checkpoint, adjust_learning_rate
import src.models as models

import datasets as datasets
from options import Options
import numpy as np

def train(args):
    args.seed = 1
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    MODEL_NAME = "SLBR"
    
    args.dataset = args.dataset.lower()
    if args.dataset == 'clwd':
        dataset_func = datasets.CLWDDataset
    elif args.dataset == 'lvw':
        dataset_func = datasets.LVWDataset
    else:
        raise ValueError("Not known dataset:\t{}".format(args.dataset))

    train_loader = torch.utils.data.DataLoader(dataset_func('train',args),batch_size=args.train_batch, shuffle=True,
        num_workers=args.workers, pin_memory=True)
    
    val_loader = torch.utils.data.DataLoader(dataset_func('val',args),batch_size=args.test_batch, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    lr = args.lr
    data_loaders = (train_loader,val_loader)

#     model = models.__dict__[args.models](datasets=data_loaders, args=args)
    model = models.__dict__[MODEL_NAME](datasets=data_loaders, args=ARGS)
    print('============================ Initization Finish && Training Start =============================================')
    print(f"It will start in epoch: {model.args.start_epoch}\nIt will end in epoch: {model.args.epochs}")
    for epoch in range(model.args.start_epoch, model.args.epochs):
        lr = adjust_learning_rate(data_loaders, model, epoch, lr, args)
            
        print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))

        model.record('lr',lr, epoch)        
        model.train(epoch)
        
        
        # save model
        save_epochs = {1,10,20,50,100,200,300,400}
        if epoch in save_epochs:
            model.validate(epoch)
            model.flush() 
            
            print(f"Saving checkpoint of epoch {epoch}")
            model.save_checkpoint(filename=f"checkpoint{epoch}.pth.tar")
            
        
        # model.validate(epoch)
#         if args.freq < 0:
#             model.validate(epoch)
#             model.flush()
#             model.save_checkpoint()
            
    return model

final_model = train(ARGS)

I get the following output with the error as soon as training starts

==> creating model 
==> creating model [Finish]
==> Total params: 21.39M
==> Total devices: 1
==> Current Checkpoint: checkpoint/slbr_v1
============================ Initization Finish && Training Start =============================================
It will start in epoch: 0
It will end in epoch: 21

Epoch: 1 | LR: 0.01000000
---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
/tmp/ipykernel_32332/965729503.py in <module>
----> 1 final_model = train(ARGS)

/tmp/ipykernel_32332/1760790972.py in train(args)
     46 
     47         model.record('lr',lr, epoch)
---> 48         model.train(epoch)
     49 
     50 

/alloc/data/fury_watermark-remover-fda/discovery/SLBR-Visible-Watermark-Removal-master/src/models/SLBR.py in train(self, epoch)
    104         end = time.time()
    105         bar = Bar('Processing {} '.format(self.args.nets), max=len(self.train_loader))
--> 106         for i, batches in enumerate(self.train_loader):
    107             current_index = len(self.train_loader) * epoch + i
    108 

/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py in __next__(self)
    433         if self._sampler_iter is None:
    434             self._reset()
--> 435         data = self._next_data()
    436         self._num_yielded += 1
    437         if self._dataset_kind == _DatasetKind.Iterable and \

/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py in _next_data(self)
   1083             else:
   1084                 del self._task_info[idx]
-> 1085                 return self._process_data(data)
   1086 
   1087     def _try_put_index(self):

/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py in _process_data(self, data)
   1109         self._try_put_index()
   1110         if isinstance(data, ExceptionWrapper):
-> 1111             data.reraise()
   1112         return data
   1113 

/usr/local/lib/python3.7/site-packages/torch/_utils.py in reraise(self)
    426             # have message field
    427             raise self.exc_type(message=msg)
--> 428         raise self.exc_type(msg)
    429 
    430 

error: Caught error in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 198, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/alloc/data/fury_watermark-remover-fda/discovery/SLBR-Visible-Watermark-Removal-master/datasets/clwd_dataset.py", line 82, in __getitem__
    sample = self.get_sample(index)
  File "/alloc/data/fury_watermark-remover-fda/discovery/SLBR-Visible-Watermark-Removal-master/datasets/clwd_dataset.py", line 63, in get_sample
    img_J = cv2.cvtColor(img_J, cv2.COLOR_BGR2RGB)
cv2.error: OpenCV(4.2.0) /io/opencv/modules/imgproc/src/color.cpp:182: error: (-215:Assertion failed) !_src.empty() in function 'cvtColor'

I've narrowed it down to a problem on the batching, as the stack trace suggests. I tried a minimal example just initialising the train_loader and iterating over the batches. The same happens. I print the path being processed.

args = ARGS
train_loader = torch.utils.data.DataLoader(datasets.CLWDDataset("train", args),batch_size=args.train_batch, shuffle=True,
                                            num_workers=args.workers, pin_memory=True)

for i, batch in enumerate(train_loader):
    print(i,batch["img_path"],"\n")

And get this output with the error mentioned

0 ['media/10CLWD-train-equal-test/train/Watermarked_image/8.jpg', 'media/10CLWD-train-equal-test/train/Watermarked_image/10.jpg'] 

1 ['media/10CLWD-train-equal-test/train/Watermarked_image/5.jpg', 'media/10CLWD-train-equal-test/train/Watermarked_image/3.jpg'] 

2 ['media/10CLWD-train-equal-test/train/Watermarked_image/1.jpg', 'media/10CLWD-train-equal-test/train/Watermarked_image/6.jpg'] 

3 ['media/10CLWD-train-equal-test/train/Watermarked_image/4.jpg', 'media/10CLWD-train-equal-test/train/Watermarked_image/2.jpg'] 

---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
/tmp/ipykernel_32332/1400065821.py in <module>
      3                                             num_workers=args.workers, pin_memory=True)
      4 
----> 5 for i, batch in enumerate(train_loader):
      6     print(i,batch["img_path"],"\n")

/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py in __next__(self)
    433         if self._sampler_iter is None:
    434             self._reset()
--> 435         data = self._next_data()
    436         self._num_yielded += 1
    437         if self._dataset_kind == _DatasetKind.Iterable and \

/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py in _next_data(self)
   1083             else:
   1084                 del self._task_info[idx]
-> 1085                 return self._process_data(data)
   1086 
   1087     def _try_put_index(self):

/usr/local/lib/python3.7/site-packages/torch/utils/data/dataloader.py in _process_data(self, data)
   1109         self._try_put_index()
   1110         if isinstance(data, ExceptionWrapper):
-> 1111             data.reraise()
   1112         return data
   1113 

/usr/local/lib/python3.7/site-packages/torch/_utils.py in reraise(self)
    426             # have message field
    427             raise self.exc_type(message=msg)
--> 428         raise self.exc_type(msg)
    429 
    430 

error: Caught error in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 198, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/alloc/data/fury_watermark-remover-fda/discovery/SLBR-Visible-Watermark-Removal-master/datasets/clwd_dataset.py", line 82, in __getitem__
    sample = self.get_sample(index)
  File "/alloc/data/fury_watermark-remover-fda/discovery/SLBR-Visible-Watermark-Removal-master/datasets/clwd_dataset.py", line 63, in get_sample
    img_J = cv2.cvtColor(img_J, cv2.COLOR_BGR2RGB)
cv2.error: OpenCV(4.2.0) /io/opencv/modules/imgproc/src/color.cpp:182: error: (-215:Assertion failed) !_src.empty() in function 'cvtColor'

Some images do get processed, but it fails with the last batch. The images missing on this example (7.jpg and 9.jpg) are on the folder, double checked. Have you run into anything similar?

Thanks for the support

Question about the optimizer

Hi, thanks for your excellent work.

I am puzzled about the settings of optimizers as follows:
`def set_optimizers(self):
self.optimizer_encoder = torch.optim.Adam(self.encoder.parameters(), lr=self.args.lr)
self.optimizer_image = torch.optim.Adam(self.coarse_decoder.parameters(), lr=self.args.lr)

    if self.refinement is not None:
        self.optimizer_refine = torch.optim.Adam(self.refinement.parameters(), lr=self.args.lr)
    
    if self.shared != 0:
        self.optimizer_shared = torch.optim.Adam(self.shared_decoder.parameters(), lr=self.args.lr)`

As can be seen, different optimizers are set to different components. However, the learning rates are the same. I wonder is there any difference to set only one optimizer to the parameters of SLBR model.

Thank you very much and looking forward your replies.

Training on custom dataset

Hi author, I want to train your SLBR model from scratch on my custom dataset, I am not able to understand how long I should train it, and how much training data is required, what much cuda it will consume during training with batch size 8?

confusion about output add input in training

reconstructed_image = (reconstructed_image + synthesized).clamp(0,1)
refine_bg = (torch.tanh(refine_bg) + synthesized).clamp(0,1)

About above code,Idon't understand why add the input image,
it is different about train and infer,i worry the output is different

请教一下

作者您好,我想请教一下要训练高分辨率去水印的话需要修改哪些地方,train.sh里面的INPUT_SIZE吗,网络模型能支持多大分辨率的,我看您训练与测试都是256的图像,我想做640、1024大小的去水印

weight file problem

Hello, I have carefully read your wonderful and unparalleled article. But while running the article code, I found that your weight file on the cloud drive is damaged. Could you please update the weight file on the cloud drive?

测试推理权重

有给出的模型权重,项目主页给的都是空的活着报错,无法下载

Freezing layers of the model

Hello.

I would like to experiment with freezing some weights of the model. I want to keep the core knowledge learned but still give it room to learn. What weights should be frozen?

If this were a simple Sequential model or a large pretrained model to be used on a downstream task, I would just freeze all the weights except the ones from the head (if the task was a classification, for instance). Is there an equivalent here?

This is the output I get from summary function from torch-summary library, if it helps to see it fast.

===========================================================================
Layer (type:depth-idx)                             Param #
===========================================================================
├─CoarseEncoder: 1-1                               --
|    └─ModuleList: 2-1                             --
|    |    └─DownConv: 3-1                          28,896
|    |    └─DownConv: 3-2                          129,792
|    |    └─DownConv: 3-3                          517,632
├─SharedBottleNeck: 1-2                            --
|    └─ModuleList: 2-2                             --
|    |    └─DownConv: 3-4                          2,065,408
|    |    └─DownConv: 3-5                          8,259,584
|    └─ModuleList: 2-3                             --
|    |    └─UpConv: 3-6                            4,130,048
|    └─ModuleList: 2-4                             --
|    |    └─ECABlock: 3-7                          3
|    └─ModuleList: 2-5                             --
|    |    └─ECABlock: 3-8                          3
├─CoarseDecoder: 1-3                               --
|    └─Conv2d: 2-6                                 99
|    └─ModuleList: 2-7                             --
|    |    └─MBEBlock: 3-9                          970,147
|    |    └─MBEBlock: 3-10                         243,923
|    |    └─MBEBlock: 3-11                         61,675
|    └─ModuleList: 2-8                             --
|    |    └─ECABlock: 3-12                         3
|    |    └─ECABlock: 3-13                         3
|    |    └─ECABlock: 3-14                         3
|    └─ModuleList: 2-9                             --
|    |    └─SMRBlock: 3-15                         1,135,782
|    |    └─SMRBlock: 3-16                         284,758
|    |    └─SMRBlock: 3-17                         71,598
|    └─ModuleList: 2-10                            --
|    |    └─ECABlock: 3-18                         3
|    |    └─ECABlock: 3-19                         3
|    |    └─ECABlock: 3-20                         3
├─Refinement: 1-4                                  --
|    └─Sequential: 2-11                            --
|    |    └─Conv2d: 3-21                           1,184
|    |    └─InstanceNorm2d: 3-22                   --
|    |    └─LeakyReLU: 3-23                        --
|    └─ResDownNew: 2-12                            --
|    |    └─DownConv: 3-24                         36,992
|    └─ResDownNew: 2-13                            --
|    |    └─DownConv: 3-25                         129,280
|    └─ResDownNew: 2-14                            --
|    |    └─DownConv: 3-26                         516,608
|    └─Sequential: 2-15                            --
|    |    └─Conv2d: 3-27                           1,056
|    └─Sequential: 2-16                            --
|    |    └─Conv2d: 3-28                           2,080
|    |    └─LeakyReLU: 3-29                        --
|    |    └─Conv2d: 3-30                           9,248
|    |    └─LeakyReLU: 3-31                        --
|    └─Sequential: 2-17                            --
|    |    └─Conv2d: 3-32                           8,256
|    |    └─LeakyReLU: 3-33                        --
|    |    └─Conv2d: 3-34                           36,928
|    |    └─LeakyReLU: 3-35                        --
|    └─ModuleList: 2-18                            --
|    |    └─CFFBlock: 3-36                         895,104
|    |    └─CFFBlock: 3-37                         895,104
|    |    └─CFFBlock: 3-38                         895,104
|    └─Sequential: 2-19                            --
|    |    └─Conv2d: 3-39                           64,544
|    |    └─InstanceNorm2d: 3-40                   --
|    |    └─LeakyReLU: 3-41                        --
|    |    └─Conv2d: 3-42                           99
===========================================================================
Total params: 21,390,953
Trainable params: 21,390,953
Non-trainable params: 0

I've tried just finetunning the model without freezing, of course, but I've seen it degrading in some examples. As in, the checkpoint you provide removes the watermark but after training I see some parts of the watermark coming back.

Thanks for your help :)

size mismatch

RuntimeError: Error(s) in loading state_dict for SLBR:
size mismatch for shared_decoder.up_im_atts.0.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 256]).
size mismatch for shared_decoder.up_mask_atts.0.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 256]).
size mismatch for coarse_decoder.atts_bg.0.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 128]).
size mismatch for coarse_decoder.atts_bg.1.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 64]).
size mismatch for coarse_decoder.atts_bg.2.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 32]).
size mismatch for coarse_decoder.atts_mask.0.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 128]).
size mismatch for coarse_decoder.atts_mask.1.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 64]).
size mismatch for coarse_decoder.atts_mask.2.conv.weight: copying a param with shape torch.Size([1, 1, 3]) from checkpoint, the shape in current model is torch.Size([1, 1, 32]).

How to train a watermark removal model on two datasets simultaneously

Hello! Thank you very much for your tremendous contribution.
I noticed that your code is trained on two datasets separately. I would like to try to train a model that can run on both datasets simultaneously, but I encountered difficulties. As far as I know, it is possible to merge the two datasets, but due to issues with image size and format, it may not be feasible. Therefore, I would like to ask if you have a better solution? Looking forward to your reply!

自定义CLWD训练集

使用预训练模型去水印,发现没有透明度的水印,完全没有去除效果。打算自己加入图片到训练集中来训练模型,请问CLWD训练集打水印的图片和mask是怎样生成的呢,能否提供下代码呢

About the dataset

it seems that the dataset CLWD only includes one sub folder, however, the code to train need "train" and "test" folders, could you give me some suggestion on how to split the train set and test set? Thanks a lot !

test_custom

When I run test_custom.py file on my own images, I get different results in consequent executions which is so strange. Also, results are so bad. My question is are there any parameters that I must change to get better results?

Thank you!

Question about the mask in dataset.

Is there any influence if the masks in the dataset are not accurate?

For example, I detected the watermarked area with YOLO, so the mask was generated as white squares. If I put these kinds of masks in the training process, will the result still be good or not?

Inference and Training with Images of variable sizes

Thank you for the excellent work done on this repository.

I would like to train a model (SLBR) from scratch using images of variable sizes (1200x1200, 1200x800, 1000x1000 etc.), my goal is to be able to return the image with the watermark removed in the same quality of the original image.

Do you have any recommendations on how to achieve this?

How to use without mask?

For a task without mask (just has target, input(before watermark, after watermark))
How to use this code?

As "ground-truth watermark mask M", in your paper formula(1) is needed.

Any suggestion to no ground-truth watermark?

the watermark mask cannot get exactly, but the style is almost same, across img data.

Thx!

bad performance or wrong operation?

It's a pity to see such output, I don't know if there is a problem with the operation, please don't hesitate to enlighten me, thanks a lot
(with the code test_custom.py)
image

test_custom.py error

import scripts.datasets as datasets
import scripts.machines as machines

can not find “scripts.datasets” and “scripts.machines”

对模型设计的一个疑问

为什么是这样:

        reconstructed_image = torch.tanh(im)  # [-1, 1]
        if self.long_skip:
            reconstructed_image = (reconstructed_image + synthesized).clamp(0,1)

而不是这样:

       
        if self.long_skip:
            reconstructed_image = reconstructed_image + synthesized
            reconstructed_image = torch.sigmoid(reconstructed_image)

Inference Code

Hi, Could you please add inference code for this model, for us to try on? Thanks!

Color space changes

Hello again, thank you for your support.

I want to ask why are there color space changes when loading the data, for the CLWD dataset and when doing inference with test_custom.py

Does it matter to the model the color space used? In that case, which one should you use.

I notice this when using your checkpoint for doing inference on my images. I could work around it just changing back the color space after applying the model.

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.