Spaces:
Sleeping
Sleeping
import logging | |
from fastapi import FastAPI, File, UploadFile, HTTPException | |
from fastapi.middleware.cors import CORSMiddleware | |
from paddleocr import PaddleOCR | |
from doctr.io import DocumentFile | |
from doctr.models import ocr_predictor | |
import numpy as np | |
from PIL import Image | |
import io | |
# Set up logging | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
app = FastAPI() | |
app.add_middleware( | |
CORSMiddleware, | |
allow_origins=["*"], | |
allow_credentials=True, | |
allow_methods=["*"], | |
allow_headers=["*"], | |
) | |
# Initialize models once at startup | |
ocr_model = ocr_predictor(pretrained=True) | |
paddle_ocr = PaddleOCR(lang='en', use_angle_cls=True) | |
def ocr_with_doctr(file): | |
text_output = '' | |
try: | |
logger.info("Processing PDF with Doctr...") | |
doc = DocumentFile.from_pdf(file) | |
result = ocr_model(doc) | |
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" | |
except Exception as e: | |
logger.error(f"Error processing PDF: {e}") | |
raise HTTPException(status_code=500, detail=f"Error processing PDF: {e}") | |
return text_output | |
def ocr_with_paddle(img): | |
finaltext = '' | |
try: | |
logger.info("Processing image with PaddleOCR...") | |
result = paddle_ocr.ocr(img) | |
for i in range(len(result[0])): | |
text = result[0][i][1][0] | |
finaltext += ' ' + text | |
except Exception as e: | |
logger.error(f"Error processing image: {e}") | |
raise HTTPException(status_code=500, detail=f"Error processing image: {e}") | |
return finaltext | |
async def perform_ocr(file: UploadFile = File(...)): | |
try: | |
logger.info(f"Received file: {file.filename}") | |
file_bytes = await file.read() | |
if file.filename.endswith('.pdf'): | |
logger.info("Detected PDF file") | |
text_output = ocr_with_doctr(io.BytesIO(file_bytes)) | |
else: | |
logger.info("Detected image file") | |
img = np.array(Image.open(io.BytesIO(file_bytes))) | |
text_output = ocr_with_paddle(img) | |
logger.info("OCR completed successfully") | |
return {"ocr_text": text_output} | |
except Exception as e: | |
logger.error(f"Internal server error: {e}") | |
raise HTTPException(status_code=500, detail=f"Internal server error: {e}") | |
async def test_call(): | |
return {"message": "Hi. I'm running"} | |
import uvicorn | |
# Main entry point | |
if __name__ == "__main__": | |
uvicorn.run("app:app", host="0.0.0.0", port=23333, reload=False) |