Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, APIRouter, File, UploadFile, Form | |
| from typing import Optional | |
| from PIL import Image | |
| import urllib.request | |
| from io import BytesIO | |
| import json | |
| import os | |
| from config import settings | |
| import utils | |
| from routers import inference, training | |
| from routers.donut_inference import process_document_donut | |
| app = FastAPI(openapi_url="/api/v1/sparrow-ml/openapi.json", docs_url="/api/v1/sparrow-ml/docs") | |
| app.include_router(inference.router, prefix="/api-inference/v1/sparrow-ml", tags=["Inference"]) | |
| app.include_router(training.router, prefix="/api-training/v1/sparrow-ml", tags=["Training"]) | |
| router = APIRouter() | |
| def count_values(obj): | |
| if isinstance(obj, dict): | |
| count = 0 | |
| for value in obj.values(): | |
| count += count_values(value) | |
| return count | |
| elif isinstance(obj, list): | |
| count = 0 | |
| for item in obj: | |
| count += count_values(item) | |
| return count | |
| else: | |
| return 1 | |
| async def run_inference(file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None), | |
| shipper_id: int = Form(...), model_in_use: str = Form('donut')): | |
| result = [] | |
| model_url = settings.get_model_url(shipper_id) # Get the correct model URL based on shipper_id | |
| 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, model_url) # Pass model_url to the function | |
| 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, model_url) # Pass model_url to the function | |
| # parse file name from url | |
| 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 | |
| async def get_statistics(): | |
| file_path = settings.inference_stats_file | |
| # Check if the file exists, and read its content | |
| if os.path.exists(file_path): | |
| with open(file_path, 'r') as file: | |
| try: | |
| content = json.load(file) | |
| except json.JSONDecodeError: | |
| content = [] | |
| else: | |
| content = [] | |
| return content |