Spaces:
Running
Running
File size: 2,792 Bytes
83dd2a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
import logging
from typing import Dict, Any
import json
from src.models.layoutlm import FormIQModel
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Initialize FastAPI app
app = FastAPI(
title="FormIQ API",
description="Intelligent Document Parser API",
version="1.0.0"
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize model
model = FormIQModel()
@app.get("/")
async def root():
"""Health check endpoint."""
return {"status": "healthy", "service": "FormIQ API"}
@app.post("/extract")
async def extract_information(
file: UploadFile = File(...),
confidence_threshold: float = 0.5,
document_type: str = "invoice"
) -> Dict[str, Any]:
"""Extract information from uploaded document.
Args:
file: Uploaded document image
confidence_threshold: Minimum confidence score for predictions
document_type: Type of document being processed
Returns:
Dictionary containing extracted fields and metadata
"""
try:
# Read and validate image
contents = await file.read()
image = Image.open(io.BytesIO(contents))
if image.mode != "RGB":
image = image.convert("RGB")
# Process image
extraction_results = model.predict(
image=image,
confidence_threshold=confidence_threshold
)
# Validate extraction
validation_results = model.validate_extraction(
extracted_fields=extraction_results,
document_type=document_type
)
# Combine results
response = {
"extraction": extraction_results,
"validation": validation_results,
"metadata": {
"document_type": document_type,
"confidence_threshold": confidence_threshold
}
}
return response
except Exception as e:
logger.error(f"Error processing document: {str(e)}")
raise HTTPException(
status_code=500,
detail=f"Error processing document: {str(e)}"
)
@app.get("/model-info")
async def get_model_info() -> Dict[str, Any]:
"""Get information about the current model."""
return {
"model_name": model.model.config.model_type,
"device": model.device,
"version": "1.0.0"
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000) |