116 lines
3.3 KiB
Python
116 lines
3.3 KiB
Python
import os
|
|
import base64
|
|
import logging
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
import pypdfium2 as pdfium
|
|
from fastapi import FastAPI, UploadFile, File
|
|
from fastapi.responses import JSONResponse
|
|
from paddleocr import PaddleOCR
|
|
from pydantic import BaseModel
|
|
from PIL import Image
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("paddleocr-server")
|
|
|
|
OCR_LANG = os.getenv("OCR_LANG", "fr")
|
|
|
|
print(f"[INIT] Initializing PaddleOCR (lang={OCR_LANG})...", flush=True)
|
|
ocr_engine = PaddleOCR(
|
|
use_doc_orientation_classify=False,
|
|
use_doc_unwarping=False,
|
|
use_textline_orientation=False,
|
|
lang=OCR_LANG,
|
|
)
|
|
print("[INIT] PaddleOCR ready.", flush=True)
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
class OCRRequest(BaseModel):
|
|
image: str
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "healthy", "service": "PaddleOCR Server"}
|
|
|
|
|
|
@app.post("/ocr")
|
|
def ocr_json(req: OCRRequest):
|
|
print(f"[OCR] Request received, image field length: {len(req.image)}", flush=True)
|
|
try:
|
|
img_bytes = base64.b64decode(req.image)
|
|
except Exception as e:
|
|
print(f"[OCR] Base64 decode failed: {e}", flush=True)
|
|
return JSONResponse(
|
|
status_code=400,
|
|
content={"errorCode": 1, "message": "invalid base64 image"},
|
|
)
|
|
print(f"[OCR] Decoded {len(img_bytes)} bytes, header: {img_bytes[:32].hex()}", flush=True)
|
|
return run_ocr(img_bytes)
|
|
|
|
|
|
@app.post("/ocr/upload")
|
|
def ocr_upload(file: UploadFile = File(...)):
|
|
img_bytes = file.file.read()
|
|
print(f"[OCR] Upload received, {len(img_bytes)} bytes, header: {img_bytes[:32].hex()}", flush=True)
|
|
return run_ocr(img_bytes)
|
|
|
|
|
|
def _is_pdf(data: bytes) -> bool:
|
|
return data[:5] == b"%PDF-"
|
|
|
|
|
|
def _pdf_to_images(pdf_bytes: bytes):
|
|
doc = pdfium.PdfDocument(BytesIO(pdf_bytes))
|
|
images = []
|
|
for i in range(len(doc)):
|
|
page = doc[i]
|
|
bitmap = page.render(scale=3)
|
|
pil_image = bitmap.to_pil()
|
|
images.append(pil_image.convert("RGB"))
|
|
doc.close()
|
|
return images
|
|
|
|
|
|
def run_ocr(img_bytes: bytes):
|
|
try:
|
|
if _is_pdf(img_bytes):
|
|
images = _pdf_to_images(img_bytes)
|
|
print(f"[OCR] PDF: {len(images)} pages", flush=True)
|
|
else:
|
|
images = [Image.open(BytesIO(img_bytes)).convert("RGB")]
|
|
|
|
pages = []
|
|
for i, image in enumerate(images):
|
|
img_array = np.array(image)
|
|
print(f"[OCR] Page {i+1}/{len(images)}: {img_array.shape}", flush=True)
|
|
result = list(ocr_engine.predict(img_array))
|
|
|
|
texts, scores, boxes, polys = [], [], [], []
|
|
for r in result:
|
|
raw = r._to_json()
|
|
data = raw.get("res", raw)
|
|
texts.extend(data.get("rec_texts", []))
|
|
scores.extend(float(s) for s in data.get("rec_scores", []))
|
|
boxes.extend(data.get("rec_boxes", []))
|
|
polys.extend(data.get("rec_polys", []))
|
|
|
|
pages.append({
|
|
"rec_texts": texts,
|
|
"rec_scores": scores,
|
|
"rec_boxes": boxes,
|
|
"rec_polys": polys,
|
|
})
|
|
|
|
return {"errorCode": 0, "result": {"ocrResults": pages}}
|
|
|
|
except Exception as e:
|
|
logger.exception("OCR failed")
|
|
return JSONResponse(
|
|
status_code=500,
|
|
content={"errorCode": 2, "message": str(e)},
|
|
)
|