超详细记录在pycharm中训练测试DBnet(pytorch)

论文链接：https://arxiv.org/pdf/1911.08947.pdf
项目链接：https://github.com/WenmuZhou/DBNet.pytorch

项目介绍

DBnet是基于分割的文本检测方法，能够更好的检测自然场景下中不同形状的文字

1. 网络结构
   

2. 导入：下载代码后解压并导入pycharm中
   

3. 修改部分代码：
   由于作者输入有误，需要在代码中做部分修改
   3.1将’–save_resut’修改为’–save_result’
   3.2
   在训练中还会出现No module named ‘torchvision.models.utils‘错误，因此只需要将 from torchvision.models.utils import load_state_dict_from_url 改为 from torch.hub import load_state_dict_from_url即可
   

4. 环境安装
   打开README在找到以下代码：
   pip install -r requirement.txt
   并在终端输入按回车即可
   
   
   不过个人是一个一个命令输入安装也比较快
   打开requirement.txt，并在terminal中输入相应命令即可：
   例如：pip install addict -i https://mirrors.aliyun.com/pypi/simple/
   
   输入并按回车即可安装
   

5. 数据集：文章采用的是icdar2015数据集
   链接：https://rrc.cvc.uab.es/?ch=4&com=tasks
   选择Task 4.1: Text Localization，下载即可
   不过需要使用邮箱注册
   下载好后解压到相应文件夹，并按照如下方式放到项目中：
   

6. 数据处理：作者是在Ubuntu系统下的处理的，而我使用的是windows系统，因此输入以下数据代码

import os
def get_images(img_path):
    '''
    find image files in data path
    :return: list of files found
    '''
    files = []
    exts = ['jpg', 'png', 'jpeg', 'JPG', 'PNG']
    for parent, dirnames, filenames in os.walk(img_path):
        for filename in filenames:
            for ext in exts:
                if filename.endswith(ext):
                    files.append(os.path.join(parent, filename))
                    break
    print('Find {} images'.format(len(files)))
    return sorted(files)

def get_txts(txt_path):
    '''
    find gt files in data path
    :return: list of files found
    '''
    files = []
    exts = ['txt']
    for parent, dirnames, filenames in os.walk(txt_path):
        for filename in filenames:
            for ext in exts:
                if filename.endswith(ext):
                    files.append(os.path.join(parent, filename))
                    break
    print('Find {} txts'.format(len(files)))
    return sorted(files)

if __name__ == '__main__':
    import json

    img_train_path = r'F:\apy\DBNet\datasets\train\img\ch4_training_images'
    img_test_path = r'F:\apy\DBNet\datasets\test\img\ch4_test_images'
    train_files = get_images(img_train_path)
    test_files = get_images(img_test_path)

    txt_train_path = r'F:\apy\DBNet\datasets\train\gt\ch4_training_localization_transcription_gt'
    txt_test_path = r'F:\apy\DBNet\datasets\test\gt\Challenge4_Test_Task1_GT'
    train_txts = get_txts(txt_train_path)
    test_txts = get_txts(txt_test_path)
    n_train = len(train_files)
    n_test = len(test_files)
    assert len(train_files) == len(train_txts) and len(test_files) == len(test_txts)
    # with open('train.txt', 'w') as f:
    with open('./datasets/train.txt', 'w') as f:
        for i in range(n_train):
            line = train_files[i] + '\t' + train_txts[i] + '\n'
            f.write(line)
    with open('./datasets/test.txt', 'w') as f:
        for i in range(n_test):
            line = test_files[i] + '\t' + test_txts[i] + '\n'
            f.write(line)

其中并在代码中修改相应的文件路径即可

10. 训练：
    先在config/icdar2015_resnet18_FPN_DBhead_polyLR.yaml里面配置相应参数
    
    在terminal中输入：
    python tools/train.py --config_file "config/icdar2015_resnet18_FPN_DBhead_polyLR.yaml"
    大概训练了三天左右
    

11. 测试：在tools/predict.py中修改以下参数并运行即可
    
    结果如下：
    
    
    以上就是完整的训练测试过程！希望对大家有用！

参考链接：
[1]: https://bbs.huaweicloud.com/blogs/345205
[2]: https://github.com/WenmuZhou/DBNet.pytorch
[3]: https://rrc.cvc.uab.es/?ch=4&com=tasks