一、赛题解读
1、赛题分析
赛题任务需要对添加了水印的图像，将水印擦除掉，还原原本的图的样子(图1)。

与手写文字擦除任务(图2)一个比较大的区别是：水印占据面积很大，因此对水印擦除后，还需要对被擦除的区域进行一个填补，这个是该项目的难点所在。

结论，本任务单纯使用语义分割效果不佳，需要使用带有生成能力的img2img式模型。

2、数据处理
数据处理的方式决定了模型的设计，也会对预测的精度产生较大的影响。
(1)、为了显示的引导模型进行预测，需要结合gt和img做差值来生成mask(如下图)。(参考Erasenet论文对比结果，带有预测mask的模型的psnr要普遍高于纯的img2img的)。从GoogLeNet也可以得到启示，添加了预测mask的分支可以更有效的实现梯度传递。
2)、由于本次比赛数据集过大，1841张本体图像，每张本体图生成551张带水印的图像，一共1841x551张，100多G。其实到后面就会发现，这个任务模型推理出mask的位置是比较简单的，因为mask是十分规律的，但是生成依然做的不够好，所以将数据集从1841x551削减到1841x20(不到10G)，使得可以在aistudio上就可以加载进行训练。

(3)、参考手写文字擦除，我们同样将图片进行裁剪(随机裁剪至512, 512大小)，对密集预测型任务不使用resize。
3、模型训练及预测
A榜用的是Erasenet，同手写文字擦除一样，我们更改了loss函数，因为这个方式比较直接效果显著(模型是需要训练的，并不是设计的越复杂越好，直接调整面向真实数据的loss设计可以有效改变模型训练的轨迹)。模型结构图如下：
预测时：先将图像重叠分块到512x512大小，对每个小块取值res = pre_image * mask + image * (1 – mask)。 也就是对模型预测为mask的地方取模型的输出，对预测为非mask的地方取输入图片的输出。这样在非mask的地方就可以保证像素差接近0（因为jpg图像本身有一些噪点，一般达不到0）。

B榜对模型进行了一次调优，方法是将网络最开始下采样和精修部分下采样的卷积替换成了SwinT模块，就像在我之前Swin那个项目里一样，将Swin和CNN成功的结合起来，做到又快又好，最终B榜分数也比较高。下图展示了原Erasenet和带swin的Erasenet改在验证集上的表现，psnr分别是31.418，33.042。
4、模型优化
(1)、损失函数调整： 在所有超参数的调整中，我们把损失函数放在首要位置，因为其直面数据集。

调整mask损失，不仅使用bce，也使用l1。 增加image_loss，该任务对生成的要求更高，而且mask十分规整，因此加大image_loss的权重，增加到1.5，其他的则相对的调整到0.5。最后将所有loss相加，psnr有了一个很大的提升。

(2)、优化器调整： 在所有超参数的调整中，我们把优化器模式放在第二位置，因为其决定着训练能否达到当前模型最优。

每次重新调用优化器，就相当于对模型加载了一
个经过预训练的模型。稳定后续训练，并且逃离局部平坦区域(梯度接近0)。
(3)、结构调整： 在所有超参数的调整中，我们把网络结构放在第三位置，因为其难调整，需要训练到平稳才能看出模型的好坏。

因为该任务，生成是比较困难的，尤其是彩色水印叠加在彩色图像上之后，虽然可以检测到mask，但是生成的效果不佳。

二、本项目使用说明
由于本次比赛，水印擦除挑战赛评分提交的是模型加模型参数文件，因此就没有再加入对A，B榜的图片进行推理的notebook。因此本项目只关心于如何训练模型。

首先运行下面cell的代码，对数据集进行解压。

In [ ]

解压文件

!unzip -oq data/data145795/train_dataset.zip -d ./dataset
!unzip -oq data/data145795/valid_dataset.zip -d ./dataset
1、数据处理
数据处理的方式决定了模型的设计，也会对预测的精度产生较大的影响。与手写文字擦除任务一个比较大的区别是：水印占据面积很大，因此对水印擦除后，还需要对被擦除的区域进行一个填补，这个是该项目的难点所在。为了显示的引导模型进行预测，需要结合gt和img做差值来生成mask。如下图，从左向右依次为img,gt,mask(用自己的代码生成的，参考generate_mask.py，代码中图片路径供参考，是在本地电脑进行处理的):
另一方面，由于本次比赛数据集过大，1841张本体图像，每张本体图生成551张带水印的图像，一共1841x551张，100多G。其实到后面就会发现，这个任务模型推理出mask的位置是比较简单的，因为mask是十分规律的，但是生成依然做的不够好，所以要扩充数据集最好是找到1841张本体图像的分布然后进行扩充。

虽然机器学习定理告诉我们，训练数据量越多模型效果越好，越不容易过拟合；但这是有前提的，因为我们无法做到全批量梯度下降，真实的训练过程我们只会一次一个小batch的训练，最早期的batch对模型的梯度影响必然会被后期的batch洗掉一部分，反向传播决定了模型不能进行增量学习。所以，在显存不大的情况下，过大训练数据集起到的作用得不偿失，将数据集控制在20G之内既加快了项目打开的速度，也不会掉精度。

参考手写文字擦除，我们同样将图片进行裁剪(随机裁剪至512, 512大小)，对密集预测型任务不使用resize。

总结一下：在数据处理部分，我们一共使用了三种策略， 1、缩减数据集100G–>10G 2、生成mask引导模型训练 3、随机裁剪至512x512大小

2、模型搭建
A榜用的是Erasenet，模型代码参考了https://aistudio.baidu.com/aistudio/projectdetail/3439691 ， 同手写文字擦除一样，我们更改了loss函数，因为这个方式比较直接效果显著(模型是需要训练的，并不是设计的越复杂越好，直接调整面向真实数据的loss设计可以有效改变模型训练的轨迹)。模型结构图如下：模型数据流向大体如上，loss的地方做了一定的修改。

B榜对模型进行了一次调优，方法是将网络最开始下采样和精修部分下采样的卷积替换成了SwinT模块，就像在我之前Swin那个项目里一样，将Swin和CNN成功的结合起来，做到又快又好，最终B榜分数也比较高。下图展示了原Erasenet和带swin的Erasenet改在验证集上的表现，psnr分别是31.418，33.042。
再次印证了SwinT单个模块的强大力量！
3、训练模型
运行trainstr.ipynb可以训练原始erasenet，训练日志log和最好的模型都已包含在项目中，用visualdl即可可视化。虽然最后不会用这个模型提交，但还是放在这，可以起到一个参考的作用，因为batchsize达到28，所以训练起来是要比erasenet改快一点的。

erasenet改是分两部分训练的，开始是用的A100，运行trainswin.ipynb即可，但是A100只能训练24小时，因此将最好的模型加载再使用V100进行训练，运行trainswinv100即可。log_swin，log_swin_v100包含了完整的训练日志。我们只是因为时间紧迫才用的A100训练的，但这并不是必要的，单纯用V100多训练几天也是可以的。

In [ ]
import warnings
warnings.filterwarnings(“ignore”)

进行训练

from visualdl import LogWriter
import os
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.io import DataLoader
from dataset.data_loader import TrainDataSet, ValidDataSet
from loss.Loss import LossWithGAN_STE, LossWithSwin
from models.swin_gan import STRnet2_change
import utils
import random
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import math
%matplotlib inline

log = LogWriter(‘log_swin_v100’)
def psnr(img1, img2):
mse = np.mean((img1/1.0 - img2/1.0) ** 2 )
if mse < 1.0e-10:
return 100
return 10 * math.log10(255.0**2/mse)

训练配置字典

CONFIG = {
‘numOfWorkers’: 0,
‘modelsSavePath’: ‘train_models_swin_v100’,
‘batchSize’: 10,
‘traindataRoot’: ‘dataset/dataset’,
‘validdataRoot’: ‘dataset/valid_dataset’,
‘pretrained’: ‘train_models_swin/STE_15_43.2223.pdparams’,
‘num_epochs’: 100,
‘net’: ‘str’,
‘lr’: 1e-4,
‘lr_decay_iters’: 40000,
‘gamma’: 0.5,
‘seed’: 9420
}

设置gpu

if paddle.is_compiled_with_cuda():
paddle.set_device(‘gpu:0’)
else:
paddle.set_device(‘cpu’)

设置随机种子

random.seed(CONFIG[‘seed’])
np.random.seed(CONFIG[‘seed’])
paddle.seed(CONFIG[‘seed’])

noinspection PyProtectedMember

paddle.framework.random._manual_program_seed(CONFIG[‘seed’])

batchSize = CONFIG[‘batchSize’]
if not os.path.exists(CONFIG[‘modelsSavePath’]):
os.makedirs(CONFIG[‘modelsSavePath’])

traindataRoot = CONFIG[‘traindataRoot’]
validdataRoot = CONFIG[‘validdataRoot’]

TrainData = TrainDataSet(training=True, file_path=traindataRoot)
TrainDataLoader = DataLoader(TrainData, batch_size=batchSize, shuffle=True,
num_workers=CONFIG[‘numOfWorkers’], drop_last=True)
ValidData = ValidDataSet(file_path=validdataRoot)
ValidDataLoader = DataLoader(ValidData, batch_size=1, shuffle=True, num_workers=0, drop_last=True)

netG = STRnet2_change()

if CONFIG[‘pretrained’] is not None:
print('loaded ')
weights = paddle.load(CONFIG[‘pretrained’])
netG.load_dict(weights)

开始直接上大火

lr = 2e-3
G_optimizer = paddle.optimizer.Adam(learning_rate=lr, parameters=netG.parameters())

loss_function = LossWithGAN_STE()

print(‘OK!’)
num_epochs = CONFIG[‘num_epochs’]
mse = nn.MSELoss()
best_psnr = 0
iters = 0

for epoch_id in range(1, num_epochs + 1):

netG.train()

if epoch_id % 8 == 0:
    # 每8个epoch时重置优化器，学习率变为1/10
    lr /= 10
    paddle.optimizer.Adam(learning_rate=lr, parameters=netG.parameters())

for k, (imgs, gts, masks) in enumerate(TrainDataLoader):
    iters += 1

    fake_images, mm = netG(imgs)
    G_loss = loss_function(masks, fake_images, mm, gts)
    G_loss = G_loss.sum()

    #后向传播，更新参数的过程
    G_loss.backward()
    # 最小化loss,更新参数
    G_optimizer.step()
    # 清除梯度
    G_optimizer.clear_grad()

    # 打印训练信息
    if iters % 100 == 0:
        print('epoch{}, iters{}, loss:{:.5f}, net:{}, lr:{}'.format(
            epoch_id, iters, G_loss.item(), CONFIG['net'], G_optimizer.get_lr()
        ))
        log.add_scalar(tag="train_loss", step=iters, value=G_loss.item())

# 对模型进行评价并保存
netG.eval()
val_psnr = 0

# noinspection PyAssignmentToLoopOrWithParameter
for index, (imgs, gt) in enumerate(ValidDataLoader):
    _, _, h, w = imgs.shape
    rh, rw = h, w
    step = 512
    pad_h = step - h if h < step else 0
    pad_w = step - w if w < step else 0
    m = nn.Pad2D((0, pad_w, 0, pad_h))
    imgs = m(imgs)
    _, _, h, w = imgs.shape
    res = paddle.zeros_like(imgs)
    mm_out = paddle.zeros_like(imgs)
    mm_in = paddle.zeros_like(imgs)
    input_array = []
    i_j_list = []
    for i in range(0, h, step):
        for j in range(0, w, step):
            if h - i < step:
                i = h - step
            if w - j < step:
                j = w - step
            clip = imgs[:, :, i:i + step, j:j + step]
            input_array.append(clip[0])
            i_j_list.append((i, j))

    # 并行处理进行加速
    input_array = paddle.to_tensor(input_array)
    input_array = input_array.cuda()
    with paddle.no_grad():
        g_images, mm = netG(input_array)
    g_images, mm = g_images.cpu(), mm.cpu()
    for idx in range(len(i_j_list)):
        i, j = i_j_list[idx]
        mm_in[:, :, i:i + step, j:j + step] = mm[idx]
        g_image_clip_with_mask = imgs[:, :, i:i + step, j:j + step] * (1 - mm[idx]) + g_images[idx] * mm[idx]
        res[:, :, i:i + step, j:j + step] = g_image_clip_with_mask
        mm_out[:, :, i:i + step, j:j + step] = mm[idx]

    # for i in range(0, h, step):
    #     for j in range(0, w, step):
    #         if h - i < step:
    #             i = h - step
    #         if w - j < step:
    #             j = w - step
    #         clip = imgs[:, :, i:i + step, j:j + step]
    #         clip = clip.cuda()
    #         with paddle.no_grad():
    #             g_images_clip, mm = netG(clip)
    #         g_images_clip = g_images_clip.cpu()
    #         mm = mm.cpu()
    #         clip = clip.cpu()
    #         mm_in[:, :, i:i + step, j:j + step] = mm
    #         # mm = paddle.where(F.sigmoid(mm) > 0.5, paddle.zeros_like(mm), paddle.ones_like(mm))
    #         # g_image_clip_with_mask = clip * mm + g_images_clip * (1 - mm)
    #         g_image_clip_with_mask = clip * (1 - mm) + g_images_clip * mm
    #         res[:, :, i:i + step, j:j + step] = g_image_clip_with_mask
    #         mm_out[:, :, i:i + step, j:j + step] = mm

    res = res[:, :, :rh, :rw]
    mm_out = mm_out[:, :, :rh, :rw]
    # 改变通道
    output = utils.pd_tensor2img(res)
    target = utils.pd_tensor2img(gt)
    mm_out = utils.pd_tensor2img(mm_out)
    mm_in = utils.pd_tensor2img(mm_in)

    psnr_value = psnr(output, target)
    print('psnr: ', psnr_value)

    if index in [2, 3, 5, 7, 11]:
        fig = plt.figure(figsize=(20, 10),dpi=100)
        # 图一
        ax1 = fig.add_subplot(2, 2, 1)  # 1行 2列 索引为1
        ax1.imshow(output)
        # 图二
        ax2 = fig.add_subplot(2, 2, 2)
        ax2.imshow(mm_in)
        # 图三
        ax3 = fig.add_subplot(2, 2, 3)
        ax3.imshow(target)
        # 图四
        ax4 = fig.add_subplot(2, 2, 4)
        ax4.imshow(mm_out)
        plt.show()

    del res
    del gt
    del target
    del output

    val_psnr += psnr_value
ave_psnr = val_psnr / (index + 1)
print('epoch:{}, psnr:{}'.format(epoch_id, ave_psnr))
log.add_scalar(tag="valid_psnr", step=epoch_id, value=ave_psnr)
paddle.save(netG.state_dict(), CONFIG['modelsSavePath'] +
            '/STE_{}_{:.4f}.pdparams'.format(epoch_id, ave_psnr
            ))
if ave_psnr > best_psnr:
    best_psnr = ave_psnr
    paddle.save(netG.state_dict(), CONFIG['modelsSavePath'] + '/STE_best.pdparams')

4、模型预测
模型预测部分的代码保存在predict.py文件中，同在训练过程中对模型进行评估的处理方法是一致的，预测为mask的地方取模型的输出，预测为非mask的地方取输入图片的像素。这样在非mask的地方就可以保证像素差接近0（因为jpg图像本身有一些噪点，一般达不到0）。

In [ ]
import os
import sys
import glob
import json
import cv2
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from models.sa_gan import STRnet2

加载STRnet改

from models.swin_gan import STRnet2_change
import utils
from paddle.vision.transforms import Compose, ToTensor
from PIL import Image

netG = STRnet2_change()
weights = paddle.load(‘train_models_swin_v100/STE_12_44.8510.pdparams’)
netG.load_dict(weights)
netG.eval()

def ImageTransform():
return Compose([ToTensor(), ])

ImgTrans = ImageTransform()

def process(src_image_dir, save_dir):
image_paths = glob.glob(os.path.join(src_image_dir, “*.jpg”))
for image_path in image_paths:
# do something
img = Image.open(image_path)
inputImage = paddle.to_tensor([ImgTrans(img)])

    _, _, h, w = inputImage.shape
    rh, rw = h, w
    step = 512
    pad_h = step - h if h < step else 0
    pad_w = step - w if w < step else 0
    m = nn.Pad2D((0, pad_w, 0, pad_h))
    imgs = m(inputImage)
    _, _, h, w = imgs.shape
    res = paddle.zeros_like(imgs)
    for i in range(0, h, step):
        for j in range(0, w, step):
            if h - i < step:
                i = h - step
            if w - j < step:
                j = w - step
            clip = imgs[:, :, i:i + step, j:j + step]
            clip = clip.cuda()
            with paddle.no_grad():
                g_images_clip, mm = netG(clip)
            g_images_clip = g_images_clip.cpu()
            mm = mm.cpu()
            clip = clip.cpu()
            # mm = paddle.where(F.sigmoid(mm) > 0.5, paddle.zeros_like(mm), paddle.ones_like(mm))
            # g_image_clip_with_mask = clip * mm + g_images_clip * (1 - mm)
            g_image_clip_with_mask = g_images_clip * mm + clip * (1 - mm)
            res[:, :, i:i + step, j:j + step] = g_image_clip_with_mask
    res = res[:, :, :rh, :rw]
    output = utils.pd_tensor2img(res)


    # 保存结果图片
    save_path = os.path.join(save_dir, os.path.basename(image_path))
    cv2.imwrite(save_path, output)

if name == “main”:
assert len(sys.argv) == 3

src_image_dir = sys.argv[1]
save_dir = sys.argv[2]

if not os.path.exists(save_dir):
    os.makedirs(save_dir)

process(src_image_dir, save_dir)