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

! 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

import paddleimport numpy as npimport pandas as pdimport cv2class 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])<min_dis:
                min_dis=abs(site[0]-target_sites[0])+abs(site[1]-target_sites[1])
                best_x = site[0]
                best_y = site[1]        
        return best_x, best_y    def __len__(self):
        return len(self.train_infor)# 对dataloader进行测试'''
train_dataset=MyDateset()

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

for step, data in enumerate(train_dataloader):
    img, label = data
    print(step, img.shape, label.shape)
    break

'''

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=10scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=0.0001, T_max=max_epoch)
opt = paddle.optimizer.Adam(learning_rate=scheduler, parameters=model.parameters())

now_step=0for 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