Merge pull request #332 from tjbck/master

feat: API server file upload support

.gitignore

@@ -9,6 +9,7 @@ wandb
 report.json
 benchmark_data
 debug_data
+uploads
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

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()
+    main()