169 lines
5.7 KiB
Python
169 lines
5.7 KiB
Python
import os
|
|
import json
|
|
import base64
|
|
import asyncio
|
|
import logging
|
|
from typing import Optional, Tuple, Dict
|
|
import aiofiles # pip install aiofiles
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def find_and_read_content_list(
|
|
directory: str,
|
|
original_filename: str,
|
|
encoding: str = 'utf-8'
|
|
) -> Tuple[Optional[list], Optional[str]]:
|
|
"""根据原始文件名查找并读取对应的 _content_list.json 文件。"""
|
|
file_name, _ = os.path.splitext(original_filename)
|
|
target_filename = f"{file_name}_content_list.json"
|
|
logger.info(f"正在查找文件: {target_filename}")
|
|
|
|
# os.walk 没有原生 async 版本,放到线程里执行避免阻塞事件循环
|
|
def _walk_for_file():
|
|
for root, _, files in os.walk(directory):
|
|
if target_filename in files:
|
|
return os.path.join(root, target_filename)
|
|
return None
|
|
|
|
file_path = await asyncio.to_thread(_walk_for_file)
|
|
if file_path is None:
|
|
return None, None
|
|
|
|
try:
|
|
async with aiofiles.open(file_path, 'r', encoding=encoding) as f:
|
|
text = await f.read()
|
|
# JSON 解析是 CPU 操作,大文件可考虑放线程
|
|
return json.loads(text), file_path
|
|
except json.JSONDecodeError:
|
|
logger.error(f"文件 {file_path} 不是有效的 JSON 格式。")
|
|
return None, None
|
|
except Exception as e:
|
|
logger.error(f"读取文件 {file_path} 时发生错误: {e}")
|
|
return None, None
|
|
|
|
|
|
def extract_all_jpg_filenames(doc_data: list) -> set:
|
|
"""从文档数据中提取所有 .jpg 图片的纯文件名。纯 CPU 操作,不需要 async。"""
|
|
result_set = set()
|
|
for item in doc_data:
|
|
if 'img_path' in item and item['img_path']:
|
|
img_path = item['img_path']
|
|
if img_path.endswith('.jpg'):
|
|
result_set.add(os.path.basename(img_path))
|
|
return result_set
|
|
|
|
|
|
async def _encode_one_image(
|
|
img_file: str,
|
|
local_image_dir: str,
|
|
mime_map: dict,
|
|
) -> Tuple[str, Optional[str]]:
|
|
"""编码单张图片,返回 (文件名, data_uri 或 None)。"""
|
|
img_path = os.path.join(local_image_dir, img_file)
|
|
if not os.path.isfile(img_path):
|
|
logger.warning(f"图片不存在: {img_path}")
|
|
return img_file, None
|
|
|
|
try:
|
|
async with aiofiles.open(img_path, "rb") as f:
|
|
img_bytes = await f.read()
|
|
# base64 编码是 CPU 操作,小文件直接做即可;大文件可放线程
|
|
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
|
|
ext = img_file.lower().rsplit('.', 1)[-1]
|
|
mime_type = f"image/{mime_map.get(ext, 'png')}"
|
|
return img_file, f"data:{mime_type};base64,{img_base64}"
|
|
except Exception as e:
|
|
logger.error(f"编码图片失败 {img_file}: {str(e)}")
|
|
return img_file, None
|
|
|
|
|
|
async def encode_images_to_base64(
|
|
local_image_dir: str,
|
|
referenced_images: set,
|
|
concurrency: int = 32,
|
|
) -> Dict[str, str]:
|
|
"""并发将引用的图片编码为 base64。"""
|
|
if not os.path.exists(local_image_dir) or not referenced_images:
|
|
return {}
|
|
|
|
mime_map = {'jpg': 'jpeg', 'jpeg': 'jpeg', 'png': 'png', 'gif': 'gif', 'webp': 'webp'}
|
|
|
|
# 用 Semaphore 限制并发数,避免一次性打开几百个文件句柄
|
|
sem = asyncio.Semaphore(concurrency)
|
|
|
|
async def _bounded(img_file: str):
|
|
async with sem:
|
|
return await _encode_one_image(img_file, local_image_dir, mime_map)
|
|
|
|
results = await asyncio.gather(*(_bounded(f) for f in referenced_images))
|
|
return {name: data for name, data in results if data is not None}
|
|
|
|
|
|
async def process_document(
|
|
input_file_name: str,
|
|
search_dir: Optional[str] = "/app/mineru_output",
|
|
local_image_dir: Optional[str] = None,
|
|
concurrency: int = 32,
|
|
) -> Optional[Tuple[list, Dict[str, str]]]:
|
|
"""
|
|
异步处理文档:查找 content_list.json,提取图片引用,并发编码为 base64。
|
|
|
|
Args:
|
|
input_file_name: 原始文件名(如 "xxx.pdf")
|
|
search_dir: 搜索 _content_list.json 的根目录
|
|
local_image_dir: 图片所在目录,None 则自动推断
|
|
concurrency: 图片编码的并发数上限,默认 32
|
|
|
|
Returns:
|
|
成功: (content_list, images_dict)
|
|
失败: None
|
|
"""
|
|
content, content_path = await find_and_read_content_list(search_dir, input_file_name)
|
|
if content is None:
|
|
logger.error("❌ 未找到指定的 _content_list.json 文件。")
|
|
return None
|
|
|
|
logger.info(f"✅ 找到文件: {content_path}")
|
|
logger.info(f"内容类型: {type(content).__name__}, 长度: {len(str(content))}")
|
|
|
|
if local_image_dir is None:
|
|
local_image_dir = os.path.join(os.path.dirname(content_path), "images")
|
|
logger.info(f"图片目录: {local_image_dir}")
|
|
|
|
referenced_images = extract_all_jpg_filenames(content)
|
|
logger.info(f"引用图片数量: {len(referenced_images)}")
|
|
|
|
images_dict = await encode_images_to_base64(
|
|
local_image_dir, referenced_images, concurrency=concurrency
|
|
)
|
|
logger.info(f"成功编码 {len(images_dict)} 张图片")
|
|
|
|
return content, images_dict
|
|
|
|
|
|
async def main():
|
|
result = await process_document(
|
|
input_file_name="163-06A0014-B01001_发动机-维修手册.pdf",
|
|
search_dir="/app/mineru_output",
|
|
concurrency=32,
|
|
)
|
|
|
|
if result is None:
|
|
return 1
|
|
|
|
content, images_dict = result
|
|
print(f"\n=== 处理完成 ===")
|
|
print(f"文档段落数: {len(content)}")
|
|
print(f"图片数量: {len(images_dict)}")
|
|
if images_dict:
|
|
first_key = next(iter(images_dict))
|
|
print(f"示例图片 [{first_key}]: {images_dict[first_key][:80]}...")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit_code = asyncio.run(main())
|
|
exit(exit_code)
|