Dit Base Layout Detection

🚀 DIT-base-layout-detection

DIT-base-layout-detection 模型（cmarkea/dit-base-layout-detection）可从文档图像中提取不同的布局元素，如文本、图片、标题、脚注等。该模型基于 dit-base 在 DocLayNet 数据集上进行微调，非常适合处理要导入开放域问答（ODQA）系统的文档语料。此模型能够提取 11 种实体，包括标题、脚注、公式、列表项、页面页脚、页面页眉、图片、章节标题、表格、文本和标题。

✨ 主要特性

• 可从文档图像中精准提取多种布局元素。
• 基于大规模数据集微调，适用于处理文档语料。
• 能提取 11 种不同类型的实体。

📦 安装指南

由于文档未提供具体安装步骤，此部分跳过。

💻 使用示例

基础用法

import torch
from transformers import AutoImageProcessor, BeitForSemanticSegmentation

img_proc = AutoImageProcessor.from_pretrained(
    "cmarkea/dit-base-layout-detection"
)
model = BeitForSemanticSegmentation.from_pretrained(
    "cmarkea/dit-base-layout-detection"
)

img: PIL.Image

with torch.inference_mode():
    input_ids = img_proc(img, return_tensors='pt')
    output = model(**input_ids)

segmentation = img_proc.post_process_semantic_segmentation(
    output,
    target_sizes=[img.size[::-1]]
)

高级用法

import cv2

def detect_bboxes(masks: np.ndarray):
    r"""
    A simple bounding box detection function
    """
    detected_blocks = []
    contours, _ = cv2.findContours(
        masks.astype(np.uint8),
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )
    for contour in list(contours):
        if len(list(contour)) >= 4:
            # smallest rectangle containing all points
            x, y, width, height = cv2.boundingRect(contour)
            bounding_box = [x, y, x + width, y + height]
            detected_blocks.append(bounding_box)
    return detected_blocks

bbox_pred = []
for segment in segmentation:
    boxes, labels = [], []
    for ii in range(1, len(model.config.label2id)):
        mm = segment == ii
        if mm.sum() > 0:
            bbx = detect_bboxes(mm.numpy())
            boxes.extend(bbx)
            labels.extend([ii]*len(bbx))
    bbox_pred.append(dict(boxes=boxes, labels=labels))

📚 详细文档

性能评估

在这部分，我们将分别从语义分割和目标检测两个方面评估模型的性能。对于语义分割，我们未进行任何后处理；对于目标检测，仅应用了 OpenCV 的 findContours 函数，未进行进一步的后处理。

在语义分割方面，我们使用 F1 分数来评估每个像素的分类情况；在目标检测方面，我们基于广义交并比（GIoU）和预测边界框类别的准确性来评估性能。评估在 DocLayNet 的 PDF 评估数据集的 500 页上进行。

基准测试

现在，让我们将该模型的性能与其他模型进行比较。

示例

🔧 技术细节

文档未提供具体技术细节，此部分跳过。

📄 许可证

本项目采用 Apache-2.0 许可证。

引用

@online{DeDitLay,
  AUTHOR = {Cyrile Delestre},
  URL = {https://huggingface.co/cmarkea/dit-base-layout-detection},
  YEAR = {2024},
  KEYWORDS = {Image Processing ; Transformers ; Layout},
}