| from fastapi import FastAPI, UploadFile, File, HTTPException
|
| from paddleocr import PaddleOCR
|
| from PIL import Image
|
| import io, numpy as np
|
| import fitz
|
|
|
| app = FastAPI()
|
| ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False)
|
|
|
| @app.get("/")
|
| def health():
|
| return {"status": "ok"}
|
|
|
| @app.post("/ocr")
|
| async def run_ocr(file: UploadFile = File(...)):
|
| try:
|
| contents = await file.read()
|
| all_texts = []
|
|
|
|
|
| if file.filename.lower().endswith(".pdf"):
|
| pdf = fitz.open(stream=contents, filetype="pdf")
|
| for page in pdf:
|
| pix = page.get_pixmap(dpi=200)
|
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
| img_array = np.array(img)
|
| result = ocr.ocr(img_array, cls=True)
|
| if result and result[0]:
|
| for line in result[0]:
|
| all_texts.append({
|
| "text": line[1][0],
|
| "confidence": round(line[1][1], 4)
|
| })
|
|
|
|
|
| else:
|
| image = Image.open(io.BytesIO(contents)).convert("RGB")
|
| img_array = np.array(image)
|
| result = ocr.ocr(img_array, cls=True)
|
| if result and result[0]:
|
| for line in result[0]:
|
| all_texts.append({
|
| "text": line[1][0],
|
| "confidence": round(line[1][1], 4)
|
| })
|
|
|
| return {"filename": file.filename, "results": all_texts}
|
|
|
| except Exception as e:
|
| raise HTTPException(status_code=500, detail=str(e)) |