本文围绕水印擦除任务展开，分析其与手写文字擦除的差异及难点。介绍数据处理方式，包括生成mask、缩减数据集、随机裁剪。还阐述了模型训练及预测，A榜用Erasenet并改损失函数，B榜优化模型结构，以及模型优化和使用说明。
☞☞☞AI 智能聊天, 问答助手, AI 智能搜索, 免费无限量使用 DeepSeek R1 模型☜☜☜
赛题任务需要对添加了水印的图像，将水印擦除掉，还原原本的图的样子(图1)。
与手写文字擦除任务(图2)一个比较大的区别是：水印占据面积很大，因此对水印擦除后，还需要对被擦除的区域进行一个填补，这个是该项目的难点所在。
结论，本任务单纯使用语义分割效果不佳，需要使用带有生成能力的img2img式模型。
数据处理的方式决定了模型的设计，也会对预测的精度产生较大的影响。
(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。
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。
(1)、损失函数调整： 在所有超参数的调整中，我们把损失函数放在首要位置，因为其直面数据集。
调整mask损失，不仅使用bce，也使用l1。 增加image_loss，该任务对生成的要求更高，而且mask十分规整，因此加大image_loss的权重，增加到1.5，其他的则相对的调整到0.5。最后将所有loss相加，psnr有了一个很大的提升。
(2)、优化器调整： 在所有超参数的调整中，我们把优化器模式放在第二位置，因为其决定着训练能否达到当前模型最优。
每次重新调用优化器，就相当于对模型加载了一个经过预训练的模型。稳定后续训练，并且逃离局部平坦区域(梯度接近0)。
(3)、结构调整： 在所有超参数的调整中，我们把网络结构放在第三位置，因为其难调整，需要训练到平稳才能看出模型的好坏。
因为该任务，生成是比较困难的，尤其是彩色水印叠加在彩色图像上之后，虽然可以检测到mask，但是生成的效果不佳。
由于本次比赛，水印擦除挑战赛评分提交的是模型加模型参数文件，因此就没有再加入对A，B榜的图片进行推理的notebook。因此本项目只关心于如何训练模型。
首先运行下面cell的代码，对数据集进行解压。
# 解压文件!unzip -oq data/data145795/train_dataset.zip -d ./dataset !unzip -oq data/data145795/valid_dataset.zip -d ./dataset
数据处理的方式决定了模型的设计，也会对预测的精度产生较大的影响。与手写文字擦除任务一个比较大的区别是：水印占据面积很大，因此对水印擦除后，还需要对被擦除的区域进行一个填补，这个是该项目的难点所在。为了显示的引导模型进行预测，需要结合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大小
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。
运行trainstr.ipynb可以训练原始erasenet，训练日志log和最好的模型都已包含在项目中，用visualdl即可可视化。虽然最后不会用这个模型提交，但还是放在这，可以起到一个参考的作用，因为batchsize达到28，所以训练起来是要比erasenet改快一点的。
erasenet改是分两部分训练的，开始是用的A100，运行trainswin.ipynb即可，但是A100只能训练24小时，因此将最好的模型加载再使用V100进行训练，运行trainswinv100即可。log_swin，log_swin_v100包含了完整的训练日志。我们只是因为时间紧迫才用的A100训练的，但这并不是必要的，单纯用V100多训练几天也是可以的。
import warnings
warnings.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 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}# 设置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')模型预测部分的代码保存在predict.py文件中，同在训练过程中对模型进行评估的处理方法是一致的，预测为mask的地方取模型的输出，预测为非mask的地方取输入图片的像素。这样在非mask的地方就可以保证像素差接近0（因为jpg图像本身有一些噪点，一般达不到0）。
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 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)以上就是百度网盘AI大赛——水印智能消除赛：第8名方案的详细内容，更多请关注php中文网其它相关文章！
广告
收藏
收藏
收藏
Copyright 2014-2025 https://www.php.cn/ All Rights Reserved | php.cn | 湘ICP备2023035733号
291
