Enhance PDF processing: validate file type, implement temporary file handling, and add preprocessing functionality

src/controllers/_parser_controller.py

@@ -1,9 +1,13 @@
+import os
+
+import aiofiles
 from fastapi import APIRouter, Body, File, HTTPException, UploadFile
 from fastapi.responses import JSONResponse
 from loguru import logger
 from pydantic import BaseModel
 
 from src.services import PDFProcessorService
+from src.utils import PDFPreprocessor
 
 
 class EntityExtractorSchema(BaseModel):
@@ -18,12 +22,35 @@ class ParserController:
         self.router.add_api_route("/pdf", self.parse_pdf, methods=["POST"])
         self.router.add_api_route("/entity", self.extract_entity, methods=["POST"])
 
-    async def parse_pdf(self, file: UploadFile = File(...)):
+    async def parse_pdf(
+        self, file: UploadFile = File(...), preprocess: bool = Body(False)
+    ):
         try:
-            if not file:
-                raise HTTPException(status_code=400, detail="No file uploaded")
+            if not file.filename.lower().endswith(".pdf"):
+                raise HTTPException(
+                    status_code=400,
+                    detail="Invalid file type. Only PDF files are accepted.",
+                )
             if file.content_type != "application/pdf":
-                raise HTTPException(status_code=400, detail="Invalid file type")
+                raise HTTPException(
+                    status_code=400,
+                    detail="Invalid content type. Only PDF files are accepted.",
+                )
+            async with aiofiles.tempfile.NamedTemporaryFile(
+                prefix=file.filename + "_original_", suffix=".pdf", delete=False
+            ) as temp_pdf:
+                content = await file.read()
+                await temp_pdf.write(content)
+                await temp_pdf.flush()
+                await file.seek(0)
+                if preprocess:
+                    logger.info("Preprocessing the uploaded PDF file.")
+                    file = await PDFPreprocessor.preprocess(file=file)
+                    if file is None:
+                        raise HTTPException(
+                            status_code=400,
+                            detail="Failed to preprocess the PDF file.",
+                        )
             async with self.service as processor:
                 extracted_data = await processor.process_pdf(file)
             return JSONResponse(content={"data": extracted_data})
@@ -35,6 +62,9 @@ class ParserController:
                 status_code=500,
                 detail=str(e),
             )
+        finally:
+            if os.path.exists(temp_pdf.name):
+                os.remove(temp_pdf.name)
 
     async def extract_entity(
         self, entity_extractor_schema: EntityExtractorSchema = Body(...)

src/services/_pdf_processor_service.py

@@ -1,4 +1,5 @@
 import asyncio
+import os
 import re
 import tempfile
 from pathlib import Path
@@ -181,6 +182,9 @@ class PDFProcessorService:
         except Exception as e:
             logger.error(f"Error processing PDF: {e}")
             raise
+        finally:
+            if os.path.exists(pdf_path):
+                os.remove(pdf_path)
 
     async def extract_entity(self, text: str):
         logger.debug(f"Extracting entities from text: {text[:100]}...")

src/utils/__init__.py

@@ -1,4 +1,3 @@
-from ._model_manager import model_manager
-from ._text_extractor import TextExtractor
-
-__all__ = ["model_manager", "TextExtractor"]
+from ._model_manager import *
+from ._preprocessor import *
+from ._text_extractor import *

src/utils/_preprocessor.py

@@ -0,0 +1,38 @@
+import io
+
+import fitz
+from fastapi import UploadFile
+from PIL import Image, ImageEnhance, ImageFilter
+
+
+class PDFPreprocessor:
+    @staticmethod
+    async def preprocess(file: UploadFile) -> UploadFile:
+        content = await file.read()
+        doc = fitz.open(stream=content, filetype="pdf")
+        processed_doc = fitz.open()
+        for page in doc:
+            pix = page.get_pixmap(dpi=300, colorspace=fitz.csGRAY)
+            img = Image.open(io.BytesIO(pix.tobytes()))
+            img = img.filter(
+                ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3)
+            )
+            enhancer = ImageEnhance.Contrast(img)
+            img = enhancer.enhance(1.5)
+            buf = io.BytesIO()
+            img.save(buf, format="PNG")
+            buf.seek(0)
+            processed_doc.new_page(width=page.rect.width, height=page.rect.height)
+            processed_doc[-1].insert_image(
+                processed_doc[-1].rect, stream=buf.getvalue()
+            )
+        doc.close()
+        output_buf = io.BytesIO()
+        processed_doc.save(output_buf)
+        processed_doc.close()
+        output_buf.seek(0)
+        return UploadFile(
+            file=output_buf,
+            filename=file.filename,
+            headers=file.headers,
+        )