1614 lines
58 KiB
Python
1614 lines
58 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
from collections import Counter
|
||
import uuid
|
||
from openai import OpenAI
|
||
import re
|
||
from typing import List, Dict, Optional, Union, Any, Tuple
|
||
# 指定 JSON 文件路径
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/122-06A0014-B01001_发动机-维修手册_content_list.json'
|
||
try:
|
||
from filename_proceess_and_kgquery import get_entity
|
||
except ImportError:
|
||
get_entity = None
|
||
from pathlib import PurePath, PurePosixPath, Path as PathLib
|
||
from fastapi import FastAPI, File, Path, UploadFile, HTTPException, Form, Request, Header, Body
|
||
from fileparse_util import process_document
|
||
from io import BytesIO
|
||
from urllib.parse import unquote, urlparse
|
||
import asyncio
|
||
import base64
|
||
import binascii
|
||
import threading
|
||
import aiofiles
|
||
import httpx
|
||
import logging
|
||
from resetlevel import reset_textlevel
|
||
from rapidocr_onnxruntime import RapidOCR
|
||
from PIL import Image, ImageOps, UnidentifiedImageError
|
||
from config import (
|
||
API_URLS,
|
||
IMAGE_DIR,
|
||
MAX_SPLIT_IMAGE_BASE64_CHARS,
|
||
MAX_SPLIT_IMAGE_BYTES,
|
||
MAX_SPLIT_IMAGE_INPUT_BYTES,
|
||
OCR_CONCURRENCY,
|
||
SPLIT_IMAGE_COMPRESS_THRESHOLD_BYTES,
|
||
SPLIT_IMAGE_JPEG_QUALITY,
|
||
SPLIT_IMAGE_MAX_DIMENSION,
|
||
)
|
||
try:
|
||
from config import SHIP_MODEL_NAME
|
||
except ImportError:
|
||
SHIP_MODEL_NAME = ""
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
ocr_engine = RapidOCR()
|
||
OCR_SEMAPHORE = asyncio.Semaphore(max(1, OCR_CONCURRENCY))
|
||
_ocr_thread_lock = threading.Lock()
|
||
_image_ocr_cache: Dict[str, str] = {}
|
||
_image_ocr_cache_lock = threading.Lock()
|
||
_api_url_inflight = [0] * len(API_URLS)
|
||
_api_url_lock = threading.Lock()
|
||
IMAGE_OCR_PREFIX = "图片中文本内容为："
|
||
IMAGE_OCR_PREFIXES = (IMAGE_OCR_PREFIX, "图片包含的文字内容为：")
|
||
|
||
|
||
def _image_mime_type(suffix: str) -> str:
|
||
return {
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".png": "image/png",
|
||
".bmp": "image/bmp",
|
||
".webp": "image/webp",
|
||
".gif": "image/gif",
|
||
".tif": "image/tiff",
|
||
".tiff": "image/tiff",
|
||
}.get((suffix or "").lower(), "image/png")
|
||
|
||
|
||
def _image_save_format(suffix: str) -> tuple[str, dict]:
|
||
suffix = (suffix or "").lower()
|
||
if suffix in {".jpg", ".jpeg"}:
|
||
return "JPEG", {"quality": SPLIT_IMAGE_JPEG_QUALITY, "optimize": True}
|
||
if suffix == ".webp":
|
||
return "WEBP", {"quality": SPLIT_IMAGE_JPEG_QUALITY, "method": 6}
|
||
if suffix == ".png":
|
||
return "PNG", {"optimize": True, "compress_level": 9}
|
||
if suffix in {".tif", ".tiff"}:
|
||
return "TIFF", {"compression": "tiff_lzw"}
|
||
if suffix == ".gif":
|
||
return "GIF", {"optimize": True}
|
||
if suffix == ".bmp":
|
||
return "BMP", {}
|
||
return "PNG", {"optimize": True, "compress_level": 9}
|
||
|
||
|
||
def _compression_suffix(suffix: str) -> str:
|
||
suffix = (suffix or "").lower()
|
||
if suffix in {".jpg", ".jpeg", ".png", ".webp"}:
|
||
return suffix
|
||
if suffix in {".bmp", ".gif", ".tif", ".tiff"}:
|
||
return ".png"
|
||
return ".png"
|
||
|
||
|
||
def _save_image_to_bytes(img: Image.Image, suffix: str) -> bytes:
|
||
image_format, save_options = _image_save_format(suffix)
|
||
if image_format in {"JPEG", "WEBP", "BMP"} and img.mode not in {"RGB", "L"}:
|
||
img = img.convert("RGB")
|
||
elif image_format in {"PNG", "TIFF", "GIF"} and img.mode not in {"RGB", "RGBA", "L", "P"}:
|
||
img = img.convert("RGBA")
|
||
|
||
output = BytesIO()
|
||
img.save(output, format=image_format, **save_options)
|
||
return output.getvalue()
|
||
|
||
|
||
def _save_image_as_jpeg(img: Image.Image) -> bytes:
|
||
if img.mode not in {"RGB", "L"}:
|
||
img = img.convert("RGB")
|
||
last_bytes = b""
|
||
for quality in (SPLIT_IMAGE_JPEG_QUALITY, 82, 74):
|
||
output = BytesIO()
|
||
img.save(output, format="JPEG", quality=quality, optimize=True)
|
||
last_bytes = output.getvalue()
|
||
if len(last_bytes) <= MAX_SPLIT_IMAGE_BYTES:
|
||
break
|
||
return last_bytes
|
||
|
||
|
||
def compress_image_bytes(image_bytes: bytes, suffix: str = "") -> tuple[bytes, str]:
|
||
if not image_bytes:
|
||
raise HTTPException(status_code=400, detail="decoded image is empty")
|
||
if len(image_bytes) > MAX_SPLIT_IMAGE_INPUT_BYTES:
|
||
raise HTTPException(status_code=413, detail="image input is too large")
|
||
|
||
suffix = (suffix or "").lower()
|
||
try:
|
||
with Image.open(BytesIO(image_bytes)) as img:
|
||
img = ImageOps.exif_transpose(img)
|
||
max_side = max(img.size)
|
||
should_resize = SPLIT_IMAGE_MAX_DIMENSION > 0 and max_side > SPLIT_IMAGE_MAX_DIMENSION
|
||
should_reencode = len(image_bytes) > SPLIT_IMAGE_COMPRESS_THRESHOLD_BYTES
|
||
|
||
if not should_resize and not should_reencode:
|
||
if len(image_bytes) > MAX_SPLIT_IMAGE_BYTES:
|
||
raise HTTPException(status_code=413, detail="image is too large")
|
||
return image_bytes, suffix
|
||
|
||
if should_resize:
|
||
scale = SPLIT_IMAGE_MAX_DIMENSION / max_side
|
||
new_size = (
|
||
max(1, int(img.width * scale)),
|
||
max(1, int(img.height * scale)),
|
||
)
|
||
resample = getattr(getattr(Image, "Resampling", Image), "LANCZOS")
|
||
img = img.resize(new_size, resample)
|
||
|
||
compressed_suffix = _compression_suffix(suffix)
|
||
compressed_bytes = _save_image_to_bytes(img, compressed_suffix)
|
||
|
||
if len(compressed_bytes) > MAX_SPLIT_IMAGE_BYTES:
|
||
compressed_bytes = _save_image_as_jpeg(img)
|
||
compressed_suffix = ".jpg"
|
||
except UnidentifiedImageError as exc:
|
||
raise HTTPException(status_code=400, detail="image content is not a supported image") from exc
|
||
except HTTPException:
|
||
raise
|
||
except Exception as exc:
|
||
logger.warning(f"Failed to compress image, using original bytes: {exc}")
|
||
compressed_bytes = image_bytes
|
||
compressed_suffix = suffix
|
||
|
||
if len(compressed_bytes) >= len(image_bytes) and len(image_bytes) <= MAX_SPLIT_IMAGE_BYTES:
|
||
return image_bytes, suffix
|
||
if len(compressed_bytes) > MAX_SPLIT_IMAGE_BYTES:
|
||
raise HTTPException(status_code=413, detail="image is too large after compression")
|
||
return compressed_bytes, compressed_suffix
|
||
|
||
|
||
def prepare_split_image(image_base64: str, suffix: str = "") -> tuple[bytes, str]:
|
||
raw = (image_base64 or "").strip()
|
||
if not raw:
|
||
raise HTTPException(status_code=400, detail="image_base64 is empty")
|
||
|
||
if "," in raw and raw.split(",", 1)[0].lower().startswith("data:image"):
|
||
raw = raw.split(",", 1)[1]
|
||
|
||
raw = "".join(raw.split())
|
||
if len(raw) > MAX_SPLIT_IMAGE_BASE64_CHARS:
|
||
raise HTTPException(status_code=413, detail="image_base64 is too large")
|
||
|
||
try:
|
||
image_bytes = base64.b64decode(raw, validate=True)
|
||
except (binascii.Error, ValueError) as exc:
|
||
raise HTTPException(status_code=400, detail="image_base64 is not valid base64") from exc
|
||
|
||
return compress_image_bytes(image_bytes, suffix)
|
||
|
||
|
||
def decode_base64_image(image_base64: str, suffix: str = "") -> bytes:
|
||
image_bytes, _ = prepare_split_image(image_base64, suffix)
|
||
return image_bytes
|
||
|
||
|
||
def safe_original_filename(filename: str) -> str:
|
||
name = os.path.basename((filename or "").strip())
|
||
if not name:
|
||
raise HTTPException(status_code=400, detail="filename is empty")
|
||
safe_name = re.sub(r"[^\w.\-\u4e00-\u9fff]+", "_", name).strip("._")
|
||
if not safe_name:
|
||
raise HTTPException(status_code=400, detail="filename is invalid")
|
||
return safe_name
|
||
|
||
|
||
def _strip_markdown_image_path(img_path: str) -> str:
|
||
value = (img_path or "").strip()
|
||
match = re.match(r"!\[[^\]]*\]\((.*?)\)", value)
|
||
return match.group(1).strip() if match else value
|
||
|
||
|
||
def resolve_local_image_path(img_path: str) -> Optional[PathLib]:
|
||
value = _strip_markdown_image_path(str(img_path or ""))
|
||
if not value:
|
||
return None
|
||
|
||
parsed_path = unquote(urlparse(value).path or value)
|
||
marker = "/images/"
|
||
if marker in parsed_path:
|
||
relative_path = parsed_path.split(marker, 1)[1]
|
||
else:
|
||
normalized = parsed_path.replace("\\", "/").lstrip("/")
|
||
if normalized.startswith("images/"):
|
||
relative_path = normalized.split("/", 1)[1]
|
||
else:
|
||
relative_path = normalized
|
||
|
||
relative_path = relative_path.replace("\\", "/").lstrip("/")
|
||
if not relative_path:
|
||
return None
|
||
|
||
image_root = PathLib(IMAGE_DIR).resolve()
|
||
local_path = (image_root / PurePosixPath(relative_path)).resolve()
|
||
try:
|
||
local_path.relative_to(image_root)
|
||
except ValueError:
|
||
logger.warning(f"Image path is outside IMAGE_DIR and will be ignored: {img_path}")
|
||
return None
|
||
|
||
return local_path if local_path.is_file() else None
|
||
|
||
|
||
def extract_image_text_sync(image_path: Union[str, PathLib]) -> dict:
|
||
resolved_path = str(PathLib(image_path).resolve())
|
||
with _image_ocr_cache_lock:
|
||
cached_text = _image_ocr_cache.get(resolved_path)
|
||
if cached_text is not None:
|
||
return {"texts": [], "full_text": cached_text, "elapse": None}
|
||
|
||
with _ocr_thread_lock:
|
||
ocr_result, elapse = ocr_engine(resolved_path)
|
||
|
||
texts = [item[1] for item in (ocr_result or [])]
|
||
full_text = IMAGE_OCR_PREFIX + ",".join(texts) if texts else ""
|
||
with _image_ocr_cache_lock:
|
||
_image_ocr_cache[resolved_path] = full_text
|
||
return {
|
||
"texts": texts,
|
||
"full_text": full_text,
|
||
"elapse": elapse,
|
||
}
|
||
|
||
|
||
def build_image_markdown_with_ocr(markdown_image: str, local_image_path: Optional[PathLib]) -> str:
|
||
if any(prefix in markdown_image for prefix in IMAGE_OCR_PREFIXES):
|
||
return markdown_image
|
||
if not local_image_path:
|
||
return markdown_image
|
||
|
||
try:
|
||
ocr_data = extract_image_text_sync(local_image_path)
|
||
except Exception as exc:
|
||
logger.warning(f"Failed to OCR local image {local_image_path}: {exc}")
|
||
return markdown_image
|
||
|
||
ocr_full_text = (ocr_data.get("full_text") or "").strip()
|
||
if not ocr_full_text:
|
||
return markdown_image
|
||
return f"{markdown_image}\n{ocr_full_text}"
|
||
|
||
|
||
async def get_image_ocr_text(local_image_path: Optional[PathLib]) -> str:
|
||
if not local_image_path:
|
||
return ""
|
||
|
||
try:
|
||
ocr_data = await extract_image_text(local_image_path)
|
||
except Exception as exc:
|
||
logger.warning(f"Failed to OCR local image {local_image_path}: {exc}")
|
||
return ""
|
||
|
||
return (ocr_data.get("full_text") or "").strip()
|
||
|
||
|
||
def stringify_content_field(value: Any) -> str:
|
||
if value is None:
|
||
return ""
|
||
if isinstance(value, list):
|
||
return "\n".join(str(item) for item in value if str(item).strip())
|
||
if isinstance(value, dict):
|
||
return json.dumps(value, ensure_ascii=False)
|
||
return str(value)
|
||
|
||
|
||
def clean_table_body(table_body: Any) -> str:
|
||
table_body = stringify_content_field(table_body)
|
||
for pattern in [' colspan="1"', ' rowspan="1"', " colspan='1'", " rowspan='1'", " colspan=1", " rowspan=1"]:
|
||
table_body = table_body.replace(pattern, "")
|
||
return table_body
|
||
|
||
|
||
def join_nonempty_parts(*parts: Any) -> str:
|
||
return "\n".join(part for part in (stringify_content_field(part).strip() for part in parts) if part)
|
||
|
||
|
||
def record_to_chunk_text(ins: Dict[str, Any]) -> str:
|
||
record_type = ins.get("type")
|
||
ocr_text = ins.get("ocr_text", "")
|
||
|
||
if record_type == "text":
|
||
text = stringify_content_field(ins.get("text", ""))
|
||
level = ins.get("text_level")
|
||
if level is not None and isinstance(level, int) and level > 0:
|
||
text = "#" * min(level, 6) + " " + text
|
||
return join_nonempty_parts(ocr_text, text)
|
||
|
||
if record_type == "table":
|
||
return join_nonempty_parts(
|
||
ocr_text,
|
||
ins.get("table_caption"),
|
||
clean_table_body(ins.get("table_body", "")),
|
||
ins.get("table_footnote"),
|
||
)
|
||
|
||
if record_type == "image":
|
||
return join_nonempty_parts(
|
||
ocr_text,
|
||
ins.get("image_caption"),
|
||
ins.get("img_path"),
|
||
ins.get("image_footnote"),
|
||
)
|
||
|
||
if record_type == "chart":
|
||
return join_nonempty_parts(
|
||
ocr_text,
|
||
ins.get("chart_caption"),
|
||
ins.get("content"),
|
||
ins.get("img_path"),
|
||
ins.get("chart_footnote"),
|
||
)
|
||
|
||
if record_type == "equation":
|
||
return join_nonempty_parts(
|
||
ocr_text,
|
||
ins.get("text"),
|
||
ins.get("img_path"),
|
||
)
|
||
|
||
if record_type == "code":
|
||
return join_nonempty_parts(
|
||
ocr_text,
|
||
ins.get("code_caption"),
|
||
ins.get("code_body"),
|
||
ins.get("code_footnote"),
|
||
)
|
||
|
||
if record_type == "list":
|
||
return join_nonempty_parts(ocr_text, ins.get("list_items"))
|
||
|
||
if record_type in {"discarded", "header", "footer", "page_number"}:
|
||
return join_nonempty_parts(ocr_text, ins.get("text"))
|
||
|
||
return join_nonempty_parts(ocr_text, ins.get("text"))
|
||
|
||
|
||
def is_supported_chunk_record(ins: Dict[str, Any]) -> bool:
|
||
return ins.get("type") in {
|
||
"text",
|
||
"table",
|
||
"image",
|
||
"chart",
|
||
"equation",
|
||
"code",
|
||
"list",
|
||
"discarded",
|
||
"header",
|
||
"footer",
|
||
"page_number",
|
||
}
|
||
|
||
|
||
def append_caption_ocr(ins: Dict[str, Any], ocr_texts: List[str]) -> None:
|
||
caption_field = {
|
||
"image": "image_caption",
|
||
"table": "table_caption",
|
||
"chart": "chart_caption",
|
||
"code": "code_caption",
|
||
}.get(ins.get("type"))
|
||
if not caption_field:
|
||
caption_field = "ocr_text"
|
||
|
||
caption = ins.get(caption_field, [])
|
||
existing_text = stringify_content_field(caption)
|
||
if caption_field == "ocr_text":
|
||
existing_text = join_nonempty_parts(existing_text, ins.get("text"))
|
||
|
||
additions = []
|
||
for text in ocr_texts:
|
||
text = (text or "").strip()
|
||
text_without_prefix = text
|
||
for prefix in IMAGE_OCR_PREFIXES:
|
||
if text_without_prefix.startswith(prefix):
|
||
text_without_prefix = text_without_prefix[len(prefix):].strip()
|
||
break
|
||
|
||
if (
|
||
text
|
||
and text not in existing_text
|
||
and text_without_prefix not in existing_text
|
||
and text not in additions
|
||
):
|
||
additions.append(text)
|
||
|
||
if not additions:
|
||
return
|
||
|
||
if caption_field == "ocr_text":
|
||
ins[caption_field] = join_nonempty_parts(additions, caption)
|
||
return
|
||
|
||
if isinstance(caption, list):
|
||
ins[caption_field] = additions + caption
|
||
elif caption is None or not str(caption).strip():
|
||
ins[caption_field] = additions
|
||
else:
|
||
ins[caption_field] = "\n".join(additions) + "\n" + str(caption).lstrip()
|
||
|
||
|
||
def build_image_markdown(markdown_or_path: str, prefix: str) -> Tuple[str, Optional[PathLib]]:
|
||
img_path = str(markdown_or_path).strip()
|
||
if not img_path:
|
||
return "", None
|
||
|
||
if img_path.strip().startswith("!["):
|
||
return img_path, resolve_local_image_path(img_path)
|
||
|
||
p = PurePosixPath(img_path)
|
||
if p.parts and p.parts[0] == "images":
|
||
relative_path = str(PurePosixPath(*p.parts[1:]))
|
||
return f"", resolve_local_image_path(relative_path)
|
||
|
||
return f"", resolve_local_image_path(img_path)
|
||
|
||
|
||
def image_base64_to_data_url(image_base64: Union[str, bytes], suffix: str) -> str:
|
||
if isinstance(image_base64, bytes):
|
||
raw = base64.b64encode(image_base64).decode("ascii")
|
||
else:
|
||
raw = (image_base64 or "").strip()
|
||
if "," in raw and raw.split(",", 1)[0].lower().startswith("data:image"):
|
||
raw = raw.split(",", 1)[1]
|
||
raw = "".join(raw.split())
|
||
|
||
mime_type = _image_mime_type(suffix)
|
||
return f"data:{mime_type};base64,{raw}"
|
||
|
||
|
||
async def extract_image_text(image_path: PathLib) -> dict:
|
||
async with OCR_SEMAPHORE:
|
||
return await asyncio.to_thread(extract_image_text_sync, image_path)
|
||
|
||
def format_entity_text(data):
|
||
# 1. 定义实体类型与属性的映射关系
|
||
# 格式: '实体类型': [ ('前缀1', '属性1'), ('前缀2', '属性2') ]
|
||
type_mapping = {
|
||
'舰艇': [('舰艇为', '名称'), ('，舷号', '舷号')],
|
||
'系统': [('，系统：', '名称')],
|
||
'子系统': [('，子系统：', '名称')],
|
||
'设备': [('，设备：', '名称')]
|
||
}
|
||
|
||
text_parts = []
|
||
|
||
# 2. 遍历 entities 列表
|
||
for entity in data.get('entities', []):
|
||
entity_type = entity.get('type')
|
||
|
||
properties = entity.get('properties', {})
|
||
|
||
# 3. 如果当前实体类型在映射表中，则进行拼接
|
||
if entity_type in type_mapping:
|
||
entity_str = ""
|
||
# 遍历该实体需要提取的所有属性
|
||
for prefix, attr_key in type_mapping[entity_type]:
|
||
value = str(properties.get(attr_key, ''))
|
||
entity_str += f"{prefix}{value}"
|
||
text_parts.append(entity_str)
|
||
|
||
# 4. 将所有部分拼接成一段完整的文本
|
||
# 注意：因为前缀中已经包含了逗号，这里直接拼接即可
|
||
return ''.join(text_parts)
|
||
|
||
def get_api_url(task_id: str = "") -> tuple:
|
||
"""返回 (analyze_url, contentlist_url)，选当前 inflight 最少的实例。"""
|
||
with _api_url_lock:
|
||
idx = _api_url_inflight.index(min(_api_url_inflight))
|
||
_api_url_inflight[idx] += 1
|
||
base = API_URLS[idx]
|
||
return base
|
||
|
||
|
||
def release_api_url(url: str):
|
||
with _api_url_lock:
|
||
if url in API_URLS:
|
||
idx = API_URLS.index(url)
|
||
_api_url_inflight[idx] = max(0, _api_url_inflight[idx] - 1)
|
||
def extract_text_from_first_tr(html: str) -> str:
|
||
"""提取 HTML 表格中第一个 <tr> 内的文本（去除标签）"""
|
||
if not isinstance(html, str):
|
||
return ""
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', html, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
return ""
|
||
first_tr_content = match.group(1)
|
||
# 去除所有 HTML 标签
|
||
text = re.sub(r'<[^>]+>', '', first_tr_content)
|
||
return text.strip()
|
||
|
||
|
||
def matches_any_pattern(text: str, patterns: List[List[str]]) -> bool:
|
||
"""检查文本是否包含任意一组关键词（全部命中）"""
|
||
if not text:
|
||
return False
|
||
text_lower = text.lower()
|
||
for pattern in patterns:
|
||
if all(keyword in text_lower for keyword in pattern):
|
||
return True
|
||
return False
|
||
|
||
|
||
def is_maintenance_table(table_body: str, patterns: List[List[str]]) -> bool:
|
||
"""判断是否为维修项目表格：仅检查第一个 <tr> 是否匹配任一模式"""
|
||
first_tr_text = extract_text_from_first_tr(table_body)
|
||
return matches_any_pattern(first_tr_text, patterns)
|
||
|
||
|
||
# ==============================
|
||
# 维修项目表格的关键词模式（可按需扩展）
|
||
# ==============================
|
||
|
||
MAINTENANCE_TABLE_PATTERNS: List[List[str]] = [
|
||
["维修项目编号", "组成编码", "名称", "维修项目", "维修间隔期"],
|
||
["维修项目编号", "维修级别", "名称"],
|
||
["项目编号", "编码", "维修内容", "周期"],
|
||
["序号", "维修项", "标准", "频次"], # 示例：可继续添加业务变体
|
||
]
|
||
|
||
|
||
# ==============================
|
||
# 主函数：提取维修项目表格及其后续内容组
|
||
# ==============================
|
||
|
||
def extract_maintenance_groups_and_remaining(
|
||
records: List[Dict[str, Any]],
|
||
patterns: List[List[str]] = MAINTENANCE_TABLE_PATTERNS,
|
||
max_following_for_last: int | None = None
|
||
) -> Tuple[List[List[Dict]], List[Dict]]:
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
|
||
# 标记是否已经遇到第一个维修表格（用于跳过）
|
||
skipped_first = False
|
||
|
||
for record in records:
|
||
is_start = (
|
||
record.get("type") == "table" and
|
||
is_maintenance_table(record.get("table_body", ""), patterns)
|
||
)
|
||
|
||
if is_start:
|
||
# 如果是第一个维修表，跳过（不加入 groups，也不标记 used_ids）
|
||
if not skipped_first:
|
||
skipped_first = True
|
||
# 注意：这里不清空 current_group，因为前面不可能有 group（第一个 start）
|
||
current_group = None # 确保不会把之前的非 start 内容误加
|
||
continue # 跳过这个 record，不加入任何 group
|
||
|
||
# 从第二个维修表开始，才正常分组
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
current_group.append(record)
|
||
|
||
# 处理最后一个 group（仅当有有效 group 时）
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# 限制最后一个 group 长度
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
max_len = max_following_for_last + 1
|
||
if len(last_group) > max_len:
|
||
truncated = last_group[:max_len]
|
||
removed_ids = {r["id"] for r in last_group[max_len:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
# remaining_records 包含所有未被 used_ids 标记的记录
|
||
# 由于第一个 group 没有加入 used_ids，所以会保留在 remaining_records 中
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
def is_operation_table_by_first_tr(table_body: str) -> bool:
|
||
"""
|
||
判断 table_body 是否为操作项目表，依据是第一行 <tr> 中是否包含 '项目编号' 和 '操作项目'
|
||
"""
|
||
if not isinstance(table_body, str):
|
||
return False
|
||
|
||
# 非贪婪匹配第一个 <tr> ... </tr>
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
return False
|
||
|
||
first_tr_html = match.group(1)
|
||
|
||
# 去除所有 HTML 标签，提取纯文本
|
||
text_in_first_tr = re.sub(r'<[^>]+>', '', first_tr_html)
|
||
|
||
# 转为小写（中文无所谓大小写，但统一处理更安全）
|
||
text_clean = text_in_first_tr.strip()
|
||
|
||
# 检查是否同时包含两个关键词
|
||
has_code = "项目编号" in text_clean
|
||
has_item = "操作项目" in text_clean
|
||
return has_code and has_item
|
||
|
||
def extract_operation_groups_and_remaining(records, max_following_for_last=23):
|
||
"""
|
||
通过语义判断表格是否为操作项目表（基于第一行内容），并分组。
|
||
注意：第一个识别到的操作项目表及其后续记录组成的 group 会被跳过（不加入 groups），
|
||
但这些记录仍保留在 remaining_records 中。
|
||
|
||
普通起始标记：组 = [start, ..., next_start - 1]
|
||
最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
|
||
"""
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
first_group_skipped = False # 标记是否已跳过第一个 group
|
||
|
||
for record in records:
|
||
is_start = (
|
||
record.get("type") == "table"
|
||
and is_operation_table_by_first_tr(record.get("table_body", ""))
|
||
)
|
||
|
||
if is_start:
|
||
if current_group is not None:
|
||
# 决定是否保留当前积累的 group
|
||
if not first_group_skipped:
|
||
# 跳过第一个 group：不清空 current_group，也不加入 groups 或 used_ids
|
||
first_group_skipped = True
|
||
# 注意：这里不把 current_group 加入 groups，也不更新 used_ids
|
||
else:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
# 开启新 group
|
||
current_group = [record]
|
||
else:
|
||
# 第一个起始点
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
current_group.append(record)
|
||
|
||
# 处理最后一个 group（如果存在）
|
||
if current_group is not None:
|
||
if not first_group_skipped:
|
||
# 整个数据中只有一个 group，且是第一个 → 跳过它
|
||
pass # 不加入 groups，也不加 id 到 used_ids
|
||
else:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# —————— 后处理：截断最后一个 group（如果需要） ——————
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
if len(last_group) > max_following_for_last + 1: # +1 是起始记录本身
|
||
truncated = last_group[:max_following_for_last + 1]
|
||
removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
def extract_wxanli(records):
|
||
keywords = [
|
||
"一、故障现象",
|
||
"单位:",
|
||
"故障名称:",
|
||
"一、基本情况",
|
||
"故障现象及原因:",
|
||
"故障名称：",
|
||
"实例故障"
|
||
]
|
||
|
||
# 存储切片
|
||
slices = []
|
||
# 标记已处理的记录索引
|
||
processed_indices = set()
|
||
|
||
# 查找所有包含关键词的记录位置
|
||
keyword_positions = []
|
||
for i, record in enumerate(records):
|
||
if record.get("type") == "text":
|
||
text = record.get("text", "")
|
||
# 检查是否包含任一关键词
|
||
if any(keyword in text for keyword in keywords):
|
||
keyword_positions.append(i)
|
||
|
||
# 为每个关键词位置创建切片
|
||
for idx, start_pos in enumerate(keyword_positions):
|
||
# 确定切片结束位置
|
||
if idx + 1 < len(keyword_positions):
|
||
end_pos = keyword_positions[idx + 1]
|
||
else:
|
||
end_pos = len(records)
|
||
|
||
# 收集当前切片
|
||
slice_records = []
|
||
|
||
# 添加前2个记录（如果存在且未超出边界）
|
||
for offset in [2, 1]:
|
||
prev_idx = start_pos - offset
|
||
if prev_idx >= 0:
|
||
slice_records.append(records[prev_idx])
|
||
processed_indices.add(prev_idx)
|
||
|
||
# 添加从关键词位置到下一个关键词位置的所有记录
|
||
for i in range(start_pos, end_pos):
|
||
slice_records.append(records[i])
|
||
processed_indices.add(i)
|
||
|
||
slices.append(slice_records)
|
||
|
||
# 收集未处理的剩余记录
|
||
remaining_records = [
|
||
records[i] for i in range(len(records))
|
||
if i not in processed_indices
|
||
]
|
||
|
||
return slices, remaining_records
|
||
|
||
def split_content_bbox(data, max_length=8000):
|
||
contents = []
|
||
highlight_lists = []
|
||
|
||
current_content = ""
|
||
current_highlights = []
|
||
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
if not ins.get("bbox"):
|
||
continue
|
||
# 提取通用字段
|
||
bbox = ins.get("bbox")
|
||
page_idx = ins.get("page_idx", -1)
|
||
|
||
if not is_supported_chunk_record(ins):
|
||
continue
|
||
|
||
new_text = record_to_chunk_text(ins)
|
||
if new_text.strip():
|
||
new_text += "\n"
|
||
should_record_highlight = True
|
||
|
||
# --- 关键：即使 new_text 为空，只要 should_record_highlight 为 True，就要处理 chunk 切分和 highlight 记录 ---
|
||
|
||
# 构造 highlight 项（即使 new_text 为空）
|
||
highlight_item = {"page_idx": page_idx, "bbox": bbox}
|
||
|
||
if not new_text.strip() and not should_record_highlight:
|
||
continue # 完全跳过
|
||
|
||
# 检查是否需要切分（仅当 new_text 非空时才影响长度）
|
||
would_exceed = len(current_content) + len(new_text) > max_length
|
||
|
||
# 如果 new_text 非空 且 会导致超长，则先保存当前 chunk
|
||
if new_text.strip() and would_exceed:
|
||
if current_content or current_highlights:
|
||
contents.append(current_content.rstrip('\n'))
|
||
highlight_lists.append(current_highlights)
|
||
current_content = new_text
|
||
current_highlights = [highlight_item]
|
||
else:
|
||
# 追加内容（如果 new_text 非空）
|
||
if new_text.strip():
|
||
current_content += new_text
|
||
# 无论如何，只要 should_record_highlight，就加 highlight
|
||
if should_record_highlight:
|
||
current_highlights.append(highlight_item)
|
||
|
||
# 保存最后一段
|
||
if current_content or current_highlights:
|
||
contents.append(current_content.rstrip('\n'))
|
||
highlight_lists.append(current_highlights)
|
||
|
||
return contents, highlight_lists
|
||
|
||
|
||
|
||
def _group_records_without_parent_context(records):
|
||
result = []
|
||
current_group = []
|
||
|
||
for i, ins in enumerate(records):
|
||
if 'text_level' in ins and ins['text_level']:
|
||
# 当遇到新的text_level时，如果current_group非空，则先将其添加到结果中
|
||
if current_group:
|
||
result.append(current_group)
|
||
current_group = [] # 开始新一组
|
||
current_group.append(ins)
|
||
else:
|
||
# 如果当前记录没有text_level或其值为空，则继续添加到当前组
|
||
current_group.append(ins)
|
||
|
||
# 别忘了将最后一个组添加到结果中
|
||
if current_group:
|
||
result.append(current_group)
|
||
|
||
return result
|
||
|
||
|
||
def _get_group_level(group):
|
||
if not group:
|
||
return None
|
||
|
||
level = group[0].get("text_level")
|
||
if isinstance(level, int) and level > 0:
|
||
return level
|
||
return None
|
||
|
||
|
||
def group_records(records):
|
||
"""
|
||
Split records by text_level, then prepend ancestor section content.
|
||
|
||
For example, a text_level=2 section will include the nearest text_level=1
|
||
group content before its own content. A text_level=3 section will include
|
||
the nearest level 1 and level 2 group content before its own content.
|
||
"""
|
||
base_groups = _group_records_without_parent_context(records)
|
||
result = []
|
||
section_stack = [] # [(level, merged_group)]
|
||
|
||
for group in base_groups:
|
||
level = _get_group_level(group)
|
||
if level is None:
|
||
result.append(group)
|
||
continue
|
||
|
||
while section_stack and section_stack[-1][0] >= level:
|
||
section_stack.pop()
|
||
|
||
parent_group = section_stack[-1][1] if section_stack else []
|
||
merged_group = parent_group + group
|
||
result.append(merged_group)
|
||
section_stack.append((level, merged_group))
|
||
|
||
return result
|
||
|
||
|
||
def get_chunk_bbox_otherfile(data, max_length=8000):
|
||
"""
|
||
处理没有 bbox 信息的文件数据。
|
||
每条支持的记录作为一个独立切片，不与其他记录拼接。
|
||
|
||
- 单个 table 内容超过 max_length 时，调用 chunk_reset 在 </tr> 处安全切分，
|
||
切分后的多个 chunk 共享同一组 positions
|
||
- 清理冗余 colspan="1" / rowspan="1" 属性
|
||
- bbox 缺失时填充空列表 []
|
||
|
||
返回结构与 get_chunk_bbox 一致：[{"content": str, "positions": [...]}]
|
||
"""
|
||
slices = []
|
||
|
||
for ins in data:
|
||
if not is_supported_chunk_record(ins):
|
||
continue
|
||
|
||
page_idx = ins.get("page_idx", -1)
|
||
bbox = ins.get("bbox", [])
|
||
content = record_to_chunk_text(ins)
|
||
if not content.strip():
|
||
continue
|
||
positions = [{"page_idx": page_idx, "bbox": bbox}]
|
||
|
||
# 单表超长 → 在 </tr> 处切分，多个 chunk 共享 positions
|
||
if len(content) > max_length:
|
||
sub_chunks = chunk_reset(content, max_length)
|
||
for sub in sub_chunks:
|
||
slices.append({
|
||
"content": sub.rstrip('\n'),
|
||
"positions": positions
|
||
})
|
||
else:
|
||
slices.append({
|
||
"content": content.rstrip('\n'),
|
||
"positions": positions
|
||
})
|
||
|
||
# 按 page_idx 排序，与 get_chunk_bbox 风格一致
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get("positions")
|
||
and isinstance(item["positions"], list)
|
||
and len(item["positions"]) > 0
|
||
and "page_idx" in item["positions"][0]
|
||
]
|
||
valid_slices.sort(key=lambda item: item["positions"][-1]["page_idx"])
|
||
|
||
return valid_slices
|
||
|
||
def get_chunk_bbox(data):
|
||
for idx, item in enumerate(data, start=1):
|
||
item["id"] = idx
|
||
repair_data, clean_repair_records = extract_maintenance_groups_and_remaining(data,max_following_for_last=23)
|
||
operation_data, clean_operation_records = extract_operation_groups_and_remaining(data,max_following_for_last=23)
|
||
wxanli_data, clean_wxanli_data = extract_wxanli(data)
|
||
slices = []
|
||
if len(repair_data) > 0 :
|
||
for ins in repair_data:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
groups = group_records(clean_repair_records)
|
||
for ins in groups:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
elif len(operation_data) > 0:
|
||
for ins in operation_data:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
groups = group_records(clean_operation_records)
|
||
for ins in groups:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
|
||
elif len(wxanli_data) > 0:
|
||
for ins in wxanli_data:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
for ins in clean_wxanli_data:
|
||
if not ins.get("bbox"):
|
||
continue
|
||
page_idx = ins.get("page_idx", 0)
|
||
bbox = ins["bbox"]
|
||
positions = [{"page_idx": page_idx, "bbox": bbox}]
|
||
|
||
if not is_supported_chunk_record(ins):
|
||
continue
|
||
|
||
content = record_to_chunk_text(ins)
|
||
if not content.strip():
|
||
continue
|
||
|
||
slices.append({"content": content, "positions": positions})
|
||
|
||
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
else:
|
||
groups = group_records(data)
|
||
for ins in groups:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
return slices
|
||
|
||
def split_text_preserve_sentences(text, max_length=7500):
|
||
if len(text) <= max_length:
|
||
return [text]
|
||
|
||
# 按句子分割，保留标点
|
||
parts = re.split(r'([。！？；!?])', text)
|
||
sentences = []
|
||
for i in range(0, len(parts), 2):
|
||
s = parts[i] + (parts[i+1] if i+1 < len(parts) else '')
|
||
sentences.append(s)
|
||
|
||
chunks = []
|
||
current = ""
|
||
|
||
for sent in sentences:
|
||
if len(current) + len(sent) <= max_length:
|
||
current += sent
|
||
else:
|
||
if current:
|
||
chunks.append(current)
|
||
else:
|
||
# 单句超长，硬切（极端情况）
|
||
chunks.append(sent[:max_length])
|
||
sent = sent[max_length:]
|
||
current = sent
|
||
|
||
if current:
|
||
chunks.append(current)
|
||
|
||
return chunks
|
||
|
||
async def data_replace(data, prefix):
|
||
"""
|
||
遍历 data 列表，若元素包含 'img_path' 字段，则移除前缀 'images/' 并拼接新前缀。
|
||
使用 pathlib 安全处理路径。
|
||
"""
|
||
for ins in data:
|
||
if isinstance(ins, dict) and "img_path" in ins:
|
||
img_path = ins["img_path"]
|
||
local_image_paths = []
|
||
if isinstance(img_path, list):
|
||
markdown_images = []
|
||
for item in img_path:
|
||
if not str(item).strip():
|
||
continue
|
||
markdown_image, local_image_path = build_image_markdown(str(item), prefix)
|
||
markdown_images.append(markdown_image)
|
||
if local_image_path:
|
||
local_image_paths.append(local_image_path)
|
||
ins["img_path"] = "\n".join(markdown_images)
|
||
else:
|
||
markdown_image, local_image_path = build_image_markdown(str(img_path), prefix)
|
||
ins["img_path"] = markdown_image
|
||
if local_image_path:
|
||
local_image_paths.append(local_image_path)
|
||
|
||
ocr_texts = await asyncio.gather(
|
||
*(get_image_ocr_text(path) for path in local_image_paths)
|
||
)
|
||
append_caption_ocr(ins, ocr_texts)
|
||
return data
|
||
|
||
def filter_content_bbox(data, max_length=3000):
|
||
"""
|
||
将 data 中的 text/table/image 内容按字符长度分块（每块不超过 max_length），
|
||
返回 [{"content": str, "highlight_list": list}, ...]
|
||
"""
|
||
chunks = []
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
def _flush_chunk():
|
||
"""内部函数：将当前内容打包进 chunks，并重置"""
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list.copy()
|
||
})
|
||
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
|
||
if not is_supported_chunk_record(ins):
|
||
continue
|
||
|
||
added_content = record_to_chunk_text(ins)
|
||
if added_content.strip():
|
||
added_content += "\n"
|
||
else:
|
||
continue
|
||
|
||
# 检查加上 added_content 后是否超限
|
||
if len(current_content) + len(added_content) > max_length:
|
||
# 如果当前块非空，先 flush
|
||
if current_content:
|
||
_flush_chunk()
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
# 如果单个 added_content 超过 max_length，仍需保留（避免丢弃）
|
||
if len(added_content) > max_length:
|
||
# 强制截断或直接放入（这里选择直接放入，避免信息丢失）
|
||
current_content = added_content[:max_length]
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
_flush_chunk()
|
||
# 剩余部分？此处简化处理：忽略超长部分（或可递归拆分，但复杂）
|
||
continue
|
||
else:
|
||
current_content = added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
else:
|
||
current_content += added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
|
||
# 处理最后一块
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list
|
||
})
|
||
|
||
return chunks
|
||
|
||
def chunk_reset(content,MAX_length=8000):
|
||
chunks = []
|
||
start = 0
|
||
while start < len(content):
|
||
end = start + MAX_length
|
||
if end >= len(content):
|
||
chunks.append(content[start:])
|
||
break
|
||
# 尽量在 </tr> 后面切分
|
||
snippet = content[start:end]
|
||
last_tr_end = snippet.rfind("</tr>")
|
||
if last_tr_end != -1:
|
||
end = start + last_tr_end + len("</tr>")
|
||
else:
|
||
# 如果找不到 </tr>，退回到最近的换行或空白
|
||
last_newline = snippet.rfind("\n")
|
||
if last_newline != -1:
|
||
end = start + last_newline
|
||
else:
|
||
# 实在不行就硬切
|
||
pass
|
||
chunks.append(content[start:end])
|
||
start = end
|
||
return chunks
|
||
def chunk_check(data,MAX_length=8000):
|
||
slice_reset = []
|
||
for ins in data:
|
||
content = ins["content"]
|
||
positions = ins["positions"]
|
||
if len(content) <= 8000:
|
||
slice_reset.append(ins)
|
||
else:
|
||
# 使用你的 chunk_reset 函数切分 content
|
||
chunks = chunk_reset(content, MAX_length)
|
||
for chunk in chunks:
|
||
slice_reset.append({
|
||
"content": chunk,
|
||
"positions": positions # 共享原始位置信息
|
||
})
|
||
return slice_reset
|
||
|
||
def find_ship_info_by_hull(json_file_path, data):
|
||
"""
|
||
从NER结果中提取舷号，再从JSON文件中查找对应的舰船信息
|
||
:param json_file_path: JSON文件路径
|
||
:param data: get_entity() 返回的NER结果字典
|
||
:return: 包含 model_name 和 ship_name 的字典，未找到返回 None
|
||
"""
|
||
# 1. 提取舷号
|
||
hulls = [
|
||
entity['properties']['舷号']
|
||
for entity in data.get('entities', [])
|
||
if 'properties' in entity and '舷号' in entity['properties']
|
||
]
|
||
if not hulls:
|
||
return None
|
||
|
||
target_hull = hulls[0]
|
||
|
||
# 2. 根据舷号查找舰船信息
|
||
try:
|
||
with open(json_file_path, 'r', encoding='utf-8') as f:
|
||
ship_data = json.load(f)
|
||
|
||
for item in ship_data:
|
||
if str(item.get('hull_number')) == str(target_hull):
|
||
return {
|
||
"model_name": item.get('model_name'),
|
||
"ship_name": item.get('ship_name')
|
||
}
|
||
|
||
return None
|
||
|
||
except FileNotFoundError:
|
||
print(f"错误：找不到文件 {json_file_path}")
|
||
return None
|
||
except json.JSONDecodeError:
|
||
print("错误：JSON 文件格式不正确")
|
||
return None
|
||
# def merge_short_slices(slices, min_length=30,filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
|
||
"""
|
||
合并过短的切片：
|
||
- 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
|
||
- 若处于末尾无下一个切片，则反向合并到上一个切片末尾
|
||
- content 用换行拼接，positions 顺序拼接
|
||
|
||
Args:
|
||
slices: [{"content": str, "positions": [...]}]
|
||
min_length: 短切片的字符长度阈值（含）
|
||
filename: 用于实体提取的文件名
|
||
|
||
Returns:
|
||
合并后的 slices 列表
|
||
"""
|
||
if not slices:
|
||
return slices
|
||
|
||
result = []
|
||
pending_contents = [] # 缓存等待合并到"下一个"的短切片 content
|
||
pending_positions = [] # 缓存对应的 positions
|
||
|
||
for ins in slices:
|
||
content = ins.get("content", "") or ""
|
||
positions = ins.get("positions", []) or []
|
||
|
||
if len(content) <= min_length:
|
||
# 暂存，等到下一个正常长度的切片再合并
|
||
pending_contents.append(content)
|
||
pending_positions.extend(positions)
|
||
else:
|
||
# 正常切片：把暂存的短切片合并到它的开头
|
||
if pending_contents:
|
||
merged_prefix = "\n".join(pending_contents)
|
||
content = merged_prefix + ("\n" if merged_prefix else "") + content
|
||
positions = pending_positions + positions
|
||
pending_contents = []
|
||
pending_positions = []
|
||
|
||
result.append({
|
||
"content": content,
|
||
"positions": positions
|
||
})
|
||
|
||
# 收尾：如果末尾还有未合并的短切片（后面没有正常切片可合并）
|
||
# 则反向合并到上一个切片末尾
|
||
if pending_contents:
|
||
merged_suffix = "\n".join(pending_contents)
|
||
if result:
|
||
last = result[-1]
|
||
last["content"] = (last["content"] or "") + ("\n" if last["content"] else "") + merged_suffix
|
||
last["positions"] = (last["positions"] or []) + pending_positions
|
||
else:
|
||
# 极端情况：所有切片都很短，整体作为一个切片返回
|
||
result.append({
|
||
"content": merged_suffix,
|
||
"positions": pending_positions
|
||
})
|
||
|
||
try:
|
||
final_result = get_entity(filename)
|
||
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
|
||
if xinghao is None: # 确保 xinghao 为 None 时不会报错
|
||
xinghao = {"model_name": "", "ship_name": ""}
|
||
print("未找到匹配的舰船信息")
|
||
|
||
# 安全处理 xinghao 为 None 的情况
|
||
model_name = ""
|
||
if xinghao and 'model_name' in xinghao:
|
||
model_name = xinghao['model_name']
|
||
|
||
other_data = format_entity_text(final_result)
|
||
logger.info(f"文件名实体提取成功: {other_data}")
|
||
except Exception as e:
|
||
logger.warning(f"文件名实体提取失败: {e}")
|
||
other_data = ""
|
||
model_name = "" # 确保异常时 model_name 有定义
|
||
|
||
# 将提取的信息（如舰艇名、型号）注入到每个切片的开头
|
||
for ins in result:
|
||
content = ins['content']
|
||
# 如果没有提取到有效数据，则跳过注入
|
||
|
||
newline_idx = content.find('\n')
|
||
info = other_data
|
||
if model_name:
|
||
info += f'，型号为{model_name}'
|
||
suffix = f'({info})'
|
||
|
||
# 将信息插入到第一行末尾
|
||
if newline_idx == -1:
|
||
ins['content'] = content + suffix
|
||
else:
|
||
ins['content'] = content[:newline_idx] + suffix + content[newline_idx:]
|
||
|
||
return result
|
||
|
||
|
||
def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
|
||
"""
|
||
合并过短的切片：
|
||
- 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
|
||
- 若处于末尾无下一个切片，则反向合并到上一个切片末尾
|
||
- content 用换行拼接，positions 顺序拼接
|
||
|
||
Args:
|
||
slices: [{"content": str, "positions": [...]}]
|
||
min_length: 短切片的字符长度阈值（含）
|
||
filename: 用于实体提取的文件名
|
||
|
||
Returns:
|
||
合并后的 slices 列表
|
||
"""
|
||
if not slices:
|
||
return slices
|
||
|
||
result = []
|
||
pending_contents = [] # 缓存等待合并到"下一个"的短切片 content
|
||
pending_positions = [] # 缓存对应的 positions
|
||
|
||
for ins in slices:
|
||
content = ins.get("content", "") or ""
|
||
positions = ins.get("positions", []) or []
|
||
|
||
if len(content) <= min_length:
|
||
# 暂存，等到下一个正常长度的切片再合并
|
||
pending_contents.append(content)
|
||
pending_positions.extend(positions)
|
||
else:
|
||
# 正常切片：把暂存的短切片合并到它的开头
|
||
if pending_contents:
|
||
merged_prefix = "\n".join(pending_contents)
|
||
content = merged_prefix + ("\n" if merged_prefix else "") + content
|
||
positions = pending_positions + positions
|
||
pending_contents = []
|
||
pending_positions = []
|
||
|
||
result.append({
|
||
"content": content,
|
||
"positions": positions
|
||
})
|
||
|
||
# 收尾：如果末尾还有未合并的短切片（后面没有正常切片可合并）
|
||
# 则反向合并到上一个切片末尾
|
||
if pending_contents:
|
||
merged_suffix = "\n".join(pending_contents)
|
||
if result:
|
||
last = result[-1]
|
||
last["content"] = (last["content"] or "") + ("\n" if last["content"] else "") + merged_suffix
|
||
last["positions"] = (last["positions"] or []) + pending_positions
|
||
else:
|
||
# 极端情况：所有切片都很短，整体作为一个切片返回
|
||
result.append({
|
||
"content": merged_suffix,
|
||
"positions": pending_positions
|
||
})
|
||
|
||
try:
|
||
final_result = get_entity(filename)
|
||
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
|
||
if xinghao is None: # 确保 xinghao 为 None 时不会报错
|
||
xinghao = {"model_name": "", "ship_name": ""}
|
||
print("未找到匹配的舰船信息")
|
||
|
||
# 安全处理 xinghao 为 None 的情况
|
||
model_name = ""
|
||
if xinghao and 'model_name' in xinghao:
|
||
model_name = xinghao['model_name']
|
||
|
||
other_data = format_entity_text(final_result)
|
||
logger.info(f"文件名实体提取成功: {other_data}")
|
||
except Exception as e:
|
||
logger.warning(f"文件名实体提取失败: {e}")
|
||
other_data = ""
|
||
model_name = "" # 确保异常时 model_name 有定义
|
||
|
||
# 将提取的信息（如舰艇名、型号）注入到每个切片的开头
|
||
for ins in result:
|
||
content = ins['content']
|
||
# 如果没有提取到有效数据，则跳过注入
|
||
|
||
newline_idx = content.find('\n')
|
||
info = other_data
|
||
if model_name:
|
||
info += f'，型号为{model_name}'
|
||
|
||
# ===== 修改部分：拼接文件名 + 实体信息 =====
|
||
# 构建前缀：#文件名为：xxx.pdf(实体信息）
|
||
filename_prefix = f"#文件名为：{filename}"
|
||
if info:
|
||
filename_prefix += f"({info}）"
|
||
# 注意：这里使用全角右括号"）"以匹配你的示例，如需半角请改为 ")"
|
||
|
||
# 将前缀插入到 content 的最开头
|
||
if content:
|
||
ins['content'] = filename_prefix + "\n" + content
|
||
else:
|
||
ins['content'] = filename_prefix
|
||
# ==========================================
|
||
|
||
return result
|
||
def find_and_read_content_list(directory, original_filename, encoding='utf-8'):
|
||
"""
|
||
根据原始文件名，在指定目录及其子目录下查找对应的 _content_list.json 文件并返回内容。
|
||
"""
|
||
# 1. 去除原始文件的后缀，并拼接新的后缀
|
||
# 加上 只获取不带后缀的文件名部分
|
||
base_name = os.path.splitext(original_filename)
|
||
file_name = base_name[0]
|
||
target_filename = f"{file_name}_content_list.json"
|
||
|
||
print(f"正在查找文件: {target_filename}")
|
||
|
||
# 2. 使用 os.walk 递归遍历目录
|
||
for root, dirs, files in os.walk(directory):
|
||
if target_filename in files:
|
||
file_path = os.path.join(root, target_filename)
|
||
|
||
try:
|
||
with open(file_path, 'r', encoding=encoding) as f:
|
||
json_data = json.load(f)
|
||
return json_data
|
||
except json.JSONDecodeError:
|
||
print(f"文件 {file_path} 不是有效的 JSON 格式。")
|
||
return None
|
||
except Exception as e:
|
||
print(f"读取文件 {file_path} 时发生错误: {e}")
|
||
return None
|
||
|
||
return None
|
||
|
||
async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Optional[Dict[str, Any]]:
|
||
"""异步调用 PDF 分析服务并处理响应"""
|
||
doc_result = await process_document(filename)
|
||
if doc_result is not None:
|
||
content_list, images = doc_result
|
||
if content_list and images:
|
||
replaced_content_textlevel = reset_textlevel(content_list)
|
||
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
slices_check = merge_short_slices(slices_check, min_length=30,filename=filename)
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
api_url = get_api_url()
|
||
try:
|
||
async with aiofiles.open(pdf_path, "rb") as f:
|
||
file_content = await f.read()
|
||
files = {"file": (filename, file_content, "application/pdf")}
|
||
|
||
async with httpx.AsyncClient(timeout=60 * 20) as client:
|
||
response = await client.post(api_url, files=files)
|
||
|
||
if response.status_code != 200:
|
||
logger.error(f"❌ API 请求失败，状态码: {response.status_code}, 响应: {response.text}")
|
||
return None
|
||
|
||
result = response.json()
|
||
data = result.get("data", {})
|
||
content_list = data.get("content_list", [])
|
||
images = data.get("images", {})
|
||
replaced_content_textlevel = reset_textlevel(content_list)
|
||
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
slices_check = merge_short_slices(slices_check, min_length=30,filename=filename)
|
||
print(11111111111111111111111111111111111111)
|
||
print(len(images))
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
except json.JSONDecodeError:
|
||
logger.error("❌ 响应不是有效的 JSON 格式")
|
||
logger.debug(response.text)
|
||
return {"slices": [], "images": {}}
|
||
except Exception as e:
|
||
logger.error(f"处理 PDF 文件时出错: {e}", exc_info=True)
|
||
return None
|
||
finally:
|
||
release_api_url(api_url)
|
||
|
||
async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -> Optional[Dict[str, Any]]:
|
||
"""异步调用 Office 文件分析服务并处理响应"""
|
||
# 优先尝试走本地/缓存处理（保持与你原逻辑一致）
|
||
doc_result = await process_document(filename)
|
||
if doc_result is not None:
|
||
content_list, images = doc_result
|
||
if content_list and images:
|
||
replaced_content_textlevel = reset_textlevel(content_list)
|
||
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
api_url = get_api_url()
|
||
analyze_other_url = api_url.replace("/analyze-pdf", "/analyze-otherfile")
|
||
|
||
mime_type, _ = mimetypes.guess_type(pdf_path.name)
|
||
content_type = mime_type or "application/octet-stream"
|
||
|
||
try:
|
||
async with aiofiles.open(pdf_path, "rb") as f:
|
||
file_content = await f.read()
|
||
|
||
files = {"file": (pdf_path.name, file_content, content_type)}
|
||
|
||
async with httpx.AsyncClient(timeout=60 * 20) as client:
|
||
response = await client.post(analyze_other_url, files=files)
|
||
|
||
if response.status_code != 200:
|
||
logger.error(f"❌ Office 文件 API 请求失败，状态码: {response.status_code}, 响应: {response.text}")
|
||
return None
|
||
|
||
result = response.json()
|
||
data = result.get("data", {})
|
||
content_list = data.get("content_list", [])
|
||
images = data.get("images", {})
|
||
|
||
# 后续的数据替换与切片处理逻辑
|
||
replaced_content_textlevel = reset_textlevel(content_list)
|
||
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
except json.JSONDecodeError:
|
||
error_text = response.text if 'response' in locals() else "请求未到达服务器"
|
||
logger.error(f"❌ 响应不是有效的 JSON 格式: {error_text}")
|
||
return {"slices": [], "images": {}}
|
||
except Exception as e:
|
||
logger.error(f"处理 Office 文件时出错: {e}", exc_info=True)
|
||
return None
|
||
finally:
|
||
release_api_url(api_url)
|
||
|
||
if __name__ == "__main__":
|
||
|
||
filepath = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\舰船抗沉损管训练仿真系统研究_content_list.json"
|
||
try:
|
||
with open(filepath, 'r', encoding='utf-8') as file:
|
||
data = json.load(file)
|
||
except Exception as e:
|
||
print(f"读取 JSON 文件时出错：{e}")
|
||
data = []
|
||
# print(len(data))
|
||
start_chars = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "0", "第", "十",
|
||
"一", "二", "三", "四", "五", "六", "七", "八", "九"]
|
||
replaced_content_textlevel = reset_textlevel(data)
|
||
data = asyncio.run(data_replace(data=replaced_content_textlevel,prefix="/api/v1/knowledge/files/images/"))
|
||
slices = get_chunk_bbox(data=data)
|
||
slices_1 =chunk_check(slices, 8000)
|
||
slices_2 = merge_short_slices(slices_1, min_length=30) # ← 新增这一行
|
||
|
||
|
||
for ins in slices_2[:30]:
|
||
print(ins)
|