浏览代码

feat(web_api): Return images in api

shniubobo 9 月之前
父节点
当前提交
3b98efbb3d
共有 1 个文件被更改,包括 17 次插入0 次删除
  1. 17 0
      projects/web_api/app.py

+ 17 - 0
projects/web_api/app.py

@@ -1,5 +1,7 @@
 import json
 import json
 import os
 import os
+from base64 import b64encode
+from glob import glob
 from io import StringIO
 from io import StringIO
 from typing import Tuple, Union
 from typing import Tuple, Union
 
 
@@ -136,6 +138,12 @@ def process_pdf(
     return infer_result, pipe_result
     return infer_result, pipe_result
 
 
 
 
+def encode_image(image_path: str) -> str:
+    """Encode image using base64"""
+    with open(image_path, "rb") as f:
+        return b64encode(f.read()).decode()
+
+
 @app.post(
 @app.post(
     "/pdf_parse",
     "/pdf_parse",
     tags=["projects"],
     tags=["projects"],
@@ -150,6 +158,7 @@ async def pdf_parse(
     return_layout: bool = False,
     return_layout: bool = False,
     return_info: bool = False,
     return_info: bool = False,
     return_content_list: bool = False,
     return_content_list: bool = False,
+    return_images: bool = False,
 ):
 ):
     """
     """
     Execute the process of converting PDF to JSON and MD, outputting MD and JSON files
     Execute the process of converting PDF to JSON and MD, outputting MD and JSON files
@@ -243,6 +252,14 @@ async def pdf_parse(
             data["info"] = middle_json
             data["info"] = middle_json
         if return_content_list:
         if return_content_list:
             data["content_list"] = content_list
             data["content_list"] = content_list
+        if return_images:
+            image_paths = glob(f"{output_image_path}/*.jpg")
+            data["images"] = {
+                os.path.basename(
+                    image_path
+                ): f"data:image/jpeg;base64,{encode_image(image_path)}"
+                for image_path in image_paths
+            }
         data["md_content"] = md_content  # md_content is always returned
         data["md_content"] = md_content  # md_content is always returned
 
 
         # Clean up memory writers
         # Clean up memory writers