1.源码下载

https://github.com/open-mmlab/mmocr/tree/v1.0.1

2.训练数据集制作

2.1官方数据集下载

参考连接:
数据准备 (Beta) — MMOCR 1.0.1 文档

下载数据集的命令行:

cd mmocr-1.0.1

#使用该脚本为 ICDAR2015 数据集准备文本检测任务所需的数据。
python tools/dataset_converters/prepare_dataset.py icdar2015 --task textrecog --overwrite-cfg

2.2 数据集转换

将标签文件的.txt文件转换为.json文件

python tools/dataset_converters/textrecog/data_migrator.py \
/home/boshi/ljj-project/Pycharm/openMMlab/mmocr-1.0.1/data/cache/ic15_textrec/ic15_textrecog_dataset/test.txt \
/home/boshi/ljj-project/Pycharm/openMMlab/mmocr-1.0.1/data/cache/ic15_textrec/ic15_textrecog_dataset/test.json --format txt

注意1:这里面有一个坑,原始的官方.txt文件当中多加入了一个逗号(修改后再转json)

原始的txt文件:
word_1.png, "JOINT"
word_2.png, "yourself"
word_3.png, "154"

修改后的txt文件:

word_1.png "JOINT"
word_2.png "yourself"
word_3.png "154"

注意2:这里面还有一个坑,原始的官方.txt文件中只有文件的名字没有包含路径,正常是需要添加文件的路径,不想添加路径的话了一将所有文件放置在同一文件夹中:
├── test.json
├── word_1000.png
├── word_1001.png
├── word_1002.png
 

2.3.制作数据集文件

修改数据集的配置文件:

文件地址:

cd mmocr-1.0.1/configs/textrecog/_base_/datasets/icdar2015.py

文件具体内容

#刚才制作的图片与数据集标签.json的存放地址
icdar2015_textrecog_data_root = '/home/boshi/ljj-project/Pycharm/openMMlab/mmocr-1.0.1/data/cache/ic15_textrec/ic15_textrecog_dataset'

#训练集
icdar2015_textrecog_train = dict(
    type='OCRDataset',
    data_root=icdar2015_textrecog_data_root,
    ann_file='test.json',  #训练集标签文件名称,图省事我用的同一个
    pipeline=None)

#验证集
icdar2015_textrecog_test = dict(
    type='OCRDataset',
    data_root=icdar2015_textrecog_data_root,
    ann_file='test.json',  #验证集标签文件名称
    test_mode=True,
    pipeline=None)

3.模型训练

模型训练采用的是命令行训练,主要调用配置文件:

参考:https://github.com/open-mmlab/mmocr/blob/main/docs/zh_cn/get_started/quick_run.md

python tools/train.py \
configs/textrecog/aster/aster_resnet45_6e_st_mj.py

修改模型的配置文件

# training schedule for 1x
_base_ = [
    '_base_aster.py',
    '../_base_/datasets/icdar2015.py',   #这个就是刚才的数据集的配置文件
    '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_adamw_cos_6e.py',
]

# dataset settings
train_list = [
    _base_.icdar2015_textrecog_train,    #我用的数据集是这个但是我不知道这个是做什么用的,官方是这么写的
]
test_list = [
    _base_.icdar2015_textrecog_test,
]

# epoch
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=2, val_interval=2)  

#下面就是慢慢优化的参数了
default_hooks = dict(logger=dict(type='LoggerHook', interval=50))

train_dataset = dict(
    type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
    type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)

train_dataloader = dict(
    batch_size=16,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    dataset=train_dataset)

auto_scale_lr = dict(base_batch_size=1024)

test_dataloader = dict(
    batch_size=16,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True,
    drop_last=False,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=test_dataset)

val_dataloader = test_dataloader

val_evaluator = dict(
    dataset_prefixes=['IC15'])
test_evaluator = val_evaluator

直接执行命令行就可以训练模型了

4.制作自己的数据集

首先转为我们自己下载的.txt文件的格式,后面就按照官方的手册走就好了

4.1 制作.txt标签

创建后的文件层次结构如下:

├── 1.5
│   ├── 11_100_0_cam_100_4_92584.jpg
│   ├── 17_89_0_cam_89_0_50723.jpg
│   └── 19_100_0_cam_100_4_37014.jpg
├── 150
│   ├── 11_96_0_cam_96_0_64814.jpg
│   ├── 2025_07_19_17_56_01_656OBB2.jpg
│   └── 2025_07_19_23_21_35_969OBB2.jpg
注意:这是由于我的数据集类相较为单一而命名的,正常情况下应该首先用python脚本获取每一张图像的名字,而后手动填写每一张图像中的文字信息

import os


def generate_label_files(data_root=".", train_output="train_label.txt", val_output="test_label.txt"):
    """
    生成训练集和验证集的标签文件,格式:文件名 标签

    Args:
        data_root (str): 数据根目录(包含 train/ 和 val/ 子目录),默认当前目录 "."
        train_output (str): 训练集标签输出文件,默认 "train_label.txt"
        val_output (str): 验证集标签输出文件,默认 "val_label.txt"
    """

    def _write_labels(data_dir, output_file):
        if not os.path.exists(data_dir):
            print(f"警告: 目录 {data_dir} 不存在,跳过生成标签文件。")
            return

        with open(output_file, 'w') as f:
            for root, _, files in os.walk(data_dir):
                label = os.path.basename(root)  # 子文件夹名作为标签
                for file in files:
                    if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                        f.write(f"{file} {label}\n")  # 仅写入文件名和标签

    # 生成训练集标签文件
    train_dir = os.path.join(data_root, "train")
    _write_labels(train_dir, train_output)
    print(f"训练集标签文件已生成: {train_output}")

    # 生成验证集标签文件
    val_dir = os.path.join(data_root, "val")
    _write_labels(val_dir, val_output)
    print(f"验证集标签文件已生成: {val_output}")


if __name__ == "__main__":
    # 示例:手动指定数据根目录
    custom_data_root = "dataset-cls"  # 替换为你的数据目录
    generate_label_files(data_root=custom_data_root)

4.2 将.txt标签转换为.json标签

#进入mmocr的环境
conda activate mmocr

#执行如下指令(.txt文件与.json文件的地址改为自己的数据集的地址)
python tools/dataset_converters/textrecog/data_migrator.py \
/home/boshi/ljj-project/Pycharm/dataset-prosecss/指针刻度识别-OCR/test_label.txt \
/home/boshi/ljj-project/Pycharm/dataset-prosecss/指针刻度识别-OCR/test_label.json --format txt

4.3 将.json与图片放在同一文件夹

├── 7_39_0_cam_39_2_77625.jpg
├── 7_39_0_cam_39_2_77626.jpg
├── 7_39_0_cam_39_2_77627.jpg
├── 7_39_0_cam_39_2_7762.jpg
├── 7_42_0_cam_42_1_7762.jpg
├── test_label.json
└── train_label.json

该文件文件夹就是我们的数据集,后面训练可以使用此文件夹

5.训练自己的数据集

5.1 修改dataset配置文件

由于我这个是参考icdar2015制作的,因此我就基于icadr2015的配置文件进行修改

icdar2015配置文件路径:

mocr-1.0.1/configs/textrecog/_base_/datasets/icdar2015.py

文件内部配置信息:

#刚才制作的图片与数据集标签.json的存放地址
icdar2015_textrecog_data_root = '/home/boshi/ljj-project/Pycharm/openMMlab/mmocr-1.0.1/data/shuzi-dataset'

#训练集
icdar2015_textrecog_train = dict(
    type='OCRDataset',
    data_root=icdar2015_textrecog_data_root,
    ann_file='train_label.json',  #训练集标签文件名称,图省事我用的同一个
    pipeline=None)

#验证集
icdar2015_textrecog_test = dict(
    type='OCRDataset',
    data_root=icdar2015_textrecog_data_root,
    ann_file='test_label.json',  #验证集标签文件名称
    test_mode=True,
    pipeline=None)

5.2 修改模型训练的配置文件

以aster_resnet45_6e_st_mj.py文件为例,文件位置:

mmocr-1.0.1/configs/textrecog/aster/aster_resnet45_6e_st_mj.py

文件内部配置信息:

# 配置文件说明:1倍训练周期的训练计划(1x表示基准训练配置)
_base_ = [  # 基础配置文件列表(继承配置)
    '_base_aster.py',  # ASTER模型架构配置文件
    '../_base_/datasets/icdar2015.py',   # ICDAR2015数据集配置文件
    '../_base_/default_runtime.py',      # 默认运行时配置(日志/checkpoint等)
    '../_base_/schedules/schedule_adamw_cos_6e.py',  # AdamW优化器+cosine学习率调度配置
]

# 数据集设置
train_list = [  # 训练集定义
    _base_.icdar2015_textrecog_train,    # 使用基础配置中的ICDAR2015文本识别训练集路径
]
test_list = [  # 测试集定义
    _base_.icdar2015_textrecog_test,     # 使用基础配置中的ICDAR2015文本识别测试集路径
]

# 训练周期配置
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=2, val_interval=2)  # 基于周期的训练循环,最大2个epoch,每2个epoch验证一次

# 默认钩子设置(用于训练过程中的辅助功能)
default_hooks = dict(logger=dict(type='LoggerHook', interval=50))  # 日志记录钩子,每50次迭代记录一次日志

# 数据集组合配置
train_dataset = dict(  # 训练数据集定义
    type='ConcatDataset',  # 数据集类型:多数据集拼接
    datasets=train_list,   # 使用的数据集列表(此处为训练集)
    pipeline=_base_.train_pipeline  # 使用基础配置中的训练数据处理流程
)
test_dataset = dict(  # 测试数据集定义
    type='ConcatDataset',   # 数据集类型:多数据集拼接
    datasets=test_list,     # 使用的数据集列表(此处为测试集)
    pipeline=_base_.test_pipeline  # 使用基础配置中的测试数据处理流程
)

# 训练数据加载器配置
train_dataloader = dict(
    batch_size=16,         # 每批加载16个样本
    num_workers=4,         # 使用4个工作进程加载数据
    persistent_workers=True,  # 保持工作进程存活(避免重复初始化)
    pin_memory=True,       # 将数据复制到CUDA固定内存(加速GPU传输)
    sampler=dict(type='DefaultSampler', shuffle=True),  # 默认采样器,训练时打乱数据顺序
    dataset=train_dataset  # 使用上面定义的训练数据集
)

# 学习率自动缩放配置
auto_scale_lr = dict(base_batch_size=1024)  # 基础batch size为1024(用于学习率自动调整)

# 测试数据加载器配置
test_dataloader = dict(
    batch_size=16,          # 每批加载16个样本
    num_workers=4,          # 使用4个工作进程加载数据
    persistent_workers=True,  # 保持工作进程存活
    pin_memory=True,        # 将数据复制到CUDA固定内存
    drop_last=False,        # 不丢弃最后不足batch的数据
    sampler=dict(type='DefaultSampler', shuffle=False),  # 测试时不打乱数据顺序
    dataset=test_dataset    # 使用上面定义的测试数据集
)

# 验证数据加载器配置(直接复用测试数据加载器)
val_dataloader = test_dataloader  # 验证集使用与测试集相同的配置

# 验证评估器配置
val_evaluator = dict(
    dataset_prefixes=['IC15']  # 评估指标前缀(IC15表示ICDAR2015数据集)
)

# 测试评估器配置(直接复用验证评估器)
test_evaluator = val_evaluator  # 测试评估使用与验证相同的配置

5.3 训练模型

python tools/train.py configs/textrecog/aster/aster_resnet45_6e_st_mj.py

5.4 模型检测效果验证

python tools/test.py configs/textrecog/aster/aster_resnet45_6e_st_mj.py work_dirs/aster_resnet45_6e_st_mj/epoch_2.pth --show-dir test_result/

程序会生成一个test_result文件夹,推理得到所有测试集图片的检测结果

Logo

更多推荐