★★★ 本文源自AlStudio社区精品项目,【点击此处】查看更多精品内容 >>>

一、赛题背景

为了帮助人们消除带有水印的图片日常生活中带有水印的图片很常见,即使是PS专家,也很难快速且不留痕迹的去除水印。在上次比赛中我们举办了去水印模型大赛,本次我们提高了难度,提高了在水印种类的广泛性,同时在模型性能上做了更高的要求,向各位选手发起挑战。
此外,由于本次比赛是第二次举办,在模型性能上有更高的要求,因此本次比赛的目标是兼顾性能与复杂度,尽可能做到又快又准。

二、数据分析

本次比赛的数据集大致可分为两类,第一类是上期比赛使用的包含文字型水印的图片,第二类是本期新增的包含图片logo类水印的图片。

官方只给出了生成文字类水印图片的脚本,我们对此进行了一定的修改使其可以生成图片logo类水印,然后通过网络查找下载了100张logo类图片用于生成此类水印。

三、模型设计

网络模型,参考了上期第一名的方案,并且结合上次通用场景去除手写的经验。仍然采用带有分割头的两阶段擦除网络。

模型:

四、训练细节

学习率调度:StepDecay

为了获得更好的精度,也为了节省训练时间,本次最终仅对图片logo类水印图片进行训练,而文字类水印图片的消除采用上期冠军的最优模型。

训练大致分为两阶段:

第一阶段训练100Epoch,lr=1e-4,采用l1loss;

第二阶段训练100Epoch,lr=2e-5,采用PSNRloss;

由于网络较深,为了模型真正收敛,后期模型改进后又多次进行多阶段的训练,没有最好只有更好,就是练!

五、测试细节

测试trick:

两类型测试,为了加快推理时间,根据尺寸大小对图像进行分类,其中小于1400x1400的图像padding到128的整数倍直接预测,对于大于1400x1400的图像进行交错切分成1024x1024的patch,重叠区域尺寸为112进行分块预测,并且采用翻转测试增强,最后取平均。

两阶段测试,为了提升水印去除的精度,采用两阶段进行预测,即先去除文字类型水印,再去除图片logo类型水印。

六、文件

loss:定义损失函数
model:定义网络模型
predict_code:提交测试文件
gen: 生成水印脚本
dataloader.py: 加载数据脚本
predict.py: 测试脚本
train.py: 训练脚本

七、代码运行:

# 解压数据集文件
import os

!ls /home/aistudio/data/
!unzip -qo /home/aistudio/data/data193099/test_real.zip -d /home/aistudio/data/  #解压测试集
!unzip -qo /home/aistudio/data/data193099/test_make_A.zip -d /home/aistudio/data/  #解压测试集
!unzip -qo /home/aistudio/data/data193099/data_qushuiyin_tu.zip -d /home/aistudio/data/  #解压测试集
!unzip -qo /home/aistudio/data/data193106/gts.zip -d /home/aistudio/data/  #解压测试集
#!tar -zxvf /home/aistudio/data/data189925/data_qushuiyin.tar.gz -C /home/aistudio//data #解压训练集
!ls /home/aistudio/data/
print("load down")

训练模型

# step 1,
%cd ~
!python work/train.py \
    --modelsSavePath /home/aistudio/work/savepath1 \
    --batchSize 10 \
    --loadSize 512 \
    --loss l1 \
    --lr 1e-3 \
    --num_epochs 300 \
    --lr_decay_iters 10 \
    --logPath /home/aistudio/work/log1
# step 2
%cd ~
!python work/train.py \
    --modelsSavePath savepath2 \
    --batchSize 10 \
    --loadSize 512 \
    --loss psnr \
    --pretrained /home/aistudio/work/savepath1/best.pdparams \
    --lr 5e-5 \
    --num_epochs 300 \
    --lr_decay_iters 10 \
    --logPath /home/aistudio/work/log2

生成测试结果

# 评估A榜
%cd ~
!python work/predict_code/predict.py /home/aistudio/data/test_make_A /home/aistudio/work/result_A
#!python work/predict.py /home/aistudio/data/dehw_testB_dataset /home/aistudio/work/resultB
# 评估真实场景
%cd ~
!python work/predict_code/predict.py /home/aistudio/data/test_real /home/aistudio/work/result_real
#!python work/predict.py /home/aistudio/data/dehw_testB_dataset /home/aistudio/work/resultB

此文章为转载
原为链接

Logo

学大模型,用大模型上飞桨星河社区!每天8点V100G算力免费领!免费领取ERNIE 4.0 100w Token >>>

更多推荐