# -*- coding: utf-8 -*-
|
#
|
# @author: lyg, ym
|
# @date: 2025-5-8
|
# @version: 1
|
# @description: docx文档拆分器，根据段落拆分，将图片和表格转换为json数据。
|
import docx
|
import docx.table
|
import json
|
from PIL import Image
|
import io
|
import re
|
import typing
|
|
from knowledgebase.doc.image_to_text import ImageToText
|
from knowledgebase.doc.models import ParagraphInfo
|
from knowledgebase.log import Log
|
from bs4 import BeautifulSoup
|
|
class DocSplit:
|
"""
|
docx文档拆分器，根据段落拆分，将图片和表格转换为json数据。
|
1.封装段落信息
|
2.将图片转换为自然语言描述
|
3.将表格转换为json格式
|
4.将段落按照文档标题级别组合成树形结构
|
|
"""
|
|
def __init__(self, docx_file: str, docx_type: str):
|
"""
|
docx文档拆分
|
:param docx_file: 要拆分的docx文件路径
|
:param docx_type: 文档类型
|
"""
|
self.docx_file = docx_file
|
self.docx_type = docx_type
|
self.image_to_text = ImageToText()
|
self.paragraphs: list[ParagraphInfo] = []
|
self.paragraph_tree: list[ParagraphInfo] = []
|
|
def table_to_json(self, table: docx.table.Table):
|
"""
|
将表格转换为 JSON 格式
|
|
:param table: docx.table.Table - 要转换的表格对象
|
:return list - 表格数据，以 JSON 格式表示
|
"""
|
headers = []
|
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
|
if cell._element.xml.find("w:ins")!=-1:
|
soup = BeautifulSoup(cell._element.xml, "xml")
|
text = ''.join([x.get_text() for x in soup.find_all("w:t")])
|
# row_data[headers[row_idx]] = text
|
row_data.append(text)
|
row_idx += 1
|
|
table_data.append(row_data)
|
return table_data
|
|
def split(self):
|
"""
|
将文档拆分成段落，并返回段落列表
|
|
:return: list[ParagraphInfo] - 段落列表
|
"""
|
Log.info(f"开始拆分文档：{self.docx_file}")
|
document = docx.Document(self.docx_file)
|
table_cnt = 0
|
paragraph_cnt = 0
|
|
for element in document.element.body:
|
if element.tag.endswith('p'): # 段落
|
# 获取标题多级编号
|
paragraph = document.paragraphs[paragraph_cnt]
|
p_text = paragraph.text
|
is_toc = paragraph.style.name.startswith('TOC') and '目' in p_text and '录' in p_text
|
try:
|
num = element.pPr.numPr.numId.val
|
level = element.pPr.numPr.ilvl.val
|
except:
|
num = 0
|
level = 0
|
if p_text:
|
title_level = self.get_title_level(paragraph)
|
self.paragraphs.append(ParagraphInfo(p_text, title_level, num, level, is_toc))
|
# 检查是否是图片，如果是图片则转换为文本
|
img_data = self.get_image_text(paragraph)
|
if img_data:
|
text = self.gen_text_from_img(img_data)
|
text = f"```图片（以下内容为图片描述）\n{text}\n```"
|
self.paragraphs.append(ParagraphInfo(text, 0, num, level))
|
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\n" + json.dumps(table_data, indent=4, ensure_ascii=False) + "\n```", 0))
|
else:
|
continue
|
# 去除目录
|
self.remove_toc(self.paragraphs)
|
# 生成标题编号
|
Log.info(f"开始生成标题编号和列表编号")
|
self.gen_title_num(self.paragraphs)
|
# 生成树形结构
|
Log.info(f"开始生成树形结构")
|
self.gen_paragraph_tree(self.paragraphs)
|
|
@staticmethod
|
def remove_toc(paragraphs: [ParagraphInfo]):
|
rm_list = []
|
for p in paragraphs:
|
if p.is_toc:
|
rm_list.append(p)
|
elif rm_list and p.title_level == 1:
|
break
|
elif rm_list:
|
rm_list.append(p)
|
for p in rm_list:
|
paragraphs.remove(p)
|
|
@staticmethod
|
def get_image_text(paragraph):
|
"""
|
获取段落中的图片描述
|
:param paragraph: 段落
|
:return: 图片内容描述信息
|
"""
|
# 遍历段落中的所有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)
|
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)
|
return None
|
|
@staticmethod
|
def gen_title_num(paragraphs: list[ParagraphInfo]):
|
"""
|
生成标题编号和列表编号
|
标题级别从1-9，0表示正文
|
|
:param paragraphs: list[ParagraphInfo] - 段落列表
|
:return: None
|
"""
|
MAX_TITLE_LEVEL = 9 # 定义为常量，便于统一管理和扩展
|
title_levels = [0] * MAX_TITLE_LEVEL # 初始化为全0
|
|
list_counters = [0] * MAX_TITLE_LEVEL
|
|
def format_number(level: int, value: int) -> str:
|
# 使用映射方式简化逻辑
|
if level < 0 or level > 4:
|
return str(value)
|
formats = {
|
0: lambda v: f"({v})",
|
1: lambda v: f"{v})",
|
2: lambda v: f"({chr(96 + v)})",
|
3: lambda v: f"{chr(96 + v)})",
|
4: lambda v: chr(96 + v),
|
}
|
return formats[level](value)
|
|
for p in paragraphs:
|
if p.title_level > 0:
|
title_levels[p.title_level - 1] += 1
|
for i in range(p.title_level, MAX_TITLE_LEVEL):
|
title_levels[i] = 0
|
p.title_num = '.'.join([str(x) for x in title_levels[:p.title_level]])
|
list_counters = [0] * MAX_TITLE_LEVEL
|
else:
|
# 处理列表编号
|
if p.num > 0:
|
level = p.num_level
|
|
# 校验 level 合法性
|
if level < 0 or level >= MAX_TITLE_LEVEL:
|
continue
|
list_counters[level] += 1
|
|
# 重置当前层级之后的计数器
|
for l in range(level + 1, MAX_TITLE_LEVEL):
|
list_counters[l] = 0
|
|
# 当前层级递增并赋值
|
p.title_num = format_number(level, list_counters[level])
|
else:
|
list_counters = [0] * MAX_TITLE_LEVEL
|
|
@staticmethod
|
def get_title_level(paragraph) -> int:
|
"""
|
获取段落标题级别
|
|
:param paragraph: docx.paragraph.Paragraph - 要获取标题级别的段落对象
|
:return: int - 标题级别，0 表示非标题
|
"""
|
style = paragraph.style
|
if style and style.name.startswith('Heading'):
|
# 获取标题级别
|
level = int(style.name.split(' ')[1])
|
return level
|
elif style.base_style and style.base_style.name.startswith('Heading'):
|
level = int(style.base_style.name.split(' ')[1])
|
return level
|
else:
|
return 0
|
|
@staticmethod
|
def image_convert(_in: bytes) -> bytes:
|
"""
|
将图片转换为png格式的bytes
|
:param _in: bytes - 图片数据
|
:return: bytes - png格式的图片数据
|
"""
|
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):
|
"""
|
利用LLM将图片转为文本
|
:param img_data: bytes - 图片数据
|
:return: str - 文本
|
"""
|
return ''
|
return self.image_to_text.gen_text_from_img(img_data)
|
|
def gen_paragraph_tree(self, paragraphs: typing.List[ParagraphInfo]):
|
"""
|
生成段落树结构，根据title_level划分段落树
|
|
:param paragraphs: list[ParagraphInfo] - 段落列表（会被原地修改）
|
"""
|
if not paragraphs:
|
return
|
|
stack = []
|
result = []
|
_paragraphs = []
|
|
def merge_paragraph_text(info: ParagraphInfo):
|
text_nodes = [child for child in info.children if child.title_level == 0]
|
info.text += '\n' + '\n'.join([child.full_text for child in text_nodes])
|
info.children = [child for child in info.children if child.title_level > 0]
|
|
for p in paragraphs:
|
if p.title_level == 1:
|
result.append(p)
|
# 清理栈顶比当前级别低或相等的节点
|
while stack and p.title_level != 0 and stack[-1].title_level >= p.title_level:
|
_p = stack.pop()
|
merge_paragraph_text(_p)
|
|
if p.title_level > 0:
|
if len(stack):
|
stack[-1].children.append(p)
|
stack.append(p)
|
_paragraphs.append(p)
|
elif len(stack):
|
stack[-1].children.append(p)
|
else:
|
# 非标题段落直接加入结果
|
result.append(p)
|
|
while stack:
|
merge_paragraph_text(stack.pop())
|
|
# 替换原始列表内容，避免多次 remove 操作
|
self.paragraphs[:] = _paragraphs
|
self.paragraph_tree = result
|
