Update app.py

app.py

@@ -1,8 +1,9 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile,HTTPException
 from fastapi.responses import JSONResponse
 import shutil
 import os
-
+from fastapi.middleware.cors import CORSMiddleware
+import uuid
 os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics"
 os.environ["XDG_CACHE_HOME"] = "/tmp"
@@ -18,15 +19,39 @@ app = FastAPI()
 
 UPLOAD_DIR = "/tmp/uploads"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
+ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg"}
 
 @app.post("/extract-invoice")
 async def extract_invoice(file: UploadFile = File(...)):
-    file_location = os.path.join(UPLOAD_DIR, file.filename)
-    with open(file_location, "wb") as f:
-        shutil.copyfileobj(file.file, f)
-
-    extracted_data = extract_invoice_data_from_image(file_location)
-
-    os.remove(file_location)
-
-    return JSONResponse(content=extracted_data)
+    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.")
+
+    # Save file to disk
+    unique_filename = f"{uuid.uuid4().hex}{file_ext}"
+    file_location = os.path.join(UPLOAD_DIR, unique_filename)
+
+    try:
+        with open(file_location, "wb") as f:
+            shutil.copyfileobj(file.file, f)
+
+        # Process the file with your model
+        extracted_data = extract_invoice_data_from_image(file_location)
+
+        return JSONResponse(content={
+            "success": True,
+            "message": "Invoice data extracted successfully.",
+            "data": extracted_data
+        })
+
+    except Exception as ex:
+        return JSONResponse(
+            status_code=500,
+            content={"success": False, "error": f"Internal Server Error: {str(ex)}"}
+        )
+
+    finally:
+        # Clean up the file if it exists
+        if os.path.exists(file_location):
+            os.remove(file_location)