mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-29 15:31:05 +08:00
76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
|
"""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
|