Update app.py

app.py

@@ -4,6 +4,10 @@ import shutil
 import os
 from fastapi.middleware.cors import CORSMiddleware
 import uuid
+from pdf2image import convert_from_path
+from PIL import Image
+from model_utils import extract_invoice_data_from_image
+
 os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics"
 os.environ["XDG_CACHE_HOME"] = "/tmp"
@@ -13,13 +17,11 @@ os.makedirs("/tmp/matplotlib", exist_ok=True)
 os.makedirs("/tmp/ultralytics", exist_ok=True)
 os.makedirs("/tmp/fontconfig", exist_ok=True)
 
-from model_utils import extract_invoice_data_from_image
-
 app = FastAPI()
 
 UPLOAD_DIR = "/tmp/uploads"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
-ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg"}
+ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg",".pdf"}
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -28,23 +30,38 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+def resize_to_640(img: Image.Image) -> Image.Image:
+    w_percent = 640 / float(img.size[0])
+    h_size = int((float(img.size[1]) * float(w_percent)))
+    return img.resize((640, h_size), Image.LANCZOS)
+
+
 @app.post("/extract-invoice")
 async def extract_invoice(file: UploadFile = File(...)):
     file_ext = os.path.splitext(file.filename)[-1].lower()
 
     if file_ext not in ALLOWED_EXTENSIONS:
-        raise HTTPException(status_code=400, detail="Please upload Jpeg, Jpg or Png images only.")
+        raise HTTPException(status_code=400, detail="Supported formats: .png, .jpg, .jpeg, .pdf")
 
-    # Save file to disk
     unique_filename = f"{uuid.uuid4().hex}{file_ext}"
-    file_location = os.path.join(UPLOAD_DIR, unique_filename)
+    file_path = os.path.join(UPLOAD_DIR, unique_filename)
 
     try:
-        with open(file_location, "wb") as f:
+        # Save uploaded file
+        with open(file_path, "wb") as f:
             shutil.copyfileobj(file.file, f)
 
-        # Process the file with your model
-        extracted_data = extract_invoice_data_from_image(file_location)
+        if file_ext == ".pdf":
+            # Convert PDF's first page to image
+            images = convert_from_path(file_path, dpi=300)
+            img = resize_to_640(images[0])
+            image_path = os.path.join(UPLOAD_DIR, f"{uuid.uuid4().hex}.png")
+            img.save(image_path)
+        else:
+            image_path = file_path
+
+        # Run inference
+        extracted_data = extract_invoice_data_from_image(image_path)
 
         return JSONResponse(content={
             "success": True,
@@ -59,6 +76,8 @@ async def extract_invoice(file: UploadFile = File(...)):
         )
 
     finally:
-        # Clean up the file if it exists
-        if os.path.exists(file_location):
-            os.remove(file_location)
+        # Clean up temp files
+        if os.path.exists(file_path):
+            os.remove(file_path)
+        if file_ext == ".pdf" and os.path.exists(image_path):
+            os.remove(image_path)