From 3b98efbb3dbc77d1eb1a3c295ed3845352a6d05e Mon Sep 17 00:00:00 2001
From: shniubobo <shniubobo@outlook.com>
Date: Fri, 14 Feb 2025 18:37:54 +0000
Subject: [PATCH] feat(web_api): Return images in api

---
 projects/web_api/app.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/projects/web_api/app.py b/projects/web_api/app.py
index b849b778..8ad230f3 100644
--- a/projects/web_api/app.py
+++ b/projects/web_api/app.py
@@ -1,5 +1,7 @@
 import json
 import os
+from base64 import b64encode
+from glob import glob
 from io import StringIO
 from typing import Tuple, Union
 
@@ -136,6 +138,12 @@ def process_pdf(
     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(
     "/pdf_parse",
     tags=["projects"],
@@ -150,6 +158,7 @@ async def pdf_parse(
     return_layout: bool = False,
     return_info: 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
@@ -243,6 +252,14 @@ async def pdf_parse(
             data["info"] = middle_json
         if return_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
 
         # Clean up memory writers