使用Layout-Parser加载和可视化COCO格式的版面标注数据-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00187/article/details/148505970

使用Layout-Parser加载和可视化COCO格式的版面标注数据

layout-parser A Unified Toolkit for Deep Learning Based Document Image Analysis 项目地址: https://gitcode.com/gh_mirrors/la/layout-parser

概述

本文将详细介绍如何使用Layout-Parser工具包加载和可视化COCO格式的版面标注数据。COCO(Common Objects in Context)是一种常用的计算机视觉数据集格式，广泛应用于目标检测、实例分割等任务。在文档分析领域，PubLayNet等数据集也采用COCO格式存储版面标注信息。

准备工作

在开始之前，需要完成以下准备工作：

获取PubLayNet数据集，包含标注文件和对应的图像文件
安装必要的Python依赖库

建议使用验证集(validation set)进行初步尝试，因为训练集通常数据量较大。将下载的数据解压后，按照以下目录结构存放：

data/
  publaynet/
    annotations/  # 存放标注JSON文件
    val/         # 存放验证集图像

安装依赖库：

pip install pycocotools layoutparser opencv-python

加载COCO标注数据

首先导入必要的库：

from pycocotools.coco import COCO
import layoutparser as lp
import random 
import cv2

定义一个辅助函数，将COCO格式的标注转换为LayoutParser对象：

def load_coco_annotations(annotations, coco=None):
    """
    将COCO标注转换为LayoutParser布局对象
    
    参数:
        annotations: COCO标注列表
        coco: COCO标注对象实例(可选)，用于将类别ID转换为类别名称
    """
    layout = lp.Layout()
    
    for ele in annotations:
        x, y, w, h = ele['bbox']
        layout.append(
            lp.TextBlock(
                block = lp.Rectangle(x, y, w+x, h+y),
                type  = ele['category_id'] if coco is None else coco.cats[ele['category_id']]['name'],
                id = ele['id']
            )
        )
    return layout

可视化标注数据

加载COCO标注文件并初始化COCO对象：

COCO_ANNO_PATH = 'data/publaynet/annotations/val.json'
COCO_IMG_PATH = 'data/publaynet/val'
coco = COCO(COCO_ANNO_PATH)

定义不同版面元素的颜色映射：

color_map = {
    'text': 'red',
    'title': 'blue',
    'list': 'green',
    'table': 'purple',
    'figure': 'pink',
}

随机选择一张图像进行可视化：

for image_id in random.sample(coco.imgs.keys(), 1):
    image_info = coco.imgs[image_id]
    annotations = coco.loadAnns(coco.getAnnIds([image_id]))
    
    image = cv2.imread(f'{COCO_IMG_PATH}/{image_info["file_name"]}')
    layout = load_coco_annotations(annotations, coco)
    
    viz = lp.draw_box(image, layout, color_map=color_map)
    display(viz)

增强可视化效果

可以进一步丰富可视化信息，例如显示元素ID和类型：

lp.draw_box(image, 
           [b.set(id=f'{b.id}/{b.type}') for b in layout],
           color_map=color_map,
           show_element_id=True, 
           id_font_size=10, 
           id_text_background_color='grey',
           id_text_color='white')

模型预测与结果可视化

Layout-Parser还支持加载预训练模型对图像进行版面分析。以下是使用预训练模型进行预测的示例：

加载预训练模型：

model = lp.Detectron2LayoutModel(
    'lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config',
    extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],
    label_map={0: "text", 1: "title", 2: "list", 3:"table", 4:"figure"}
)

运行预测并可视化结果：

layout_predicted = model.detect(image)

lp.draw_box(image, 
           [b.set(id=f'{b.type}/{b.score:.2f}') for b in layout_predicted],
           color_map=color_map,
           show_element_id=True,
           id_font_size=10,
           id_text_background_color='grey',
           id_text_color='white')