更新 fileparse_util.py

增加图片ocr文本提取功能

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)