if content != image convert it into image
This commit is contained in:
@@ -12,7 +12,8 @@ RUN pip install --no-cache-dir \
|
|||||||
uvicorn \
|
uvicorn \
|
||||||
python-multipart \
|
python-multipart \
|
||||||
Pillow \
|
Pillow \
|
||||||
numpy
|
numpy \
|
||||||
|
pypdfium2
|
||||||
|
|
||||||
COPY server.py /workspace/server.py
|
COPY server.py /workspace/server.py
|
||||||
|
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import logging
|
|||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
import pypdfium2 as pdfium
|
||||||
from fastapi import FastAPI, UploadFile, File
|
from fastapi import FastAPI, UploadFile, File
|
||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
from paddleocr import PaddleOCR
|
from paddleocr import PaddleOCR
|
||||||
@@ -58,32 +59,50 @@ def ocr_upload(file: UploadFile = File(...)):
|
|||||||
return run_ocr(img_bytes)
|
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):
|
def run_ocr(img_bytes: bytes):
|
||||||
try:
|
try:
|
||||||
image = Image.open(BytesIO(img_bytes))
|
if _is_pdf(img_bytes):
|
||||||
if image.mode != "RGB":
|
images = _pdf_to_images(img_bytes)
|
||||||
image = image.convert("RGB")
|
print(f"[OCR] PDF: {len(images)} pages", flush=True)
|
||||||
img_array = np.array(image)
|
else:
|
||||||
print(f"[OCR] Image decoded: {img_array.shape}", flush=True)
|
images = [Image.open(BytesIO(img_bytes)).convert("RGB")]
|
||||||
except Exception as e:
|
|
||||||
print(f"[OCR] Image decode FAILED: {e}", flush=True)
|
|
||||||
return JSONResponse(
|
|
||||||
status_code=400,
|
|
||||||
content={"errorCode": 1, "message": f"failed to decode image: {e}"},
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
result = list(ocr_engine.predict(img_array))
|
|
||||||
|
|
||||||
pages = []
|
pages = []
|
||||||
for r in result:
|
for i, image in enumerate(images):
|
||||||
raw = r._to_json()
|
img_array = np.array(image)
|
||||||
data = raw.get("res", raw)
|
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({
|
pages.append({
|
||||||
"rec_texts": data.get("rec_texts", []),
|
"rec_texts": texts,
|
||||||
"rec_scores": [float(s) for s in data.get("rec_scores", [])],
|
"rec_scores": scores,
|
||||||
"rec_boxes": data.get("rec_boxes", []),
|
"rec_boxes": boxes,
|
||||||
"rec_polys": data.get("rec_polys", []),
|
"rec_polys": polys,
|
||||||
})
|
})
|
||||||
|
|
||||||
return {"errorCode": 0, "result": {"ocrResults": pages}}
|
return {"errorCode": 0, "result": {"ocrResults": pages}}
|
||||||
|
|||||||
Reference in New Issue
Block a user