Spaces:
Running
Running
File size: 1,899 Bytes
f3270e6 |
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 |
# Copyright (C) 2021-2025, Mindee.
# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status
from app.schemas import KIEElement, KIEIn, KIEOut
from app.utils import get_documents, resolve_geometry
from app.vision import init_predictor
router = APIRouter()
@router.post("/", response_model=list[KIEOut], status_code=status.HTTP_200_OK, summary="Perform KIE")
async def perform_kie(request: KIEIn = Depends(), files: list[UploadFile] = [File(...)]):
"""Runs docTR KIE model to analyze the input image"""
try:
predictor = init_predictor(request)
content, filenames = await get_documents(files)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
out = predictor(content)
results = [
KIEOut(
name=filenames[i],
orientation=page.orientation,
language=page.language,
dimensions=page.dimensions,
predictions=[
KIEElement(
class_name=class_name,
items=[
dict(
value=prediction.value,
geometry=resolve_geometry(prediction.geometry),
objectness_score=round(prediction.objectness_score, 2),
confidence=round(prediction.confidence, 2),
crop_orientation=prediction.crop_orientation,
)
for prediction in page.predictions[class_name]
],
)
for class_name in page.predictions.keys()
],
)
for i, page in enumerate(out.pages)
]
return results
|