paddleocr / main.py
Yitshak's picture
Update main.py
c9331c8 verified
Raw
History Blame Contribute Delete
1.74 kB
from fastapi import FastAPI, UploadFile, File, HTTPException
from paddleocr import PaddleOCR
from PIL import Image
import io, numpy as np
import fitz # pymupdf
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 = []
# Handle PDF
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)
})
# Handle image (JPG, PNG, etc)
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))