百度网盘AI大赛——图像处理挑战赛： 通用场景手写文字擦除基线方案

赛事链接：百度网盘AI大赛——图像处理挑战赛： 通用场景手写文字擦除

1、大赛背景

  百度网盘AI大赛——图像处理挑战赛是百度网盘开放平台面向开发者发起的图像处理挑战赛事，旨在基于个人云存储的生态能力开放，通过比赛机制，鼓励选手结合当下的前沿图像处理与计算机视觉技术，设计合适模型，并提升模型的效果，助力互联网用户更便捷地进行数字生活、学习和工作，为中国开源生态建设贡献力量。本次图像处理挑战赛以线上比赛的形式进行，参赛选手在规定时间内提交基于评测数据集产出的结果文件，榜单排名靠前并通过代码复查的队伍可获得高额奖金。百度网盘开放平台致力于为全球AI开发者和爱好者提供专业、高效的AI学习和开发环境，挖掘培养AI人才，助力技术产业生态发展。

  随着技术发展，OCR扫描在学习、办公等众多场景中被使用，通过技术和算法，对扫描获得的纸张文档上的手写笔迹还原修复，恢复文件本身的样子，使得人们的使用体验越来越便捷。上一期比赛，我们举办了试卷场景下的手写文字擦除，帮助学生党们擦除试卷上的笔迹。本次比赛，我们诚邀各位选手并拓宽场景：不限于试卷，对通用文件上的手写笔迹进行擦除后还原文件，帮助更多人解决扫描上的问题。

2、赛题任务与数据说明

  随着技术发展，OCR扫描在学习、办公等众多场景中被使用，通过技术和算法，对扫描获得的纸张文档上的手写笔迹还原修复，恢复文件本身的样子，使得人们的使用体验越来越便捷。上一期比赛，我们举办了试卷场景下的手写文字擦除，帮助学生党们擦除试卷上的笔迹。本次比赛，我们诚邀各位选手并拓宽场景：不限于试卷，对通用文件上的手写笔迹进行擦除后还原文件，帮助更多人解决扫描上的问题。

以下数据分析来自于项目：百度网盘AI大赛——通用场景手写文字擦除BaseLine感谢大佬一心炼银

本赛题为img2img的图像处理任务，因此数据集也全部以图像格式给出。这里我挑选数据集中两个具有代表性的图像作为示例。

• 如下图1所示为第一类，从左到右依次为手写图片，mask，真实图片。这类图片的特点是

1、尺寸小，基本上宽高都在1000像素点以内

2、数量多，这类图片数据集一共大概25G左右十分庞大

3、手写文字有时会覆盖在文档图片上方，因此不能简单的将文字区域像素设置为固定值就完事

• 如下图2所示为第二类，从左到右依次为手写图片，mask（自己使用代码进行生成，本次比赛数据集官方没有提供这类图片的mask数据），真实图片。这类图片的特点是

1、尺寸巨大，有的宽高可以达到5000像素点，如果想要提高性能分，可以重点关注一下如何对该类图片进行处理

2、数量少，但是在A榜中的占比却高达0.4，不像训练数据集中的不到0.02。

3、是真实的拍照得到的图片，这类图片是先有手写图片，而后才有的真实图片(与第一类图片不一样，第一类是先有真实图片，而后才有的手写图片)

3、评价标准

  评价指标为 PSNR 和 MSSSIM；

4、解题思路与基线搭建

  通过对比赛赛题的分析以及比赛相关介绍，我们将本赛题解释为一种图像分割任务，所以本项目仍然可以参考一下之前赛题的一些方法，所以在本项目使用了deeplabv3+网络，对手写文字图像进行分割。

4.1 deeplabv3+简介

  到目前为止，DeepLab系列都是在降采样8倍的尺度上进行预测的，导致了边界效果不甚理想。考虑到卷积网络的特征，DeepLab v3的网络的特征并没有包含过多的浅层特征，为了解决这个问题，DeepLab v3+借鉴了FPN等网络的encoder-decoder架构，实现了Feature Map跨block的融合。DeepLab v3+的另一个改进点在于使用了分组卷积来加速。

参考资料：

deeplab v3+

Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation

4.2 数据读取器构建

!cd data/data154717 && tar xvf baseline.tar
!cd data/data154717/baseline/ && unzip -q '*.zip'
# ! pip install paddlex

import os,cv2
import numpy as np
import fnmatch
import matplotlib.pyplot as plt
from PIL import Image


np.set_printoptions(threshold=np.sys.maxsize)

images_path = 'data/data154717/baseline/01/images/'
masks_path = 'data/data154717/baseline/01/mask'
masked_path = 'masked'

def find_files(directory, pattern):
    """
    [description]
    Arguments:
        directory {[type]} -- [description]
        pattern {[type]} -- [*.[jJPp][PpNn]*]

    Returns:
        [type] -- [description]
    """
    matches = []
    for root, dirs, files in os.walk(directory):
        for basename in files:
            if fnmatch.fnmatch(basename, pattern):
                filename = os.path.join(root, basename)
                matches.append(filename)
    return matches

img_list = find_files(masks_path, pattern="*.[jZJPp][PpNn]*")
num = 0
for img_item in img_list:
    data = Image.open(img_item)
    mask1 = data.convert('L')
    mask1 = np.array(mask1)
    mask = mask1.copy()
    # width = mask.shape[0] #获取宽度
    # height = mask.shape[1] #获取长度
    mask = np.array(mask)
    
    mask[mask1 == 0] = 1
    # mask[mask1 == 2] = 0
    mask[mask1 == 255] = 0
    
    mask = Image.fromarray(mask.astype(np.uint8), mode='P')
    name = img_item.split('/')[-1]
    mask.save(os.path.join(masked_path, name.split('.')[0] + '.png'))
   
    num += 1 
    

from PIL import Image
import os
"""生成deeplabv3数据读取txt文件"""
if not os.path.exists('train_path/'):
    os.makedirs('train_path/')

#取得所有图片路径，按顺序存入列表
path_list = []
for index in os.listdir('masked/'):
       
    mask_path = 'masked/' + index
    image_path = 'data/data154717/baseline/01/images/' + index.split('.')[0] +'.jpg'

    img = Image.open(image_path)
    mask = Image.open(mask_path)
    # print(mask.shape)
    item = str(image_path + ' ' + mask_path + '\n')
    path_list.append(item)
    
train_list = path_list[0:int(0.8 * num)] 
val_list = path_list[int(0.8 * num):]

#写入txt文件
train_txt = open('train_path/train.txt', mode='w')
train_txt.writelines(train_list)
train_txt.close()

val_txt = open('train_path/val.txt', mode='w')
val_txt.writelines(val_list)
val_txt.close()

labels_txt = open('train_path/labels.txt', mode='w')
labels_txt.writelines(['background\n', 'hand_write'])
labels_txt.close()

all_data_txt = open('train_path/all_data.txt', mode='w')
all_data_txt.writelines(train_list)
all_data_txt.writelines(val_list)
all_data_txt.close()

4.4 训练模型并保存

  本部分采用了PaddleX进行训练，飞桨全流程开发工具，集飞桨核心框架、模型库、工具及组件等深度学习开发所需全部能力于一身，打通深度学习开发全流程。PaddleX同时提供简明易懂的Python API，及一键下载安装的图形化开发客户端。用户可根据实际生产需求选择相应的开发方式，获得飞桨全流程开发的最佳体验。

"""导入包"""
import paddlex as pdx
import random
import numpy as np
import paddle

from paddlex import transforms as T

"""定义数据处理方式"""
train_transforms = T.Compose([
    T.Resize(target_size=512),
    T.MixupImage(alpha=1.5, beta=1.5, mixup_epoch=-20),
    T.RandomDistort(brightness_range=0.5, brightness_prob=0.5, contrast_range=0.5, contrast_prob=0.5, 
    saturation_range=0.5, saturation_prob=0.5, hue_range=18, 
    hue_prob=0.5, random_apply=True, count=4, shuffle_channel=False),
    T.RandomBlur(prob=0.5),
    T.Normalize(
        mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

eval_transforms = T.Compose([
    T.Resize(target_size=512),
    T.Normalize(
        mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

"""直接加载eval_dataset，后面预测就不用传入transfroms了"""
train_dataset = pdx.datasets.SegDataset(
    data_dir='',
    file_list='train_path/all_data.txt',
    label_list='train_path/labels.txt',
    transforms=train_transforms,
    shuffle=True)

eval_dataset = pdx.datasets.SegDataset(
    data_dir='',
    file_list='train_path/val.txt',
    label_list='train_path/labels.txt',
    transforms=eval_transforms)

"""训练模型并保存，可视化训练过程"""
# 使用随机增强不要固定种子
# random.seed(2048)
# np.random.seed(2048)
# paddle.seed(2048)

num_classes = len(train_dataset.labels)
model = pdx.seg.DeepLabV3P(num_classes=num_classes, use_mixed_loss=True)
model.train(
    num_epochs=1,
    train_dataset=train_dataset,
    train_batch_size=16,
    eval_dataset=eval_dataset,
    learning_rate=0.1,
    lr_decay_power=0.9,
    save_interval_epochs=1,
    save_dir='model/deeplab_augument_alldata3')

4.5 使用PaddleX进行预测

from PIL import Image
import paddlex as pdx


import os,cv2
import numpy as np
import fnmatch
import matplotlib.pyplot as plt
from PIL import Image

test_images_path = 'data/data154717/baseline/dehw_testA_dataset'

def find_files(directory, pattern):
    """
    [description]
    Arguments:
        directory {[type]} -- [description]
        pattern {[type]} -- [*.[jJPp][PpNn]*]

    Returns:
        [type] -- [description]
    """
    matches = []
    for root, dirs, files in os.walk(directory):
        for basename in files:
            if fnmatch.fnmatch(basename, pattern):
                filename = os.path.join(root, basename)
                matches.append(filename)
    return matches

# image_path = 'data/data154440/deblur_test'
test_img_list = find_files(test_images_path, pattern="*.[jZJPp][PpNn]*")
model = pdx.load_model('model/deeplab_augument_alldata3/best_model')
for img in test_img_list:

    save_name = img.split('/')[-1]
    save_name = save_name.split('.')[0]
    img = Image.open(img)
    img = np.array(img)
    result = model.predict(img)
    mask = result['label_map']
    # img = Image.fromarray(img)
    
    mask = np.array(mask)
    img[mask == 1] = (255,255,255)

    img = Image.fromarray(img)
    # break
    # img.save('resulte/'+ save_name + '.png')  #保存

4.6 模型导出

  由于提交的环境不支持PaddleX，所以这里需要将模型导出，直接使用飞桨环境进行预测。

!paddlex --export_inference --model_dir=./model/deeplab_augument_alldata3/best_model/ --save_dir=./inference_model

import os
import sys
import glob
import json
import cv2
import numpy as np
from PIL import Image
import paddle

import pandas as pd

import fnmatch
import matplotlib.pyplot as plt
from PIL import Image


np.set_printoptions(threshold=np.sys.maxsize)

test_images_path = 'data/data154717/baseline/dehw_testA_dataset'

def find_files(directory, pattern):
    """
    [description]
    Arguments:
        directory {[type]} -- [description]
        pattern {[type]} -- [*.[jJPp][PpNn]*]

    Returns:
        [type] -- [description]
    """
    matches = []
    for root, dirs, files in os.walk(directory):
        for basename in files:
            if fnmatch.fnmatch(basename, pattern):
                filename = os.path.join(root, basename)
                matches.append(filename)
    return matches


use_gpu = True
# model_file_path="Hapi_MyCNN"
paddle.set_device('gpu:0') if use_gpu else paddle.set_device('cpu')
model = paddle.jit.load('inference_model/inference_model/model')

model.eval() #训练模式

test_img_list = find_files(test_images_path, pattern="*.[jZJPp][PpNn]*")
# model = pdx.load_model('model/deeplab_augument_alldata3/best_model')
for img in test_img_list:

    save_name = img.split('/')[-1]
    save_name = save_name.split('.')[0]
    img = Image.open(img)
   
    img = np.array(img).astype('float32')
    img = img.transpose((2, 0, 1))
    # img = img / 255.
    img = img[np.newaxis,:, : ,:]
    img = paddle.to_tensor(img)
    # 
    # print(img.shape)
    result = model(img)
    
    mask = result[0]
    # # img = Image.fromarray(img)
    mask = paddle.squeeze(mask)
    # print(mask.shape)
    mask = np.array(mask)
    img = paddle.squeeze(img)
    img = np.array(img)
    img = img.transpose((1, 2, 0))
    
    img = np.uint8(img)
    # mask[mask == 1] = 255
    mask = np.uint8(mask)
    img[mask == 1] = (255,255,255)
    
    dst = cv2.inpaint(img,mask,3,cv2.INPAINT_TELEA)
    img = Image.fromarray(np.uint8(dst))

    img.save('resulte1/'+ save_name + '.png')  #保存

4.7 生成提交文件

! zip result.zip -r inference_model predict.py

4.8 预测结果展示

5、总结

  本项目提供了百度网盘AI大赛——图像处理挑战赛： 通用场景手写文字擦除基线方案。该比赛属于image2image任务，此类型任务实际上还有更多的解决方案，比如GAN等，在本项目中将其处理为图像分割任务。

  一些提升方案（未验证）：

1. 本项目使用了deeplabv3+网络，可对网络进行修改，例如Loss等；
2. 本项目中未进行数据增强的一些操作，大家可以尝试一些数据增强操作。

最后，祝大家都能有好成绩！

PS：本项目仅迭代了一轮，大家可以增加迭代轮数，肯定会涨点。

本项目仅为搬运，原作链接：https://aistudio.baidu.com/aistudio/projectdetail/4267688?contributionType=1&shared=1