Update routers/inference.py

routers/inference.py

@@ -3,67 +3,59 @@ from typing import Optional
 from PIL import Image
 import urllib.request
 from io import BytesIO
-from config import settings
 import utils
 import os
 import json
+from config import Settings
 from routers.donut_inference import process_document_donut
 
-
 router = APIRouter()
 
 def count_values(obj):
     if isinstance(obj, dict):
-        count = 0
-        for value in obj.values():
-            count += count_values(value)
-        return count
+        return sum(count_values(v) for v in obj.values())
     elif isinstance(obj, list):
-        count = 0
-        for item in obj:
-            count += count_values(item)
-        return count
+        return sum(count_values(i) for i in obj)
     else:
         return 1
 
-
 @router.post("/inference")
 async def run_inference(
     file: Optional[UploadFile] = File(None),
     image_url: Optional[str] = Form(None),
-    shipper_id: str = Form(...),
-    model_in_use: Optional[str] = Form(None)
+    model_in_use: str = Form('donut'),
+    shipper_id: str = Form(...)
 ):
-    
+    # Dynamically load config based on shipper ID
+    settings = Settings(shipper_id=shipper_id)
+
     result = []
+    processing_time = 0
+
     if file:
-        # Ensure the uploaded file is a JPG image
         if file.content_type not in ["image/jpeg", "image/jpg"]:
             return {"error": "Invalid file type. Only JPG images are allowed."}
 
         image = Image.open(BytesIO(await file.read()))
-        processing_time = 0
         if model_in_use == 'donut':
-            result, processing_time = process_document_donut(image)
+            result, processing_time = process_document_donut(image, settings)
         utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file.filename, settings.model])
-        print(f"Processing time: {processing_time:.2f} seconds")
     elif image_url:
-        # test image url: https://raw.githubusercontent.com/katanaml/sparrow/main/sparrow-data/docs/input/invoices/processed/images/invoice_10.jpg
         with urllib.request.urlopen(image_url) as url:
             image = Image.open(BytesIO(url.read()))
-
-        processing_time = 0
         if model_in_use == 'donut':
-            result, processing_time = process_document_donut(image)
-        # parse file name from url
+            result, processing_time = process_document_donut(image, settings)
         file_name = image_url.split("/")[-1]
         utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file_name, settings.model])
-        print(f"Processing time inference: {processing_time:.2f} seconds")
     else:
         result = {"info": "No input provided"}
 
-    return result
-
+    return {
+        "shipper_id": shipper_id,
+        "model": settings.model,
+        "processor": settings.processor,
+        "result": result
+    }
 
 @router.get("/statistics")
 async def get_statistics():