From eb4f8021223e33dbe00726e66dd7d243440da9db Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 4 Nov 2024 00:05:19 -0800
Subject: [PATCH] feat: API server file upload support

---
 .gitignore       |   1 +
 marker_server.py | 148 +++++++++++++++++++++++++++++++++++++----------
 2 files changed, 119 insertions(+), 30 deletions(-)

diff --git a/.gitignore b/.gitignore
index 933a116e..93a6edea 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,6 +9,7 @@ wandb
 report.json
 benchmark_data
 debug_data
+uploads
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/marker_server.py b/marker_server.py
index 4387838d..58155949 100644
--- a/marker_server.py
+++ b/marker_server.py
@@ -14,12 +14,20 @@
 from typing import Optional, Annotated
 import io
 
-from fastapi import FastAPI, Body
+from fastapi import FastAPI, Body, Form, File, UploadFile, HTTPException
 from marker.convert import convert_single_pdf
 from marker.models import load_all_models
 
 app_data = {}
 
+
+UPLOAD_DIRECTORY = "./uploads"  # Directory to store uploaded files
+
+# Ensure the upload directory exists
+if not os.path.exists(UPLOAD_DIRECTORY):
+    os.makedirs(UPLOAD_DIRECTORY)
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     if app.state.LOCAL:
@@ -33,10 +41,11 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
+
 @app.get("/")
 async def root():
     return HTMLResponse(
-"""
+        """
 <h1>Marker API</h1>
 <ul>
     <li><a href="/docs">API Documentation</a></li>
@@ -48,59 +57,124 @@ async def root():
 
 class CommonParams(BaseModel):
     filepath: Annotated[
-        str,
-        Field(description="The path to the PDF file to convert.")
+        Optional[str], Field(description="The path to the PDF file to convert.")
     ]
     max_pages: Annotated[
         Optional[int],
-        Field(description="The maximum number of pages in the document to convert.", example=None)
+        Field(
+            description="The maximum number of pages in the document to convert.",
+            example=None,
+        ),
     ] = None
     langs: Annotated[
         Optional[str],
-        Field(description="The optional languages to use if OCR is needed, comma separated.  Must be either the names or codes from from https://github.com/VikParuchuri/surya/blob/master/surya/languages.py.", example=None)
+        Field(
+            description="The optional languages to use if OCR is needed, comma separated.  Must be either the names or codes from from https://github.com/VikParuchuri/surya/blob/master/surya/languages.py.",
+            example=None,
+        ),
     ] = None
     force_ocr: Annotated[
         bool,
-        Field(description="Force OCR on all pages of the PDF.  Defaults to False.  This can lead to worse results if you have good text in your PDFs (which is true in most cases).")
+        Field(
+            description="Force OCR on all pages of the PDF.  Defaults to False.  This can lead to worse results if you have good text in your PDFs (which is true in most cases)."
+        ),
     ] = False
     paginate: Annotated[
         bool,
-        Field(description="Whether to paginate the output.  Defaults to False.  If set to True, each page of the output will be separated by a horizontal rule that contains the page number (2 newlines, {PAGE_NUMBER}, 48 - characters, 2 newlines).")
+        Field(
+            description="Whether to paginate the output.  Defaults to False.  If set to True, each page of the output will be separated by a horizontal rule that contains the page number (2 newlines, {PAGE_NUMBER}, 48 - characters, 2 newlines)."
+        ),
     ] = False
     extract_images: Annotated[
         bool,
-        Field(description="Whether to extract images from the PDF.  Defaults to True.  If set to False, no images will be extracted from the PDF.")
+        Field(
+            description="Whether to extract images from the PDF.  Defaults to True.  If set to False, no images will be extracted from the PDF."
+        ),
     ] = True
 
 
 @app.post("/marker")
-async def convert_pdf(
-    params: CommonParams
-):
+async def convert_pdf(params: CommonParams):
     if app.state.LOCAL:
         print(f"Converting {params.filepath} locally.")
-        assert all([
-            params.extract_images is True,
-            params.paginate is False,
-        ]), "Local conversion API does not support image extraction or pagination."
+        assert all(
+            [
+                params.extract_images is True,
+                params.paginate is False,
+            ]
+        ), "Local conversion API does not support image extraction or pagination."
         return await convert_pdf_local(params)
     else:
         print(f"Converting {params.filepath} using the Datalab API.")
         return await convert_pdf_remote(params)
 
 
+@app.post("/marker/upload")
+async def convert_pdf_upload(
+    max_pages: Optional[int] = Form(default=None),
+    langs: Optional[str] = Form(default=None),
+    force_ocr: Optional[bool] = Form(default=False),
+    paginate: Optional[bool] = Form(default=False),
+    extract_images: Optional[bool] = Form(default=True),
+    file: UploadFile = File(
+        ..., description="The PDF file to convert.", media_type="application/pdf"
+    ),
+):
+    params = CommonParams(
+        filepath=None,
+        max_pages=max_pages,
+        langs=langs,
+        force_ocr=force_ocr,
+        paginate=paginate,
+        extract_images=extract_images,
+    )
+
+    print(f"Converting the uploaded PDF file: {file.filename}")
+    return await convert_pdf_from_upload(file, params)
+
+
+async def convert_pdf_from_upload(file: UploadFile, params: CommonParams):
+    # Check that the uploaded file is a PDF
+    if file.content_type != "application/pdf":
+        raise HTTPException(status_code=400, detail="Only PDF files are allowed.")
+
+    try:
+        # Save the uploaded file in the ./uploads directory
+        upload_path = os.path.join(UPLOAD_DIRECTORY, file.filename)
+
+        with open(upload_path, "wb") as upload_file:
+            file_contents = await file.read()
+            upload_file.write(file_contents)
+
+        # Proceed with the conversion using the saved file path
+        params.filepath = upload_path
+        if app.state.LOCAL:
+            return await convert_pdf_local(params)
+        else:
+            return await convert_pdf_remote(params)
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+        }
+    finally:
+        # Clean up the uploaded file after processing
+        if os.path.exists(upload_path):
+            os.remove(upload_path)
+
+
 async def convert_pdf_remote(params: CommonParams):
     with open(params.filepath, "rb") as f:
         filedata = f.read()
 
     filename = os.path.basename(params.filepath)
     form_data = {
-        'file': (filename, filedata, 'application/pdf'),
-        'max_pages': (None, params.max_pages),
-        'langs': (None, params.langs),
-        'force_ocr': (None, params.force_ocr),
-        'paginate': (None, params.paginate),
-        'extract_images': (None, params.extract_images),
+        "file": (filename, filedata, "application/pdf"),
+        "max_pages": (None, params.max_pages),
+        "langs": (None, params.langs),
+        "force_ocr": (None, params.force_ocr),
+        "paginate": (None, params.paginate),
+        "extract_images": (None, params.extract_images),
     }
 
     headers = {"X-API-Key": app.state.API_KEY}
@@ -128,7 +202,7 @@ async def convert_pdf_local(params: CommonParams):
             app_data["models"],
             max_pages=params.max_pages,
             langs=params.langs,
-            ocr_all_pages=params.force_ocr
+            ocr_all_pages=params.force_ocr,
         )
     except Exception as e:
         return {
@@ -146,16 +220,30 @@ async def convert_pdf_local(params: CommonParams):
         "markdown": full_text,
         "images": encoded,
         "metadata": metadata,
-        "success": True
+        "success": True,
     }
 
 
 def main():
-    parser = argparse.ArgumentParser(description='Convert PDFs to markdown.')
-    parser.add_argument('--port', type=int, default=8000, help='Port to run the server on')
-    parser.add_argument('--host', type=str, default="127.0.0.1", help='Host to run the server on')
-    parser.add_argument('--api_key', type=str, default=None, help='API key for the Datalab API.  If not specified, API will run locally.')
-    parser.add_argument("--datalab_url", type=str, default="https://api.datalab.to/api/v1/marker", help="The URL for the Datalab API")
+    parser = argparse.ArgumentParser(description="Convert PDFs to markdown.")
+    parser.add_argument(
+        "--port", type=int, default=8000, help="Port to run the server on"
+    )
+    parser.add_argument(
+        "--host", type=str, default="127.0.0.1", help="Host to run the server on"
+    )
+    parser.add_argument(
+        "--api_key",
+        type=str,
+        default=None,
+        help="API key for the Datalab API.  If not specified, API will run locally.",
+    )
+    parser.add_argument(
+        "--datalab_url",
+        type=str,
+        default="https://api.datalab.to/api/v1/marker",
+        help="The URL for the Datalab API",
+    )
 
     args = parser.parse_args()
 
@@ -172,4 +260,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()