PubTabNet数据集介绍(有干货)

前言

PubTabNet是IBM公司公布的基于图像的表格识别数据集。
其包含了568k+表格图片,其标注数据是HTML的表格结构,下载压缩包磁盘存储大小10G+。
GitHub相关地址
IBM的下载地址
相关论文:
Image-based table recognition: data, model, and evaluation
此篇论文的核心在于通过Encoder - double Decoder实现表格结构与单元格的识别

2020-12-24 15_28_04-Start.png

目前暂时没有GitHub复现项目。

PubTabNet数据

表格图片

该数据集的表格都是PDF截图,清晰度不是很高
示例如下:


PMC2838834_005_00.png

训练结构数据

训练结构数据位于.jsonl文件中,该文件每一行都是一条json数据,可以通过jsonlines库逐条读取。
数据结构:

  • imgid: 图像id
  • html: html单元格详细描述
    cells:单元格列表
    cells列表成员:
    tokens:单元格文本字符级别信息,
    bbox:单元格文本范围的bounding box,这个并不是单元格范围,而是单元格内文本范围!数字坐标说明:x_min, y_min, x_max, y_max,即文本范围对角线两个点的坐标。
    该坐标就是图片本身的坐标,可以放心食用
    structure字典
    structure字典成员:
    tokens:对应cells的HTML格式
  • split:表示训练数据或验证数据,分别为train与val
  • filename:图像名称,如PMC2838834_005_00.png

示例如下(只贴出两个单元格以及部分HTML结构数据):

{
  "imgid": 4,
  "html": {
    "cells": [
      {
        "tokens": [
          "<b>",
          "M",
          "a",
          "i",
          "n",
          " ",
          "c",
          "e",
          "l",
          "l",
          "u",
          "l",
          "a",
          "r",
          " ",
          "p",
          "r",
          "o",
          "c",
          "e",
          "s",
          "s",
          "</b>"
        ],
        "bbox": [
          1,
          4,
          76,
          13
        ]
      },
      {
        "tokens": [
          "<b>",
          "M",
          "o",
          "d",
          "u",
          "l",
          "a",
          "t",
          "e",
          "d",
          " ",
          "p",
          "a",
          "t",
          "h",
          "w",
          "a",
          "y",
          "s",
          "</b>"
        ],
        "bbox": [
          92,
          4,
          167,
          13
        ]
      }
    ],
    "structure": {
      "tokens": [
        "<thead>",
        "<tr>",
        "<td>",
        "</td>",
        "<td>",
        "</td>",
        "<td",
        " colspan=\"2\"",
        ">",
        "</td>",
        "<td",
        " colspan=\"3\"",
        ">",
        "</td>",
        "</tr>",
        "...",
        "..."
      ]
    }
  },
  "split": "train",
  "filename": "PMC2838834_005_00.png"
}

位于GitHun的PubTabNet相关代码,只有读取数据,将数据转换为HTML的功能,并没有表格识别相关的代码。

PubTabNet转换为SciTSR训练数据格式

因为之前在研究SciTSR数据,所以需要将PubTabNet的训练集数据转换为SciTSR的格式,从而扩大训练集,即我们需要根据.jsonl中的数据,提炼出chunk, structure数据来。
SciTSR的介绍链接在此
这一段是干货,因为这意味着不同数据集可以通用了。

import tqdm
import os
import jsonlines
import json
import logging
import sys
import re

sys.path.insert(0, os.path.abspath('../'))


class Transform:
    def transform(self, pubtabnet_data_file: str):
        if pubtabnet_data_file is None or not os.path.exists(pubtabnet_data_file):
            logging.error('No PubTabNet data file!')
            return
        root_path = os.path.dirname(pubtabnet_data_file)

        with open(pubtabnet_data_file, encoding='utf-8') as reader:
            for img in jsonlines.Reader(reader):
                # if img['filename'] == 'PMC5577841_001_00.png':
                # print(img)
                img_filename = img['filename']

                img_type = img['split']
                gfte_structure_folder = '{0}/{1}/structure/'.format(root_path, img_type)
                if not os.path.exists(gfte_structure_folder):
                    os.makedirs(gfte_structure_folder)
                gfte_chunk_folder = '{0}/{1}/chunk/'.format(root_path, img_type)
                if not os.path.exists(gfte_chunk_folder):
                    os.makedirs(gfte_chunk_folder)

                print('Handle image: {0}'.format(img_filename))
                cells = img['html']['cells']
                structure_list = img['html']['structure']['tokens']
                gfte_structure = self.get_row_col_position(structure_list)

                gfte_structure_file = '{0}.json'.format(img_filename.replace('.png', ''))
                gfte_chunk_file = '{0}.chunk'.format(img_filename.replace('.png', ''))
                structure_data_dict, chunk_data_dict = self.transfer_GFTE_data(cells, gfte_structure)
                if isinstance(structure_data_dict, dict) and \
                        isinstance(chunk_data_dict, dict):
                    with open(os.path.join(gfte_structure_folder,
                                           gfte_structure_file), "w", encoding='utf-8') as write_file:
                        json.dump(structure_data_dict, write_file)

                    with open(os.path.join(gfte_chunk_folder,
                                           gfte_chunk_file), "w", encoding='utf-8') as write_file:
                        json.dump(chunk_data_dict, write_file)

    def get_row_col_position(self, structure_list: list):
        row_index = 0
        col_index = 0
        gfte_structure = []
        for index, table_element in enumerate(structure_list):
            if table_element == '<tr>':
                col_index = 0
            if table_element == '</tr>':
                row_index += 1
            if table_element == '<td>':
                structure_dict = {'start_row': row_index,
                                  'end_row': row_index,
                                  'start_col': col_index,
                                  'end_col': col_index}
                col_index += 1
                gfte_structure.append(structure_dict)
                continue
            if table_element == '<td' and index < len(structure_list) - 1:
                next_element = structure_list[index + 1]
                next_element_split = next_element.strip().replace('"', '').split('=')
                if len(next_element_split) == 2:
                    span_number = int(next_element_split[1])
                    if next_element_split[0] == 'colspan':
                        structure_dict = {'start_row': row_index,
                                          'end_row': row_index,
                                          'start_col': col_index,
                                          'end_col': col_index + span_number - 1}
                        col_index += span_number
                        gfte_structure.append(structure_dict)
                        continue
                    if next_element_split[0] == 'rowspan':
                        structure_dict = {'start_row': row_index,
                                          'end_row': row_index + span_number - 1,
                                          'start_col': col_index,
                                          'end_col': col_index}
                        col_index += 1
                        gfte_structure.append(structure_dict)
                        continue
        return gfte_structure

    def transfer_GFTE_data(self, cell_list: list, structure_list: list):
        logging.info('Construct GFTE structure data')
        if cell_list is None or \
                structure_list is None or \
                not isinstance(cell_list, list) or \
                not isinstance(structure_list, list) or \
                len(cell_list) != len(structure_list):
            logging.error('Cell is not pair of structure!')
            return -1, -1
        structure_data_dict = {'cells': []}
        chunk_data_dict = {'chunks': []}
        for index, (cell, structure) in enumerate(zip(cell_list, structure_list)):
            structure_data = {}
            structure_data['id'] = index
            raw_text = ''.join(cell['tokens']).strip()
            structure_data['tex'] = raw_text
            pure_text = re.sub(r'\<.*?\>', '', raw_text).strip()
            structure_data['content'] = pure_text.split()
            structure_data.update(structure)
            structure_data_dict['cells'].append(structure_data)

            chunk_data = {}
            if cell.get('bbox', None) is not None:
                chunk_data['pos'] = [cell['bbox'][0],
                                     cell['bbox'][2],
                                     cell['bbox'][1],
                                     cell['bbox'][3]]
            else:
                chunk_data['pos'] = [0, 0, 0, 0]
            chunk_data['text'] = raw_text
            chunk_data_dict['chunks'].append(chunk_data)
        for index, (structure_data, chunk_data) in enumerate(zip(structure_data_dict['cells'],
                                                                 chunk_data_dict['chunks'])):
            # x0, x1, y0, y1
            if chunk_data['pos'] == [0, 0, 0, 0]:
                id = structure_data['id']
                start_row = structure_data['start_row']
                end_row = structure_data['end_row']
                start_col = structure_data['start_col']
                end_col = structure_data['end_col']
                x0 = None
                x1 = None
                y0 = None
                y1 = None
                for structure, chunk in zip(structure_data_dict['cells'],
                                            chunk_data_dict['chunks']):
                    if structure['id'] != id and chunk['pos'] != [0, 0, 0, 0]:
                        if y0 is None and structure['start_row'] == start_row:
                            y0 = chunk['pos'][2]
                        if y1 is None and structure['end_row'] == end_row:
                            y1 = chunk['pos'][3]
                        if x0 is None and structure['start_col'] == start_col:
                            x0 = chunk['pos'][0]
                        if x1 is None and structure['end_col'] == end_col:
                            x1 = chunk['pos'][1]
                    if x0 is not None and x1 is not None and y0 is not None and y1 is not None:
                        chunk_data['pos'] = [x0, x1, y0, y1]
                        break
        return structure_data_dict, chunk_data_dict


if __name__ == "__main__":
    transform = Transform()
    transform.transform('/data/pubtabnet/PubTabNet_2.0.0.jsonl')
    # transform.transform('/data/scitsr/examples/PubTabNet_Examples.jsonl')

UNET所需数据准备

UNET非常适合图像分割,对于表格来说,如果能够通过UNET将表格中行与列进行标注,则能够方便根据表格结构提取各个单元格信息。
事实上,UNET预测的是“线”,训练数据类似于如下格式,即两点一线,标签为0代表横线,标签为1代表竖线

{
   "label": "0",
   "line_color": [
    0,
    0,
    128
   ],
   "fill_color": [
    0,
    0,
    128
   ],
   "points": [
    [
     0,
     0
    ],
    [
     503,
     0
    ]
   ]
  },
  {
   "label": "0",
   "line_color": [
    0,
    0,
    128
   ],
   "fill_color": [
    0,
    0,
    128
   ],
   "points": [
    [
     0,
     275
    ],
    [
     503,
     275
    ]
   ]
  },
  {
   "label": "1",
   "line_color": [
    0,
    0,
    128
   ],
   "fill_color": [
    0,
    0,
    128
   ],
   "points": [
    [
     0,
     0
    ],
    [
     0,
     276
    ]
   ]
  }

训练数据转换代码,则是根据SciTSR的训练数据:chunk与structure得到。
如下代码实现了:

  1. 根据每行每列的最大与最小坐标,以及邻接单元格信息,尝试绘制表格中的横线与竖线。而这仅仅是根据原数据中的单元格中文本框的坐标计算得到的
  2. 将得到的线绘制到单独一张图片,构成训练数据示意图片
import os
import numpy as np
from PIL import Image
import cv2
import json
from glob import glob
import traceback


# Returns if columns belong to same table or not
def sameTable(ymin_1, ymin_2, ymax_1, ymax_2):
    min_diff = abs(ymin_1 - ymin_2)
    max_diff = abs(ymax_1 - ymax_2)

    if min_diff <= 5 and max_diff <= 5:
        return True
    elif min_diff <= 4 and max_diff <= 7:
        return True
    elif min_diff <= 7 and max_diff <= 4:
        return True
    return False


def draw_cell_line(root_path: str = r'/data/scitsr/train',
                   img_path: str = 'img',
                   is_scitsr: bool = True,
                   output_unet: bool = False,
                   draw_img_amount: int = None,
                   file_list: list = None,
                   is_middle: bool = False,
                   draw_thickness: bool = False,
                   draw_line_img: bool = True):
    """
    依赖PubTabNet的图像数据,以及SciTSR的数据格式,转换为table_net需要的数据结构
    SciTSR的数据格式,通过pubtabnet_format_transform.py的transform方法进行转换
    即绘制横线与竖线,横线标签为0,竖线标签为1
    并将数据保存为json文件
    :param root_path:
    :param is_scitsr:
    :param output_unet:
    :param draw_img_amount:
    :param file_list:
    :param is_middle: 是否从两个文本区域中间划线,这种绘线方式与实际表格线段更贴近
    :return:
    """
    directory = os.path.join(root_path, img_path)
    chunk_directory = os.path.join(root_path, 'chunk')
    structure_directory = os.path.join(root_path, 'structure')

    if is_middle:
        final_cell_directory = os.path.join(root_path, 'cell_mask_img_middle')
        final_table_cell_directory = os.path.join(root_path, 'table_cell_mask_img_middle')
        table_net_data_directory = os.path.join(root_path, 'table_net_data_middle')
    else:
        final_cell_directory = os.path.join(root_path, 'cell_mask_img')
        final_table_cell_directory = os.path.join(root_path, 'table_cell_mask_img')
        table_net_data_directory = os.path.join(root_path, 'table_net_data')

    if not os.path.exists(final_cell_directory):
        os.makedirs(final_cell_directory)
    if not os.path.exists(final_table_cell_directory):
        os.makedirs(final_table_cell_directory)
    if not os.path.exists(table_net_data_directory):
        os.makedirs(table_net_data_directory)
    final_table_directory = os.path.join(root_path, 'table_mask_img')
    if not os.path.exists(final_table_directory):
        os.makedirs(final_table_directory)

    unet_img_directory = os.path.join(root_path, 'img_unet')
    if output_unet:
        if not os.path.exists(unet_img_directory):
            os.makedirs(unet_img_directory)
        final_cell_directory = os.path.join(root_path, 'mask_unet')
        if not os.path.exists(final_cell_directory):
            os.makedirs(final_cell_directory)
    if file_list is None:
        files = os.listdir(directory)
    else:
        files = file_list
    count = 1
    table_net_data_file_list = []
    for index, file in enumerate(files):
        if file_list is None and \
                draw_img_amount is not None and \
                isinstance(draw_img_amount, int) and \
                draw_img_amount > 0 and \
                count == draw_img_amount:
            break
        print('Handle the {0} file: {1}'.format(count, file))
        try:
            file_path = os.path.join(directory, file)
            chunk_file = os.path.join(chunk_directory,
                                      file.replace('.png', '.chunk').replace('.jpg', '.chunk'))
            structure_file = os.path.join(structure_directory,
                                          file.replace('.png', '.json').replace('.jpg', '.json'))
            table_net_data_file = os.path.join(table_net_data_directory,
                                               file.replace('.png', '.json').replace('.jpg', '.json'))
            if not os.path.exists(chunk_file) or not os.path.exists(structure_file):
                print('No structure data for {0}'.format(file))
                continue
            count += 1
            img = cv2.imread(file_path)
            height, width = img.shape[:2]
            # Create grayscale image array
            col_mask = np.ones((height, width), dtype=np.int32) * 255

            with open(chunk_file, encoding='utf-8') as f:
                chunk_data_list = json.load(f)['chunks']
            with open(structure_file, encoding='utf-8') as f:
                structure_data = json.load(f)
                structure_data_list = structure_data['cells']

            # is_changed = fix_structure_data(structure_data_list)
            # if is_changed:
            #     with open(structure_file,"w", encoding='utf-8') as json_file:
            #         json.dump(structure_data, json_file, ensure_ascii=False)

            table_xmin = 0
            table_xmax = width
            table_ymin = 0
            table_ymax = height

            # width_ratio = (table_xmax - table_xmin) / width
            # height_ration = (table_ymax - table_ymin) / height

            row_col_bound_dict = {'row': {}, 'col': {}}
            # 先得到每行的ymin与ymax,每列的xmin与xmax
            for s_index, structure in enumerate(structure_data_list):
                index = structure['id']
                bndbox = chunk_data_list[index]['pos']
                xmin = int(int(bndbox[0]) - table_xmin)
                xmax = int(int(bndbox[1]) - table_xmin)
                # SciTSR的坐标,x0, x1是顺序的,但是y0, y1是倒序的,即坐标系y轴是从下往上的。
                # 所以相对坐标应该是,cell_ymin = -(y2 - ymax) cell_ymax = -(y1 - ymax)
                if is_scitsr:
                    ymin = int(abs(int(bndbox[3]) - table_ymax))
                    ymax = int(abs(int(bndbox[2]) - table_ymax))
                else:
                    ymin = int((int(bndbox[2]) - table_ymin))
                    ymax = int((int(bndbox[3]) - table_ymin))

                start_row = structure['start_row']
                if row_col_bound_dict['row'].get(start_row, None) is None:
                    row_col_bound_dict['row'][start_row] = {'ymin': ymin}
                else:
                    if row_col_bound_dict['row'][start_row].get('ymin', None) is None:
                        row_col_bound_dict['row'][start_row]['ymin'] = ymin
                    else:
                        if row_col_bound_dict['row'][start_row]['ymin'] > ymin:
                            row_col_bound_dict['row'][start_row]['ymin'] = ymin

                end_row = structure['end_row']
                if row_col_bound_dict['row'].get(end_row, None) is None:
                    row_col_bound_dict['row'][end_row] = {'ymax': ymax}
                else:
                    if row_col_bound_dict['row'][end_row].get('ymax', None) is None:
                        row_col_bound_dict['row'][end_row]['ymax'] = ymax
                    else:
                        if row_col_bound_dict['row'][end_row]['ymax'] < ymax:
                            row_col_bound_dict['row'][end_row]['ymax'] = ymax

                start_col = structure['start_col']
                if row_col_bound_dict['col'].get(start_col, None) is None:
                    row_col_bound_dict['col'][start_col] = {'xmin': xmin}
                else:
                    if row_col_bound_dict['col'][start_col].get('xmin', None) is None:
                        row_col_bound_dict['col'][start_col]['xmin'] = xmin
                    else:
                        if row_col_bound_dict['col'][start_col]['xmin'] > xmin:
                            if start_col - 1 >= 0:
                                if row_col_bound_dict['col'].get(start_col - 1, None) is not None:
                                    if row_col_bound_dict['col'][start_col - 1].get('max', None) is not None and \
                                            row_col_bound_dict['col'][start_col - 1]['xmax'] < xmin:
                                        row_col_bound_dict['col'][start_col]['xmin'] = xmin
                                    else:
                                        row_col_bound_dict['col'][start_col]['xmin'] = xmin
                                else:
                                    row_col_bound_dict['col'][start_col]['xmin'] = xmin

                end_col = structure['end_col']
                if row_col_bound_dict['col'].get(end_col, None) is None:
                    row_col_bound_dict['col'][end_col] = {'xmax': xmax}
                else:
                    if row_col_bound_dict['col'][end_col].get('xmax', None) is None:
                        row_col_bound_dict['col'][end_col]['xmax'] = xmax
                    else:
                        if row_col_bound_dict['col'][end_col]['xmax'] < xmax:
                            row_col_bound_dict['col'][end_col]['xmax'] = xmax

            for row_index, location in row_col_bound_dict['row'].items():
                if location.get('ymin', None) is not None and \
                        location.get('ymax', None) is None and \
                        row_col_bound_dict['row'].get(row_index + 1, None) is not None and \
                        row_col_bound_dict['row'].get(row_index + 1, {}).get('ymin', None) is not None:
                    next_ymin = row_col_bound_dict['row'].get(row_index + 1, {}).get('ymin', None)
                    current_ymin = location.get('ymin', None)
                    if next_ymin > current_ymin + 5:
                        location['ymax'] = next_ymin - 5
                    else:
                        location['ymax'] = next_ymin

            for col_index, location in row_col_bound_dict['col'].items():
                if location.get('xmin', None) is not None and \
                        location.get('xmax', None) is None and \
                        row_col_bound_dict['col'].get(col_index + 1, None) is not None and \
                        row_col_bound_dict['col'].get(col_index + 1, {}).get('xmin', None) is not None:
                    next_xmin = row_col_bound_dict['col'].get(col_index + 1, {}).get('xmin', None)
                    current_xmin = location.get('xmin', None)
                    if next_xmin > current_xmin + 5:
                        location['xmax'] = next_xmin - 5
                    else:
                        location['xmax'] = next_xmin
                if location.get('xmin', None) is not None and \
                        location.get('xmax', None) is not None and \
                        row_col_bound_dict['col'].get(col_index + 1, None) is not None and \
                        row_col_bound_dict['col'].get(col_index + 1, {}).get('xmin', None) is not None:
                    current_xmax = location.get('xmax', None)
                    next_xmin = row_col_bound_dict['col'].get(col_index + 1, {}).get('xmin', None)
                    if next_xmin < current_xmax + 5:
                        location['xmax'] = next_xmin - 5

            table_net_data = {'version': '3.16.7',
                              'flags': {},
                              'lineColor': [0, 255, 0, 128],
                              'fillColor': [255, 0, 0, 128],
                              'imagePath': file_path,
                              'shapes': []}
            # 首先添加表边框四条线
            table_net_data['shapes'].append({'label': '0',
                                             'line_color': [0, 0, 128],
                                             'fill_color': [0, 0, 128],
                                             'points': [[0, 0], [width, 0]],
                                             'thickness': 2})
            table_net_data['shapes'].append({'label': '0',
                                             'line_color': [0, 0, 128],
                                             'fill_color': [0, 0, 128],
                                             'points': [[0, height - 1],
                                                        [width, height - 1]],
                                             'thickness': 2})
            table_net_data['shapes'].append({'label': '1',
                                             'line_color': [0, 0, 128],
                                             'fill_color': [0, 0, 128],
                                             'points': [[0, 0],
                                                        [0, height]],
                                             'thickness': 2})
            table_net_data['shapes'].append({'label': '1',
                                             'line_color': [0, 0, 128],
                                             'fill_color': [0, 0, 128],
                                             'points': [[width - 1, 0],
                                                        [width - 1, height]],
                                             'thickness': 2})

            row_list = list(row_col_bound_dict.get('row', {}).keys())
            row_list.sort()
            col_list = list(row_col_bound_dict.get('col', {}).keys())
            col_list.sort()
            # 只添加结束行横线
            for row in row_list:
                if row == row_list[-1]:
                    continue
                same_row_list = [structure for structure
                                 in structure_data_list
                                 if structure['start_row'] == row]

                line_points = None
                same_row_list = sorted(same_row_list, key=lambda keys: keys['start_col'])
                thickness = None
                for cell in same_row_list:
                    end_row = cell['end_row']
                    if end_row == row_list[-1]:
                        continue

                    start_col = cell['start_col']
                    end_col = cell['end_col']

                    if start_col == 0:
                        xmin = 0
                    else:
                        try:
                            if is_middle:
                                xmin = round((row_col_bound_dict['col'][start_col - 1]['xmax'] +
                                              row_col_bound_dict['col'][start_col]['xmin']) / 2)
                                if xmin < row_col_bound_dict['col'][start_col - 1]['xmax']:
                                    xmin = row_col_bound_dict['col'][start_col - 1]['xmax']
                            else:
                                xmin = row_col_bound_dict['col'][start_col - 1]['xmax']
                        except:
                            xmin = row_col_bound_dict['col'][start_col - 1]['xmax']

                    if end_col == len(col_list) - 1:
                        xmax = table_xmax
                    else:
                        try:
                            if is_middle:
                                xmax = round((row_col_bound_dict['col'][end_col]['xmax'] +
                                              row_col_bound_dict['col'][end_col + 1]['xmin']) / 2)
                                if xmax < row_col_bound_dict['col'][end_col]['xmax']:
                                    xmax = row_col_bound_dict['col'][end_col]['xmax']
                            else:
                                xmax = row_col_bound_dict['col'][end_col]['xmax']
                        except:
                            xmax = row_col_bound_dict['col'][end_col]['xmax']

                    try:
                        if is_middle:
                            y = round((row_col_bound_dict['row'][end_row]['ymax'] +
                                       row_col_bound_dict['row'][end_row + 1]['ymin']) / 2)
                            temp = abs(int(row_col_bound_dict['row'][end_row + 1]['ymin'] -
                                           row_col_bound_dict['row'][end_row]['ymax']))
                            if temp < 2:
                                temp = 2
                            if thickness is None:
                                thickness = temp
                            elif thickness > temp:
                                thickness = temp
                            if y < row_col_bound_dict['row'][end_row]['ymax']:
                                y = row_col_bound_dict['row'][end_row]['ymax']
                        else:
                            y = row_col_bound_dict['row'][end_row]['ymax']
                            thickness = int(row_col_bound_dict['row'][end_row]['ymax'] -
                                            row_col_bound_dict['row'][end_row + 1]['ymin'])
                    except:
                        y = row_col_bound_dict['row'][end_row]['ymax']

                    if end_row == row:
                        if line_points is None:
                            line_points = [[xmin, y], [xmax, y]]
                        else:
                            line_points[1][0] = xmax
                    else:
                        if line_points is not None:
                            table_net_data['shapes'].append({'label': '0',
                                                             'line_color': [0, 0, 128],
                                                             'fill_color': [0, 0, 128],
                                                             'points': line_points,
                                                             'thickness': thickness})
                        line_points = [[xmin, y], [xmax, y]]
                        table_net_data['shapes'].append({'label': '0',
                                                         'line_color': [0, 0, 128],
                                                         'fill_color': [0, 0, 128],
                                                         'points': line_points,
                                                         'thickness': thickness})
                        line_points = None
                if line_points is not None:
                    if thickness is None or thickness < 2:
                        thickness = 2
                    table_net_data['shapes'].append({'label': '0',
                                                     'line_color': [0, 0, 128],
                                                     'fill_color': [0, 0, 128],
                                                     'points': line_points,
                                                     'thickness': thickness})

            # 只添加结束列竖线
            for col in col_list:
                if col == col_list[-1]:
                    continue
                same_col_list = [structure for structure
                                 in structure_data_list
                                 if structure['start_col'] == col]
                same_col_list = sorted(same_col_list, key=lambda keys: keys['start_row'])
                line_points = None
                thickness = None
                for index, cell in enumerate(same_col_list):
                    if index > 0:
                        last_cell_end_row = same_col_list[index - 1]['end_row']
                        cell_start_row = cell['start_row']
                        if cell_start_row - last_cell_end_row != 1 and line_points is not None:
                            table_net_data['shapes'].append({'label': '1',
                                                             'line_color': [0, 0, 128],
                                                             'fill_color': [0, 0, 128],
                                                             'points': line_points})
                            line_points = None
                    end_col = cell['end_col']
                    if end_col == col_list[-1]:
                        if line_points is not None:
                            table_net_data['shapes'].append({'label': '1',
                                                             'line_color': [0, 0, 128],
                                                             'fill_color': [0, 0, 128],
                                                             'points': line_points})
                            line_points = None
                        continue
                    start_row = cell['start_row']
                    end_row = cell['end_row']
                    if start_row - 1 >= 0:
                        try:
                            if is_middle:
                                ymin = round((row_col_bound_dict['row'][start_row - 1]['ymax'] +
                                              row_col_bound_dict['row'][start_row]['ymin']) / 2)
                                if ymin < row_col_bound_dict['row'][start_row - 1]['ymax']:
                                    ymin = row_col_bound_dict['row'][start_row - 1]['ymax']
                            else:
                                ymin = row_col_bound_dict['row'][start_row - 1]['ymax']
                        except:
                            ymin = row_col_bound_dict['row'][start_row - 1]['ymax']
                    else:
                        ymin = 0
                    if end_row == len(row_list) - 1:
                        ymax = table_ymax
                    else:
                        try:
                            if is_middle:
                                ymax = round((row_col_bound_dict['row'][end_row]['ymax'] +
                                              row_col_bound_dict['row'][end_row + 1]['ymin']) / 2)
                                if ymax < row_col_bound_dict['row'][end_row]['ymax']:
                                    ymax = row_col_bound_dict['row'][end_row]['ymax']
                            else:
                                ymax = row_col_bound_dict['row'][end_row]['ymax']
                        except:
                            ymax = row_col_bound_dict['row'][end_row]['ymax']

                    try:
                        if is_middle:
                            x = round((row_col_bound_dict['col'][end_col]['xmax'] +
                                       row_col_bound_dict['col'][end_col + 1]['xmin']) / 2)
                            temp = abs(int(row_col_bound_dict['col'][end_col + 1]['xmin'] -
                                           row_col_bound_dict['col'][end_col]['xmax']))
                            if temp < 2:
                                temp = 2
                            if thickness is None:
                                thickness = temp
                            elif thickness < temp:
                                thickness = temp
                            if x < row_col_bound_dict['col'][end_col]['xmax']:
                                x = row_col_bound_dict['col'][end_col]['xmax']
                        else:
                            x = row_col_bound_dict['col'][end_col]['xmax']
                            thickness = int(row_col_bound_dict['col'][end_col]['xmax'] -
                                            row_col_bound_dict['col'][end_col + 1]['xmin'])
                    except:
                        x = row_col_bound_dict['col'][end_col]['xmax']

                    if end_col == col:
                        if line_points is None:
                            line_points = [[x, ymin], [x, ymax]]
                        else:
                            line_points[1][1] = ymax
                    else:
                        if line_points is not None:
                            table_net_data['shapes'].append({'label': '1',
                                                             'line_color': [0, 0, 128],
                                                             'fill_color': [0, 0, 128],
                                                             'points': line_points,
                                                             'thickness': thickness})
                        line_points = [[x, ymin], [x, ymax]]
                        table_net_data['shapes'].append({'label': '1',
                                                         'line_color': [0, 0, 128],
                                                         'fill_color': [0, 0, 128],
                                                         'points': line_points,
                                                         'thickness': thickness})
                        line_points = None
                if line_points is not None:
                    if thickness is None or thickness < 2:
                        thickness = 2
                    table_net_data['shapes'].append({'label': '1',
                                                     'line_color': [0, 0, 128],
                                                     'fill_color': [0, 0, 128],
                                                     'points': line_points,
                                                     'thickness': thickness})

            fix_table_net_data(table_net_data)
            if draw_line_img:
                draw_line(col_mask, table_net_data['shapes'], draw_thickness=draw_thickness)
                cv2.imwrite(os.path.join(final_cell_directory, file), col_mask)

                # cv2.imwrite(os.path.join(unet_img_directory, file), img)

                draw_line(img, table_net_data['shapes'], draw_thickness=draw_thickness)
                updated = file.replace('.png', '').replace('.jpg', '') + '_line.png'
                cv2.imwrite(os.path.join(final_table_cell_directory, updated), img)
            with open(table_net_data_file, "w", encoding='utf-8') as f:
                json.dump(table_net_data, f, indent=True, ensure_ascii=False)
            table_net_data_file_list.append(table_net_data_file)
        except Exception as e:
            traceback.print_exc()
            # print(e)
    return table_net_data_file_list


def fix_table_net_data(table_net_data):
    lines = table_net_data['shapes']
    row_x1_list = sorted(list(set([line['points'][0][0] for line in lines if line['label'] == '0'])))
    row_x2_list = sorted(list(set([line['points'][1][0] for line in lines if line['label'] == '0'])))
    col_y1_list = sorted(list(set([line['points'][0][1] for line in lines if line['label'] == '1'])))
    col_y2_list = sorted(list(set([line['points'][1][1] for line in lines if line['label'] == '1'])))
    # 修正横线与表格左边框有空隙
    if len(row_x1_list) > 1:
        second_min_x = row_x1_list[1]
        x1_hori_lines = [line
                         for line in
                         table_net_data['shapes']
                         if line['points'][0][0] == second_min_x and
                         line['label'] == '0']
        for h_i, h_line in enumerate(x1_hori_lines):
            h_y = h_line['points'][0][1]
            if not [line for line in lines
                    if line['label'] == '1' and
                       line['points'][0][1] < h_y and
                       line['points'][1][1] > h_y and
                       line['points'][0][0] == second_min_x]:
                h_line['points'][0][0] = 0
    # 修正横线与表格右边框有空隙
    if len(row_x2_list) > 1:
        second_max_x = row_x2_list[-2]
        x2_hori_lines = [line
                         for line in
                         table_net_data['shapes']
                         if line['points'][1][0] == second_max_x and
                         line['label'] == '0']
        for h_i, h_line in enumerate(x2_hori_lines):
            h_y = h_line['points'][0][1]
            if not [line for line in lines
                    if line['label'] == '1' and
                       line['points'][0][1] < h_y and
                       line['points'][1][1] > h_y and
                       line['points'][1][0] == second_max_x]:
                h_line['points'][1][0] = row_x1_list[-1]

    # 修正竖线与表格上边框有空隙
    if len(col_y1_list) > 1:
        second_min_y = col_y1_list[1]
        y1_verti_lines = [line
                         for line in
                         table_net_data['shapes']
                         if line['points'][0][1] == second_min_y and
                         line['label'] == '1']
        for v_i, v_line in enumerate(y1_verti_lines):
            v_x = v_line['points'][0][0]
            if not [line for line in lines
                    if line['label'] == '0' and
                       line['points'][0][0] < v_x and
                       line['points'][1][0] > v_x and
                       line['points'][0][1] == second_min_y]:
                v_line['points'][0][1] = 0

    # 修正竖线与表格下边框有空隙
    if len(col_y2_list) > 1:
        second_max_y = col_y2_list[-2]
        y2_verti_lines = [line
                         for line in
                         table_net_data['shapes']
                         if line['points'][1][1] == second_max_y and
                         line['label'] == '1']
        for v_i, v_line in enumerate(y2_verti_lines):
            v_x = v_line['points'][0][0]
            if not [line for line in lines
                    if line['label'] == '0' and
                       line['points'][0][0] < v_x and
                       line['points'][1][0] > v_x and
                       line['points'][1][1] == second_max_y]:
                v_line['points'][1][1] = col_y2_list[-1]


def fix_structure_data(structrue_list: list):
    if structrue_list is None or len(structrue_list) == 0:
        return []
    is_changed = False
    span_row_start = None
    span_row_end = None
    span_row_cur_col = None
    handled_cells = []
    for index, cell in enumerate(structrue_list):
        if index in handled_cells:
            continue
        start_row = cell['start_row']
        end_row = cell['end_row']
        start_col = cell['start_col']
        end_col = cell['end_col']

        if end_row - start_row > 0:
            if span_row_end is None or span_row_end < end_row:
                span_row_start = start_row
                span_row_end = end_row
                span_row_cur_col = cell['end_col']
                continue
        if span_row_start is not None and span_row_end is not None and \
                (end_row >= span_row_start and end_row <= span_row_end):
            if start_col <= span_row_cur_col:
                count_round = 1
                current_end_row = None
                for idx in range(index, len(structrue_list), 1):
                    current_cell = structrue_list[idx]
                    if current_end_row is None:
                        current_end_row = current_cell['end_row']
                    if current_cell['end_row'] != current_end_row:
                        count_round = 1
                        current_end_row = current_cell['end_row']
                    if current_end_row > span_row_start and current_end_row <= span_row_end:
                        origin_end_start_sub = current_cell['end_col'] - current_cell['start_col']
                        current_cell['start_col'] = span_row_cur_col + count_round
                        current_cell['end_col'] = current_cell['start_col'] + origin_end_start_sub
                        is_changed = True
                        count_round += 1
                        handled_cells.append(idx)
                    else:
                        span_row_start = None
                        span_row_end = None
                        span_row_cur_col = None
                        break
        else:
            span_row_start = None
            span_row_end = None
            span_row_cur_col = None
    return is_changed


def get_span_row_col_img():
    root_path = r'/data/pubtabnet/train/'
    structure_directory = os.path.join(root_path, 'structure/*.json')
    structure_files = glob(structure_directory)
    span_row_img_file = os.path.join(root_path, 'span_row_img.txt')
    span_col_img_file = os.path.join(root_path, 'span_col_img.txt')

    span_row_list = []
    span_col_list = []
    for file in structure_files:
        with open(file, encoding='utf-8') as f:
            structure_data_list = json.load(f)['cells']
            if len(span_row_list) < 100:
                with_span_row = False
                for cell in structure_data_list:
                    if cell['end_row'] - cell['start_row'] > 2:
                        with_span_row = True
                        break
                if with_span_row:
                    span_row_list.append(os.path.basename(file).replace('.json', '.png') + '\n')
                    print('Find the {0} span row file: {1}'.format(len(span_row_list),
                                                                   os.path.basename(file).replace('.json', '.png')))

            if len(span_col_list) < 100:
                with_span_col = False
                for cell in structure_data_list:
                    if cell['end_col'] - cell['start_col'] > 2:
                        with_span_col = True
                        break
                if with_span_col:
                    span_col_list.append(os.path.basename(file).replace('.json', '.png') + '\n')
                    print('Find the {0} span col file: {1}'.format(len(span_col_list),
                                                                   os.path.basename(file).replace('.json', '.png')))
            if len(span_row_list) >= 100 and len(span_col_list) >= 100:
                break
    with open(span_row_img_file, mode='w', encoding='utf-8') as write:
        write.writelines(span_row_list)

    with open(span_col_img_file, mode='w', encoding='utf-8') as write:
        write.writelines(span_col_list)


def draw_line(col_mask, line_list: list, draw_thickness: bool = False):
    for line in line_list:
        points = line['points']
        thickness = line.get('thickness', 2)
        xmin = points[0][0]
        ymin = points[0][1]
        xmax = points[1][0]
        ymax = points[1][1]
        if thickness <= 2 or not draw_thickness:
            if thickness > 2:
                thickness = 2
            cv2.line(col_mask,
                     (xmin, ymin),
                     (xmax, ymax),
                     (0, 0, 128),
                     thickness=thickness,
                     lineType=0)
        else:
            if thickness > 4:
                thickness = 4
            if xmax == xmin:
                # 处理竖线
                xmin = int(xmin - (thickness / 2))
                xmax = int(xmax + (thickness / 2))
            elif ymax == ymin:
                # 处理横线
                ymin = int(ymin - (thickness / 2))
                ymax = int(ymax + (thickness / 2))
            cv2.rectangle(col_mask, (xmin, ymin), (xmax, ymax), (0, 0, 128), -1)


def draw_cell(col_mask,
              structure_data_list: list,
              row_col_bound_dict: dict,
              is_scitsr: bool,
              table_xmin,
              table_xmax,
              table_ymin,
              table_ymax):
    for structure in structure_data_list:
        start_row = structure['start_row']
        end_row = structure['end_row']
        start_col = structure['start_col']
        end_col = structure['end_col']

        xmin = row_col_bound_dict['col'].get(start_col, {}).get('xmin', None)
        xmax = row_col_bound_dict['col'].get(end_col, {}).get('xmax', None)
        ymin = row_col_bound_dict['row'].get(start_row, {}).get('ymin', None)
        ymax = row_col_bound_dict['row'].get(end_row, {}).get('ymax', None)

        if xmin is not None and xmax is not None and ymin is not None and ymax is not None:
            draw_rectangle(col_mask, xmin, xmax, ymin, ymax)
    if is_scitsr:
        draw_rectangle(col_mask, 0, table_xmax - table_xmin, 0, abs(table_ymin - table_ymax))
    else:
        draw_rectangle(col_mask, 0, table_xmax - table_xmin, 0, table_ymax - table_ymin)


def draw_rectangle(img, xmin, xmax, ymin, ymax):
    thickness = 2
    # 上横线
    cv2.line(img, (xmin, ymin), (xmax, ymin), (0, 0, 128), lineType=thickness)
    # 左竖线
    cv2.line(img, (xmin, ymin), (xmin, ymax), (0, 0, 128), lineType=thickness)
    # 右竖线
    cv2.line(img, (xmax, ymin), (xmax, ymax), (0, 0, 128), lineType=thickness)
    # 下横线
    cv2.line(img, (xmin, ymax), (xmax, ymax), (0, 0, 128), lineType=thickness)


if __name__ == '__main__':
    draw_cell_line(root_path=r'/data/pubtabnet/val/',
                   is_scitsr=False,
                   output_unet=False,
                   draw_img_amount=9000)

示例

原图如下:


PMC1181812_008_00.png

纯线段图如下:


PMC1181812_008_00.png

将纯线段与原图叠加如下:
PMC1181812_008_00_line.png

效果不是完美,是因为数据源提供的坐标仅仅是各个单元格的文本区域bounding box,之后会进一步完善