商汤系列模型使用(一)MMOCR源码下载、数据及制作、模型训练(预训练权重)、
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文件夹,推理得到所有测试集图片的检测结果
更多推荐
所有评论(0)