# -*- coding: utf-8 -*-
|
#
|
# @author: lyg, ym
|
# @date: 2025-5-8
|
# @version: 1
|
# @description: docx文档拆分器，根据段落拆分，将图片和表格转换为json数据。
|
import docx
|
import docx.table
|
import json
|
from dataclasses import dataclass
|
from PIL import Image
|
import io
|
import re
|
|
from knowledgebase.doc.image_to_text import ImageToText
|
|
|
@dataclass
|
class ParagraphInfo:
|
"""
|
段落信息
|
:param text: str - 段落文本
|
:param level: int - 段落级别，1-9级标题，0表示正文
|
:param title_no: str - 标题编号，如1.1、1.1.1等
|
"""
|
text: str
|
level: int
|
title_no: str
|
|
@property
|
def full_text(self):
|
"""
|
获取段落完整文本，包含标题编号
|
:return: str - 段落完整文本
|
"""
|
return f"{self.title_no} {self.text}"
|
|
def __init__(self, text: str, level: int):
|
"""
|
段落信息
|
:param text: str - 段落文本
|
:param level: int - 段落级别，1-9级标题，0表示正文
|
"""
|
self.text = text
|
self.level = level
|
self.title_no = ''
|
|
|
class DocSplit:
|
"""
|
docx文档拆分器，根据段落拆分，将图片和表格转换为json数据。
|
1.封装段落信息
|
2.将图片和表格转换为json
|
3.将段落按照文档标题级别组合成树形结构
|
|
"""
|
|
def __init__(self, doc_file):
|
self.doc_file = doc_file
|
self.image_to_text = ImageToText()
|
self.paragraphs:list[ParagraphInfo] = []
|
|
def table_to_json(self, table: docx.table.Table):
|
"""
|
将表格转换为 JSON 格式
|
|
:param table: docx.table.Table - 要转换的表格对象
|
:return list - 表格数据，以 JSON 格式表示
|
"""
|
table_data = []
|
headers = []
|
first_row = True
|
row: docx.table._Row
|
for row in table.rows:
|
if first_row:
|
for cell in row.cells:
|
headers.append(cell.text)
|
first_row = False
|
continue
|
row_data = {}
|
row_idx = 0
|
for cell in row.cells:
|
if cell.tables:
|
# 嵌套表格处理
|
if len(cell.tables) == 1:
|
text = self.table_to_json(cell.tables[0])
|
else:
|
text = []
|
for tbl in cell.tables:
|
tbl_json = self.table_to_json(tbl)
|
text.append(tbl_json)
|
else:
|
# 单元格文本获取
|
text = cell.text
|
row_data[headers[row_idx]] = text
|
row_idx += 1
|
|
table_data.append(row_data)
|
return table_data
|
|
def split(self):
|
"""
|
将文档拆分成段落，并返回段落列表
|
|
:return: list[ParagraphInfo] - 段落列表
|
"""
|
document = docx.Document(self.doc_file)
|
table_cnt = 0
|
paragraph_cnt = 0
|
|
for element in document.element.body:
|
if element.tag.endswith('p'): # 段落
|
# 获取标题多级编号
|
paragraph = document.paragraphs[paragraph_cnt]
|
paragraph_text = paragraph.text
|
if paragraph_text:
|
self.paragraphs.append(ParagraphInfo(paragraph_text, self.get_header_level(paragraph)))
|
# 检查是否是图片，如果是图片则转换为文本
|
img_data = self.get_image_blob(paragraph)
|
if img_data:
|
text = self.gen_text_from_img(img_data)
|
self.paragraphs.append(ParagraphInfo(text, 0))
|
paragraph_cnt += 1
|
elif element.tag.endswith('tbl'): # 表格
|
table = document.tables[table_cnt] # 获取当前表格对象
|
table_cnt += 1
|
table_data = self.table_to_json(table)
|
self.paragraphs.append(ParagraphInfo(json.dumps(table_data, indent=4, ensure_ascii=False), 0))
|
else:
|
continue
|
# 生成标题编号
|
self.gen_title_no(self.paragraphs)
|
|
@staticmethod
|
def get_image_blob(paragraph):
|
# 遍历段落中的所有Run对象（图片通常在单独的Run中）
|
for run in paragraph.runs:
|
xml = run._element.xml
|
if xml.find('v:imagedata') != -1:
|
# 使用正则表达式查找r:id属性
|
match = re.search(r'r:id="([^"]+)"', xml)
|
if match:
|
r_id = match.group(1)
|
if r_id:
|
# 获取图片信息
|
image_part = paragraph.part.rels[r_id].target_part
|
return DocSplit.image_convert(image_part.blob, "png")
|
if xml.find('wp:inline') != -1 or xml.find('wp:anchor') != -1:
|
# 使用正则表达式查找r:embed属性
|
match = re.search(r'r:embed="([^"]+)"', xml)
|
if match:
|
r_id = match.group(1)
|
if r_id:
|
# 获取图片信息
|
image_part = paragraph.part.rels[r_id].target_part
|
return DocSplit.image_convert(image_part.blob, "png")
|
return None
|
|
@staticmethod
|
def gen_title_no(paragraphs: list[ParagraphInfo]):
|
title_levels = [1, 1, 1, 1, 1, 1, 1, 1, 1]
|
for i in range(len(paragraphs)):
|
if paragraphs[i].level > 0:
|
for j in range(paragraphs[i].level - 1):
|
title_levels[j] = 1
|
paragraphs[i].title_no = '.'.join([str(x) for x in title_levels[0:paragraphs[i].level]])
|
title_levels[paragraphs[i].level - 1] += 1
|
else:
|
title_levels = [1, 1, 1, 1, 1, 1, 1, 1, 1]
|
|
@staticmethod
|
def get_header_level(paragraph) -> int:
|
if paragraph.style.base_style:
|
style = paragraph.style.base_style
|
else:
|
style = paragraph.style
|
if style and style.name.startswith('Heading'):
|
# 获取标题级别
|
level = int(style.name.split(' ')[1])
|
return level
|
else:
|
return 0
|
|
@staticmethod
|
def image_convert(_in: bytes, _out_format: str) -> bytes:
|
in_io = io.BytesIO()
|
in_io.write(_in)
|
img = Image.open(in_io, "r")
|
out_io = io.BytesIO()
|
img.save(out_io, "png")
|
out_io.seek(0)
|
return out_io.read()
|
|
def gen_text_from_img(self, img_data:bytes):
|
return self.image_to_text.gen_text_from_img(img_data)
|
|
if __name__ == '__main__':
|
doc_file = r'D:\workspace\PythonProjects\KnowledgeBase\doc\ZL格式(公开).docx'
|
doc_split = DocSplit(doc_file)
|
doc_split.split()
|
print("\n".join([x.full_text for x in doc_split.paragraphs]))
|
