转自AI Studio，原文链接：百度网盘AI大赛-图像处理挑战赛：文档检测优化赛 Baseline - 飞桨AI Studio

百度网盘AI大赛-图像处理挑战赛：文档检测优化赛

使用Resnet152回归图像中文档的拐角坐标完成百度网盘AI大赛-图像处理挑战赛：文档检测优化赛。

比赛链接

比赛介绍

生活中人们使用手机进行文档扫描逐渐成为一件普遍的事情，为了提高人们的使用体验，我们期望通过算法技术去除杂乱的拍摄背景并精准框取文档边缘，选手需要通过深度学习技术训练模型，对给定的真实场景下采集得到的带有拍摄背景的文件图片进行边缘智能识别，并最终输出处理后的扫描结果图片。

评测方式说明

1. 参赛选手提交的代码和模型只预测文档边缘的heatmap图，由后台评测脚本中预置的算法回归出文档区域的四个角的坐标点，并生成规则的四边形，与GT计算IoU值；
2. 参赛选手提交的代码和模型直接回归文档区域的四个角的坐标点，并生成规则的四边形，与GT计算IoU值。 注：两种评测方式的结果会放在一个排行榜内，评测脚本中预置的角点回归算法会存在一定的局限性，如果各位参赛选手希望获得更高的分数，建议采用第二种评测方式。

任务分析

本次比赛要求选手设计算法在给定图片中划定一块四边形区域，以尽可能与图片中的文档部分重合。

因此，本次任务可以同时看作回归问题和分割问题。

• 作为回归问题，需要设计学习器/其他方法寻找四边形的四个拐角坐标
• 作为分割问题，直接将数据中的segments作为标签训练对应的分割学习器即可

本项目将本次任务看作回归问题来处理，使用Resnet152+Linear的网络结构回归四个角的坐标。

代码部分

获取数据

In [ ]

! wget https://staticsns.cdn.bcebos.com/amis/2022-4/1649731549425/train_datasets_document_detection_0411.zip
! unzip -oq /home/aistudio/train_datasets_document_detection_0411.zip
! rm -rf __MACOSX
! rm -rf /home/aistudio/train_datasets_document_detection_0411.zip

构造数据读取器

通过paddle.io.dataset构造读取器，便于读取数据。

数据预处理包括：

1. 根据data_info的边缘轮廓信息分别提取四个拐角的位置。
2. 对图片进行resize和数据增强（调整明暗对比度等参数）

In [ ]

import paddle

import numpy as np
import pandas as pd

import cv2

class MyDateset(paddle.io.Dataset):
    def __init__(self, mode = \'train\', train_imgs_dir = \'/home/aistudio/train_datasets_document_detection_0411/images/\', train_txt = \'/home/aistudio/train_datasets_document_detection_0411/data_info.txt\'):
        super(MyDateset, self).__init__()

        self.mode = mode 
        self.train_imgs_dir = train_imgs_dir

        with open(train_txt,\'r\') as f:
            self.train_infor = f.readlines()

    def __getitem__(self, index):
        item = self.train_infor[index][:-1]
        splited = item.split(\',\')
        img_name = splited[0]

        img = cv2.imread(self.train_imgs_dir+img_name+\'.jpg\')
        h, w, c = img.shape

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 对图片进行resize，调整明暗对比度等参数
        img = paddle.vision.transforms.resize(img, (512,512), interpolation=\'bilinear\')
        if np.random.rand()<1/3:
            img = paddle.vision.transforms.adjust_brightness(img, np.random.rand()*2)
        else:
            if np.random.rand()<1/2:
                img = paddle.vision.transforms.adjust_contrast(img, np.random.rand()*2)
            else:
                img = paddle.vision.transforms.adjust_hue(img, np.random.rand()-0.5)

        img = img.transpose((2,0,1))
        img = img/255

        sites = []
        for i in range(1,len(splited),2):
            sites.append([float(splited[i])/w,float(splited[i+1])/h])
        
        label = []
        for i in range(4):
            x, y = self.get_corner(sites, i+1)
            label.append(x)
            label.append(y)

        img = paddle.to_tensor(img).astype(\'float32\')
        label = paddle.to_tensor(label).astype(\'float32\')

        return img, label

    def get_corner(self, sites, corner_flag):
        # corner_flag 1:top_left 2:top_right 3:bottom_right 4:bottom_left
        if corner_flag == 1:
            target_sites = [0,0]
        elif corner_flag == 2 :
            target_sites = [1,0]
        elif corner_flag == 3 :
            target_sites = [1,1]
        elif corner_flag == 4 :
            target_sites = [0,1]

        min_dis = 3
        best_x = 0
        best_y = 0
        for site in sites:
            if abs(site[0]-target_sites[0])+abs(site[1]-target_sites[1])

class MyNet(paddle.nn.Layer):
    def __init__(self):
        super(MyNet,self).__init__()
        self.resnet = paddle.vision.models.resnet152(pretrained=True, num_classes=0)
        self.flatten = paddle.nn.Flatten()
        self.linear = paddle.nn.Linear(2048, 8)

    def forward(self, img):
        y = self.resnet(img)
        y = self.flatten(y)
        out = self.linear(y)

        return out

model = MyNet()
model.train()

train_dataset=MyDateset()

# 需要接续之前的模型重复训练可以取消注释
# param_dict = paddle.load(\'./model.pdparams\')
# model.load_dict(param_dict)

train_dataloader = paddle.io.DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    drop_last=False)

max_epoch=10
scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=0.0001, T_max=max_epoch)
opt = paddle.optimizer.Adam(learning_rate=scheduler, parameters=model.parameters())

now_step=0
for epoch in range(max_epoch):
    for step, data in enumerate(train_dataloader):
        now_step+=1

        img, label = data
        pre = model(img)
        loss = paddle.nn.functional.square_error_cost(pre,label).mean()
        loss.backward()
        opt.step()
        opt.clear_gradients()
        if now_step%100==0:
            print(\"epoch: {}, batch: {}, loss is: {}\".format(epoch, step, loss.mean().numpy()))

paddle.save(model.state_dict(), \'model.pdparams\')

# 压缩可提交文件
! zip submit.zip model.pdparams predict.py