diff --git a/fileparse_util.py b/fileparse_util.py index d14d1fa..ef9e554 100644 --- a/fileparse_util.py +++ b/fileparse_util.py @@ -3,7 +3,7 @@ import json import base64 import asyncio import logging -from typing import Optional, Tuple, Dict +from typing import Any, Optional, Tuple, Dict import aiofiles # pip install aiofiles from config import SEARCH_DIR,IMAGE_DIR logging.basicConfig(level=logging.INFO) @@ -16,7 +16,7 @@ async def find_and_read_content_list( encoding: str = 'utf-8' ) -> Tuple[Optional[list], Optional[str]]: """根据原始文件名查找并读取对应的 _content_list.json 文件。""" - file_name, _ = os.path.splitext(original_filename) + file_name, _ = os.path.splitext(os.path.basename(original_filename)) target_filename = f"{file_name}_content_list.json" logger.info(f"正在查找文件: {target_filename}") @@ -44,6 +44,65 @@ async def find_and_read_content_list( return None, None +def find_image_record(content_list: Any, image_filename: str) -> Optional[dict]: + """Find an image record by the basename stored in ``img_path``.""" + if not isinstance(content_list, list) or not image_filename: + return None + + target_name = os.path.basename(str(image_filename).strip().replace('\\', '/')) + if not target_name: + return None + + for item in content_list: + if not isinstance(item, dict) or item.get('type') != 'image': + continue + + img_path = item.get('img_path') + if isinstance(img_path, str): + path_name = os.path.basename(img_path.strip().replace('\\', '/')) + if path_name == target_name: + return item + + + return None + + +async def find_image_record_by_pdf( + directory: str, + pdf_filename: str, + image_filename: str, + encoding: str = 'utf-8', +) -> Optional[dict]: + """Find an image record from the content list matching a PDF filename.""" + content_list, _ = await find_and_read_content_list( + directory, pdf_filename, encoding=encoding + ) + return find_image_record(content_list, image_filename) + + +async def extract_image_caption( + search_directory: str, + pdf_filename: str, + image_filename: str, +) -> list[str]: + """?? PDF ??????????? image_caption?""" + image_record = await find_image_record_by_pdf( + directory=search_directory, + pdf_filename=pdf_filename, + image_filename=image_filename, + ) + + if not image_record: + return [] + + image_caption = image_record.get("image_caption", []) + if isinstance(image_caption, str): + return [image_caption] if image_caption.strip() else [] + if isinstance(image_caption, list): + return [str(item).strip() for item in image_caption if str(item).strip()] + return [] + + async def find_and_read_md_file( directory: str, original_filename: str, @@ -198,5 +257,11 @@ async def main(): if __name__ == "__main__": - exit_code = asyncio.run(main()) - exit(exit_code) + result = asyncio.run( + extract_image_caption( + r"E:\ZKYNLP\Hjunproject\project0506\kgrag", + "163-06A0016-B01003_雷达-维修手册.pdf", + "", + ) + ) + print(result)