Refactor: migrate pdf_parser.py to golang (#16323)

### What problem does this PR solve?

Http API based on onnx model.
pdf_parser.py to golang

### Type of change

- [x] Refactoring
This commit is contained in:
Jack
2026-06-25 20:16:16 +08:00
committed by GitHub
parent c7052f4dd1
commit 304d9e02bb
98 changed files with 24591 additions and 8 deletions

View File

View File

@@ -0,0 +1,80 @@
"""DLA adapter — wraps LayoutRecognizer and converts output to wire format."""
import io
import logging
from typing import List
from PIL import Image
from deepdoc.vision import LayoutRecognizer
logger = logging.getLogger(__name__)
# OSS model label → Go dlaClassLabels index
# Go-side (internal/parser/deepdoc.go):
# var dlaClassLabels = []string{
# "title", "text", "reference", "figure", "figure caption",
# "table", "table caption", "table caption", "equation", "figure caption",
# }
# Indices 4/6/7/9 are duplicates; OSS model only produces unique labels.
DLA_CLASS_MAP = {
"title": 0,
"text": 1,
"reference": 2,
"figure": 3,
"figure caption": 4,
"table": 5,
"table caption": 6,
"equation": 8,
}
class DLAAdapter:
"""Calls LayoutRecognizer.forward() and converts bboxes to wire format."""
def __init__(self, model_dir: str, thr: float = 0.2):
self.model_dir = model_dir
self.thr = thr
self._layouter: LayoutRecognizer | None = None
def load(self):
"""Initialize the layout recognizer. Called once per worker."""
self._layouter = LayoutRecognizer("layout")
def __call__(self, image_data: bytes) -> List[List[float]]:
"""
Args:
image_data: JPEG image bytes.
Returns:
List of [x0, y0, x1, y1, score, class_id] for each detected layout region.
"""
if self._layouter is None:
raise RuntimeError("DLAAdapter.load() must be called before inference")
img = Image.open(io.BytesIO(image_data)).convert("RGB")
width, height = img.size
# forward() returns raw Recognizer output (no OCR integration)
raw_bboxes = self._layouter.forward([img], thr=self.thr, batch_size=1)[0]
result = []
for b in raw_bboxes:
label = b["type"].lower()
class_id = DLA_CLASS_MAP.get(label)
if class_id is None:
logger.warning("DLA: unknown label '%s', skipping", label)
continue
x0, y0, x1, y1 = b["bbox"]
score = float(b["score"])
# Clamp coordinates
x0 = max(0.0, min(float(x0), width))
y0 = max(0.0, min(float(y0), height))
x1 = max(0.0, min(float(x1), width))
y1 = max(0.0, min(float(y1), height))
result.append([x0, y0, x1, y1, score, float(class_id)])
return result

View File

@@ -0,0 +1,103 @@
"""OCR adapter — wraps OCR model and converts output to wire format.
Two modes:
- detect: 5-level nested JSON matching Go [][][][][]float64
- rec: 4-level nested JSON matching Go [][][][]any
"""
import logging
from typing import Any, Dict
import cv2
import numpy as np
from deepdoc.vision.ocr import OCR
logger = logging.getLogger(__name__)
# Confidence fill value — OSS recognize_batch does not return confidence scores.
_CONFIDENCE_FILL = 1.0
class OCRAdapter:
"""Calls OCR.detect() and OCR.recognize_batch(), converts to wire format."""
def __init__(self, model_dir: str):
self.model_dir = model_dir
self._ocr: OCR | None = None
def load(self):
"""Initialize the OCR model. Called once per worker."""
self._ocr = OCR()
def close(self):
"""Clean up OCR model resources."""
if self._ocr is not None:
try:
# Access internal detectors and recognizers
if hasattr(self._ocr, "detector") and self._ocr.detector is not None:
self._ocr.detector.close()
except Exception:
pass
try:
if hasattr(self._ocr, "text_recognizer") and self._ocr.text_recognizer is not None:
self._ocr.text_recognizer.close()
except Exception:
pass
self._ocr = None
def detect(self, image_data: bytes) -> Dict[str, Any]:
"""Run text detection.
Returns:
{"output": 5-level nested list} matching Go [][][][][]float64.
"""
if self._ocr is None:
raise RuntimeError("OCRAdapter.load() must be called before inference")
img = self._decode_bgr(image_data)
# OCR.detect() → [(quad_ndarray, ("", 0)), ...]
det_result = self._ocr.detect(img)
quads = []
for quad_ndarray, _ in det_result:
quad = quad_ndarray.tolist() # [[x0,y0],[x1,y1],[x2,y2],[x3,y3]]
# Convert to Python float for JSON compatibility
quad = [[float(p[0]), float(p[1])] for p in quad]
quads.append(quad)
# 5-level nesting matching Go [][][][][]float64:
# batch → page → quad → point → coord
output = [[quads]]
return {"output": output}
def recognize(self, image_data: bytes) -> Dict[str, Any]:
"""Run text recognition on a cropped text region.
Returns:
{"output": 4-level nested list} matching Go [][][][]any.
"""
if self._ocr is None:
raise RuntimeError("OCRAdapter.load() must be called before inference")
img = self._decode_bgr(image_data)
# OCR.recognize_batch() returns List[str]; single cropped image → list of 1 image
texts = self._ocr.recognize_batch([img])
items = [[text, _CONFIDENCE_FILL] for text in texts]
# 4-level nesting matching Go [][][][]any:
# batch → page → items list → pair [text, confidence]
output = [[items]]
return {"output": output}
@staticmethod
def _decode_bgr(data: bytes) -> np.ndarray:
"""Decode JPEG bytes to BGR numpy array (OCR expects BGR)."""
arr = np.frombuffer(data, np.uint8)
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
if img is None:
raise ValueError("Failed to decode image")
return img

View File

@@ -0,0 +1,75 @@
"""TSR adapter — wraps TableStructureRecognizer and converts output to wire format."""
import io
import logging
from typing import List
from PIL import Image
from deepdoc.vision.table_structure_recognizer import TableStructureRecognizer
logger = logging.getLogger(__name__)
# OSS model label → Go tsrLabels index (labels are identical)
# Go-side (internal/parser/deepdoc.go):
# var tsrLabels = []string{
# "table", "table column", "table row",
# "table column header", "table projected row header",
# "table spanning cell",
# }
TSR_CLASS_MAP = {
"table": 0,
"table column": 1,
"table row": 2,
"table column header": 3,
"table projected row header": 4,
"table spanning cell": 5,
}
class TSRAdapter:
"""Calls TableStructureRecognizer and converts elements to wire format."""
def __init__(self, model_dir: str, thr: float = 0.2):
self.model_dir = model_dir
self.thr = thr
self._tsr: TableStructureRecognizer | None = None
def load(self):
"""Initialize the TSR model. Called once per worker."""
self._tsr = TableStructureRecognizer()
def __call__(self, image_data: bytes) -> List[List[float]]:
"""
Args:
image_data: JPEG image bytes (cropped table region).
Returns:
List of [x0, y0, x1, y1, score, class_id] for each structural element.
"""
if self._tsr is None:
raise RuntimeError("TSRAdapter.load() must be called before inference")
img = Image.open(io.BytesIO(image_data)).convert("RGB")
width, height = img.size
tables = self._tsr([img], thr=self.thr)
result = []
for tbl_elements in tables:
for elem in tbl_elements:
label = elem["label"]
class_id = TSR_CLASS_MAP.get(label)
if class_id is None:
logger.warning("TSR: unknown label '%s', skipping", label)
continue
x0 = max(0.0, min(float(elem["x0"]), width))
y0 = max(0.0, min(float(elem["top"]), height))
x1 = max(0.0, min(float(elem["x1"]), width))
y1 = max(0.0, min(float(elem["bottom"]), height))
score = float(elem["score"])
result.append([x0, y0, x1, y1, score, float(class_id)])
return result