使用mindee/doctr中的数据集模块指南

使用mindee/doctr中的数据集模块指南

doctr docTR (Document Text Recognition) - a seamless, high-performing & accessible library for OCR-related tasks powered by Deep Learning. 项目地址: https://gitcode.com/gh_mirrors/do/doctr

概述

mindee/doctr项目提供了一个强大的数据集模块，可以帮助开发者快速获取和处理各种文档分析任务所需的数据集。本文将详细介绍如何使用这些数据集，包括预定义数据集和自定义数据集的加载方法。

预定义数据集

doctr提供了多种预定义数据集，覆盖了文档分析的多个方面，包括文本检测、文本识别、OCR和对象检测等任务。

文本检测数据集

文本检测数据集主要用于训练或验证文本检测模型，可以识别文档中文本的位置。以下是可用的数据集：

| 数据集名称 | 训练样本数 | 测试样本数 | 特点说明 | |------------|------------|------------|----------| | FUNSD | 149 | 50 | 表单理解数据集 | | SROIE | 626 | 360 | 收据信息提取 | | CORD | 800 | 100 | 收据数据集 | | IIIT5K | 2000 | 3000 | 场景文本数据集 | | SVT | 100 | 249 | 街景文本数据集 | | SVHN | 33402 | 13068 | 门牌号数据集 | | SynthText | 772875 | 85875 | 合成文本数据集 |

使用示例：

from doctr.datasets import CORD
# 加载矩形框
train_set = CORD(train=True, download=True, detection_task=True)
# 加载旋转框
train_set = CORD(train=True, download=True, use_polygons=True, detection_task=True)
img, target = train_set[0]

文本识别数据集

文本识别数据集包含文本内容和位置信息，用于训练或验证文本识别模型：

| 数据集名称 | 训练样本数 | 测试样本数 | 特点说明 | |------------|------------|------------|----------| | FUNSD | 21888 | 8707 | 英文文本 | | SROIE | 33608 | 19342 | 英文/仅大写 | | CORD | 19370 | 2186 | 英文文本 | | IIIT5K | 2000 | 3000 | 英文文本 | | SVT | 257 | 647 | 英文/仅大写 | | SVHN | 73257 | 26032 | 数字识别 |

使用示例：

from doctr.datasets import CORD
# 不规则裁剪
train_set = CORD(train=True, download=True, recognition_task=True)
# 旋转框裁剪(规则)
train_set = CORD(train=True, download=True, use_polygons=True, recognition_task=True)
img, target = train_set[0]

OCR数据集

OCR数据集同时包含文本检测和识别所需的信息，可以用于端到端的OCR任务训练：

from doctr.datasets import CORD
# 加载矩形框
train_set = CORD(train=True, download=True)
# 加载旋转框
train_set = CORD(train=True, download=True, use_polygons=True)
img, target = train_set[0]

对象检测数据集

用于训练文档中特定对象的检测模型，如二维码、条形码等：

| 数据集名称 | 训练样本数 | 测试样本数 | 检测类别 | |------------|------------|------------|----------| | DocArtefacts | 2700 | 300 | ["background", "qr_code", "bar_code", "logo", "photo"] |

使用示例：

from doctr.datasets import DocArtefacts
train_set = DocArtefacts(train=True, download=True)
img, target = train_set[0]

合成数据生成器

doctr还提供了合成数据生成器，可以动态生成训练数据：

# 字符生成器
from doctr.datasets import CharacterGenerator
ds = CharacterGenerator(vocab='abdef', num_samples=100)
img, target = ds[0]

# 单词生成器
from doctr.datasets import WordGenerator
ds = WordGenerator(vocab='abdef', min_chars=1, max_chars=32, num_samples=100)
img, target = ds[0]

使用自定义数据集

对于包含敏感信息的文档数据，可以使用自定义数据集格式：

from doctr.datasets import DetectionDataset, RecognitionDataset, OCRDataset

# 加载检测数据集
train_set = DetectionDataset(img_folder="/path/to/images", label_path="/path/to/labels.json")

# 加载识别数据集
train_set = RecognitionDataset(img_folder="/path/to/images", labels_path="/path/to/labels.json")

# 加载OCR数据集
train_set = OCRDataset(img_folder="/path/to/images", label_file="/path/to/labels.json")

img, target = train_set[0]

数据加载器

doctr提供了DataLoader来方便地批量加载数据：

from doctr.datasets import CORD, DataLoader

train_set = CORD(train=True, download=True)
train_loader = DataLoader(train_set, batch_size=32)
train_iter = iter(train_loader)
images, targets = next(train_iter)

总结

mindee/doctr的数据集模块提供了丰富的预定义数据集和灵活的接口，可以大大简化文档分析任务的准备工作。无论是使用现成数据集还是自定义数据，都能找到合适的解决方案。通过合理利用这些工具，开发者可以更专注于模型的设计和优化。

doctr docTR (Document Text Recognition) - a seamless, high-performing & accessible library for OCR-related tasks powered by Deep Learning. 项目地址: https://gitcode.com/gh_mirrors/do/doctr