本文围绕水印擦除任务展开，分析其与手写文字擦除的差异及难点。介绍数据处理方式，包括生成mask、缩减数据集、随机裁剪。还阐述了模型训练及预测，A榜用Erasenet并改损失函数，B榜优化模型结构，以及模型优化和使用说明。

一、赛题解读

1、赛题分析

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

免费影视、动漫、音乐、游戏、小说资源长期稳定更新！ 👉 点此立即查看 👈

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

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

       图1、左(添加了水印的图片)，右(真实图片)        图2、左(带手写的图片)，右(真实图片)

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、手写文字擦除baseline，使用了resize导致生成的图像非常的模糊，非常影响psnr)

3、模型训练及预测

A榜用的是Erasenet，同手写文字擦除一样，我们更改了loss函数，因为这个方式比较直接效果显著(模型是需要训练的，并不是设计的越复杂越好，直接调整面向真实数据的loss设计可以有效改变模型训练的轨迹)。模型结构图如下：

       图4、Erasenet模型

预测时：先将图像重叠分块到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。

       图5 从左到右依次为img, 原始erasenet，erasenet改，gt

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，代码中图片路径供参考，是在本地电脑进行处理的):

       图1.1 加了水印的图像(左)，真实图像(中)，mask(右)

另一方面，由于本次比赛数据集过大，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设计可以有效改变模型训练的轨迹)。模型结构图如下：

       图2.1 Erasenet主体结构

模型数据流向大体如上，loss的地方做了一定的修改。

B榜对模型进行了一次调优，方法是将网络最开始下采样和精修部分下采样的卷积替换成了SwinT模块，就像在我之前Swin那个项目里一样，将Swin和CNN成功的结合起来，做到又快又好，最终B榜分数也比较高。下图展示了原Erasenet和带swin的Erasenet改在验证集上的表现，psnr分别是31.418，33.042。

       图2.2 从左到右依次为img, 原始erasenet，erasenet改，gt再次印证了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 warningswarnings.filterwarnings("ignore")# 进行训练from visualdl import LogWriterimport osimport paddleimport paddle.nn as nnimport paddle.nn.functional as Ffrom paddle.io import DataLoaderfrom dataset.data_loader import TrainDataSet, ValidDataSetfrom loss.Loss import LossWithGAN_STE, LossWithSwinfrom models.swin_gan import STRnet2_changeimport utilsimport randomfrom PIL import Imageimport matplotlib.pyplot as pltimport numpy as npimport math%matplotlib inlinelog = 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}# 设置gpuif 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 PyProtectedMemberpaddle.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-3G_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 = 0iters = 0for 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 osimport sysimport globimport jsonimport cv2import paddleimport paddle.nn as nnimport paddle.nn.functional as Ffrom models.sa_gan import STRnet2# 加载STRnet改from models.swin_gan import STRnet2_changeimport utilsfrom paddle.vision.transforms import Compose, ToTensorfrom PIL import ImagenetG = 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, "*.webp"))    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)

登录后复制    

相关攻略

热门推荐