库名称: transformers
许可证: apache-2.0
任务类型: 图像分割
数据集:

• ds4sd/DocLayNet

DIT基础版布局检测模型

我们推出cmarkea/dit-base-layout-detection模型，可从文档图像中提取不同布局元素（文本、图片、标题、脚注等）。该模型基于microsoft/dit-base在DocLayNet数据集上微调而成，特别适合处理需要导入开放域问答系统(ODQA)的文档集。

本模型可识别11类元素：图片说明、脚注、公式、列表项、页脚、页眉、图片、章节标题、表格、正文文本和标题。

性能表现

我们分别评估语义分割和目标检测性能。语义分割未进行后处理，目标检测仅使用OpenCV的findContours函数而未做额外处理。

语义分割采用F1分数评估像素分类，目标检测则通过广义交并比(GIoU)和预测边界框类别准确率衡量。评估基于DocLayNet的500页PDF测试集。

基准对比

与其他模型的性能对比：

直接使用

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"""
    简易边界框检测函数
    """
    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:
            # 包含所有点的最小矩形
            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))

示例

引用

@online{DeDitLay,
  AUTHOR = {Cyrile Delestre},
  URL = {https://huggingface.co/cmarkea/dit-base-layout-detection},
  YEAR = {2024},
  KEYWORDS = {图像处理 ; 变换器 ; 布局},
}