123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
import io
|
|
import math
|
|
import statistics
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Callable
|
|
|
|
from .engine import EngineLine, OcrEngine
|
|
|
|
|
|
RENDER_DPI = 200
|
|
MAX_RENDER_PIXELS = 25_000_000
|
|
|
|
|
|
class PdfRenderError(ValueError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RenderedPage:
|
|
page: int
|
|
width: int
|
|
height: int
|
|
png: bytes
|
|
image: object
|
|
|
|
|
|
def render_pdf_pages(pdf: bytes, pages: list[int], max_pixels: int = MAX_RENDER_PIXELS) -> list[RenderedPage]:
|
|
import pypdfium2
|
|
|
|
if not pages or pages != sorted(set(pages)) or any(type(page) is not int or page < 1 for page in pages):
|
|
raise PdfRenderError("PDF pages must be unique, ordered, one-based integers")
|
|
try:
|
|
document = pypdfium2.PdfDocument(pdf)
|
|
except Exception as error:
|
|
raise PdfRenderError("PDF cannot be opened for deterministic rendering") from error
|
|
|
|
rendered: list[RenderedPage] = []
|
|
try:
|
|
for page_number in pages:
|
|
if page_number > len(document):
|
|
raise PdfRenderError(f"PDF page {page_number} does not exist")
|
|
page = document[page_number - 1]
|
|
try:
|
|
page_width, page_height = page.get_size()
|
|
width = math.ceil(page_width * RENDER_DPI / 72)
|
|
height = math.ceil(page_height * RENDER_DPI / 72)
|
|
if width * height > max_pixels:
|
|
raise PdfRenderError("Rendered page exceeds the 25 megapixels limit")
|
|
bitmap = page.render(scale=RENDER_DPI / 72)
|
|
try:
|
|
image = bitmap.to_pil()
|
|
output = io.BytesIO()
|
|
image.save(output, format="PNG")
|
|
rendered.append(RenderedPage(page_number, image.width, image.height, output.getvalue(), image))
|
|
finally:
|
|
bitmap.close()
|
|
finally:
|
|
page.close()
|
|
finally:
|
|
document.close()
|
|
return rendered
|
|
|
|
|
|
def _metrics(lines: list[EngineLine], text: str) -> dict[str, int | float]:
|
|
confidences = sorted(line.confidence for line in lines)
|
|
return {
|
|
"lineCount": len(lines),
|
|
"nonWhitespaceCharacters": sum(not character.isspace() for character in text),
|
|
"medianConfidence": statistics.median(confidences) if confidences else 0.0,
|
|
"p10Confidence": confidences[math.floor((len(confidences) - 1) * 0.1)] if confidences else 0.0,
|
|
"lowConfidenceLineRatio": sum(value < 0.5 for value in confidences) / len(confidences) if confidences else 0.0,
|
|
}
|
|
|
|
|
|
def process_pdf(
|
|
job_id: str,
|
|
document_sha256: str,
|
|
pdf: bytes,
|
|
requested_pages: list[int],
|
|
engine: OcrEngine,
|
|
processing_ms: Callable[[int], int] | None = None,
|
|
) -> dict:
|
|
results = []
|
|
for rendered in render_pdf_pages(pdf, requested_pages):
|
|
started = time.perf_counter_ns()
|
|
indexed = list(enumerate(engine.recognize(rendered.image), start=1))
|
|
indexed.sort(key=lambda item: (item[1].bbox[1], item[1].bbox[0], item[0]))
|
|
lines = [line for _, line in indexed]
|
|
text = "\n".join(line.text for line in lines)
|
|
serialized_lines = [
|
|
{
|
|
"lineId": f"p{rendered.page}-l{index:02d}-{'-'.join(map(str, line.bbox))}",
|
|
"text": line.text,
|
|
"confidence": line.confidence,
|
|
"bbox": list(line.bbox),
|
|
}
|
|
for index, line in enumerate(lines, start=1)
|
|
]
|
|
elapsed = math.ceil((time.perf_counter_ns() - started) / 1_000_000)
|
|
results.append({
|
|
"page": rendered.page,
|
|
"width": rendered.width,
|
|
"height": rendered.height,
|
|
"processingMs": processing_ms(rendered.page) if processing_ms else elapsed,
|
|
"text": text,
|
|
"metrics": _metrics(lines, text),
|
|
"lines": serialized_lines,
|
|
})
|
|
return {
|
|
"schemaVersion": "1",
|
|
"jobId": job_id,
|
|
"documentSha256": document_sha256,
|
|
"engine": {
|
|
"name": "paddleocr",
|
|
"version": "3.4.0",
|
|
"runtime": "paddlepaddle-3.2.2",
|
|
"device": "cpu",
|
|
"configVersion": "ocr-v1",
|
|
"dpi": RENDER_DPI,
|
|
},
|
|
"pages": results,
|
|
}
|