ocr_api / app.py
Genzo1010's picture
Update app.py
125b268 verified
raw
history blame
2.05 kB
import os
import asyncio
from asyncio import BoundedSemaphore
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
import numpy as np
from PIL import Image
from paddleocr import PaddleOCR
from doctr.io import DocumentFile
from doctr.models import ocr_predictor
import io
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"]
)
app.add_middleware(GZipMiddleware, minimum_size=1000)
# Initialize models once at startup
ocr_model = ocr_predictor(pretrained=True)
paddle_ocr = PaddleOCR(lang='en', use_angle_cls=True, use_gpu=False)
# Create a bounded semaphore to limit concurrent requests
semaphore = BoundedSemaphore(10)
async def ocr_with_doctr(file):
async with semaphore:
doc = DocumentFile.from_pdf(file)
result = ocr_model(doc)
text_output = ''
for page in result.pages:
for block in page.blocks:
for line in block.lines:
text_output += " ".join([word.value for word in line.words]) + "\n"
return text_output
async def ocr_with_paddle(img):
async with semaphore:
result = paddle_ocr.ocr(img)
finaltext = ''
for i in range(len(result[0])):
text = result[0][i][1][0]
finaltext +='' + text
return finaltext
async def generate_text_from_image(img):
return await ocr_with_paddle(img)
@app.post("/ocr/")
async def perform_ocr(file: UploadFile = File(...)):
file_bytes = await file.read()
if file.filename.endswith('.pdf'):
text_output = await ocr_with_doctr(io.BytesIO(file_bytes))
else:
img = Image.open(io.BytesIO(file_bytes))
img.thumbnail((1024, 1024)) # Reduce image size
text_output = await generate_text_from_image(img)
return {"ocr_text": text_output}
@app.get("/test/")
async def test_call():
return {"message": "Hi. I'm running"}