mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-13 16:38:26 +08:00
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:
0
deepdoc/server/adapters/__init__.py
Normal file
0
deepdoc/server/adapters/__init__.py
Normal file
80
deepdoc/server/adapters/dla_adapter.py
Normal file
80
deepdoc/server/adapters/dla_adapter.py
Normal 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
|
||||
103
deepdoc/server/adapters/ocr_adapter.py
Normal file
103
deepdoc/server/adapters/ocr_adapter.py
Normal 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
|
||||
75
deepdoc/server/adapters/tsr_adapter.py
Normal file
75
deepdoc/server/adapters/tsr_adapter.py
Normal 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
|
||||
Reference in New Issue
Block a user