432 lines
14 KiB
Python
432 lines
14 KiB
Python
"""Optional Apache ONNX layout region detector (fail-open)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import os
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from app.services import model_paths
|
|
from app.services.convert.idm.model import BBox
|
|
from app.services.convert.layout.coords import PageGeom, pixel_bbox_to_pdf
|
|
from app.services.convert.validation import run_with_timeout
|
|
|
|
_lock = threading.Lock()
|
|
_session = None
|
|
_session_path: str | None = None
|
|
_load_error: str | None = None
|
|
|
|
# PP-DocLayoutV3 (Paddle) class list — see models/layout inference.yml / HF card.
|
|
_PP_DOCLAYOUT_V3_LABELS = (
|
|
"abstract",
|
|
"algorithm",
|
|
"aside_text",
|
|
"chart",
|
|
"content",
|
|
"display_formula",
|
|
"doc_title",
|
|
"figure_title",
|
|
"footer",
|
|
"footer_image",
|
|
"footnote",
|
|
"formula_number",
|
|
"header",
|
|
"header_image",
|
|
"image",
|
|
"inline_formula",
|
|
"number",
|
|
"paragraph_title",
|
|
"reference",
|
|
"reference_content",
|
|
"seal",
|
|
"table",
|
|
"text",
|
|
"vertical_text",
|
|
"vision_footnote",
|
|
)
|
|
|
|
# Map raw detector labels → DocQube Region.label
|
|
_LABEL_MAP = {
|
|
0: "text",
|
|
1: "title",
|
|
2: "list",
|
|
3: "table",
|
|
4: "figure",
|
|
5: "header",
|
|
6: "footer",
|
|
"text": "text",
|
|
"title": "title",
|
|
"list": "list",
|
|
"table": "table",
|
|
"figure": "figure",
|
|
"header": "header",
|
|
"footer": "footer",
|
|
"Text": "text",
|
|
"Title": "title",
|
|
"List": "list",
|
|
"Table": "table",
|
|
"Figure": "figure",
|
|
"Header": "header",
|
|
"Footer": "footer",
|
|
# PP-DocLayoutV3 names
|
|
"abstract": "text",
|
|
"algorithm": "text",
|
|
"aside_text": "text",
|
|
"chart": "figure",
|
|
"content": "text",
|
|
"display_formula": "text",
|
|
"doc_title": "title",
|
|
"figure_title": "title",
|
|
"footer_image": "figure",
|
|
"footnote": "footer",
|
|
"formula_number": "text",
|
|
"header_image": "figure",
|
|
"image": "figure",
|
|
"inline_formula": "text",
|
|
"number": "text",
|
|
"paragraph_title": "title",
|
|
"reference": "text",
|
|
"reference_content": "text",
|
|
"seal": "figure",
|
|
"vertical_text": "text",
|
|
"vision_footnote": "footer",
|
|
}
|
|
|
|
# Default input size for PP-DocLayoutV3 ONNX (inference.yml: Resize 800x800, keep_ratio false)
|
|
_DEFAULT_TARGET = 800
|
|
|
|
|
|
@dataclass
|
|
class Region:
|
|
label: str
|
|
bbox_pdf: BBox
|
|
score: float
|
|
reading_index: int = 0
|
|
|
|
|
|
def layout_ml_enabled() -> bool:
|
|
return os.environ.get("CONVERT_LAYOUT_ML", "0").strip().lower() in ("1", "true", "yes", "on")
|
|
|
|
|
|
def _weights_path() -> Path:
|
|
explicit = (os.environ.get("CONVERT_LAYOUT_ML_WEIGHTS") or "").strip()
|
|
if explicit:
|
|
return model_paths.resolve(explicit)
|
|
int8_cand = model_paths.resolve("models/layout/v1/layout_int8.onnx")
|
|
if int8_cand.is_file():
|
|
return int8_cand
|
|
return model_paths.resolve("models/layout/v1/layout.onnx")
|
|
|
|
|
|
def _min_score() -> float:
|
|
try:
|
|
return float(os.environ.get("CONVERT_LAYOUT_ML_MIN_SCORE", "0.5"))
|
|
except ValueError:
|
|
return 0.5
|
|
|
|
|
|
def _device() -> str:
|
|
return os.environ.get("CONVERT_LAYOUT_ML_DEVICE", "cpu").strip().lower()
|
|
|
|
|
|
def _page_timeout() -> float:
|
|
try:
|
|
return float(os.environ.get("CONVERT_LAYOUT_ML_PAGE_TIMEOUT", "15"))
|
|
except ValueError:
|
|
return 15.0
|
|
|
|
|
|
def _target_size() -> int:
|
|
try:
|
|
return max(64, int(os.environ.get("CONVERT_LAYOUT_ML_INPUT_SIZE", str(_DEFAULT_TARGET))))
|
|
except ValueError:
|
|
return _DEFAULT_TARGET
|
|
|
|
|
|
def _get_session():
|
|
global _session, _session_path, _load_error
|
|
path = str(_weights_path())
|
|
with _lock:
|
|
if _session is not None and _session_path == path:
|
|
return _session
|
|
if _load_error and _session_path == path:
|
|
return None
|
|
try:
|
|
import onnxruntime as ort
|
|
|
|
if not Path(path).is_file():
|
|
_load_error = f"weights missing: {path}"
|
|
_session = None
|
|
_session_path = path
|
|
return None
|
|
|
|
sess_options = ort.SessionOptions()
|
|
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
|
# Limit ONNX intra-op threads to prevent CPU starvation for concurrent requests
|
|
max_threads = max(1, min(4, os.cpu_count() or 2))
|
|
sess_options.intra_op_num_threads = max_threads
|
|
|
|
providers = ["CPUExecutionProvider"]
|
|
if _device() == "cuda":
|
|
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
_session = ort.InferenceSession(path, sess_options=sess_options, providers=providers)
|
|
_session_path = path
|
|
_load_error = None
|
|
return _session
|
|
except Exception as exc:
|
|
_load_error = str(exc)
|
|
_session = None
|
|
_session_path = path
|
|
return None
|
|
|
|
|
|
def reset_session_for_tests() -> None:
|
|
global _session, _session_path, _load_error
|
|
with _lock:
|
|
_session = None
|
|
_session_path = None
|
|
_load_error = None
|
|
|
|
|
|
def last_load_error() -> str | None:
|
|
return _load_error
|
|
|
|
|
|
def _map_label(cls: int | str) -> str:
|
|
if isinstance(cls, (int, float)):
|
|
idx = int(cls)
|
|
if 0 <= idx < len(_PP_DOCLAYOUT_V3_LABELS):
|
|
name = _PP_DOCLAYOUT_V3_LABELS[idx]
|
|
return _LABEL_MAP.get(name, "text")
|
|
return _LABEL_MAP.get(idx, "text")
|
|
return _LABEL_MAP.get(str(cls), _LABEL_MAP.get(cls, "text"))
|
|
|
|
|
|
def preprocess_pp_doclayout(
|
|
rgb_hwc: np.ndarray,
|
|
*,
|
|
target_size: int = _DEFAULT_TARGET,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""
|
|
PP-DocLayoutV3 preprocess (matches official inference.yml):
|
|
Resize 800x800 keep_ratio=false, NormalizeImage mean=0 std=1 (÷255), Permute NCHW.
|
|
|
|
Returns (image NCHW float32, im_shape [1,2], scale_factor [1,2] as [scale_h, scale_w]).
|
|
"""
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
if rgb_hwc.ndim != 3 or rgb_hwc.shape[2] != 3:
|
|
raise ValueError(f"expected HWC RGB, got shape {getattr(rgb_hwc, 'shape', None)}")
|
|
orig_h, orig_w = int(rgb_hwc.shape[0]), int(rgb_hwc.shape[1])
|
|
target = int(target_size)
|
|
if rgb_hwc.dtype != np.uint8:
|
|
# Accept float 0-1 or 0-255
|
|
scaled = rgb_hwc * 255.0 if float(np.nanmax(rgb_hwc)) <= 1.5 else rgb_hwc
|
|
img = Image.fromarray(np.clip(scaled, 0, 255).astype(np.uint8))
|
|
else:
|
|
img = Image.fromarray(rgb_hwc)
|
|
resized = img.resize((target, target), Image.BILINEAR)
|
|
blob = np.asarray(resized).astype("float32") / 255.0
|
|
image = np.transpose(blob, (2, 0, 1))[None, ...].astype("float32")
|
|
scale_h = float(target) / float(orig_h)
|
|
scale_w = float(target) / float(orig_w)
|
|
im_shape = np.array([[float(target), float(target)]], dtype=np.float32)
|
|
scale_factor = np.array([[scale_h, scale_w]], dtype=np.float32)
|
|
return image, im_shape, scale_factor
|
|
|
|
|
|
def build_onnx_feed(sess, rgb_hwc: np.ndarray, *, target_size: int | None = None) -> dict:
|
|
"""
|
|
Build ORT feed dict from session input metadata.
|
|
|
|
PP-DocLayoutV3 expects: im_shape, image, scale_factor.
|
|
Simpler single-input models get NCHW (or NHWC if shape[-1]==3) image only.
|
|
"""
|
|
import numpy as np
|
|
|
|
t = _target_size() if target_size is None else int(target_size)
|
|
image, im_shape, scale_factor = preprocess_pp_doclayout(rgb_hwc, target_size=t)
|
|
inputs = list(sess.get_inputs())
|
|
names = {i.name for i in inputs}
|
|
feed: dict = {}
|
|
|
|
# Named PP-DocLayout / Paddle Detection I/O
|
|
if "image" in names or "scale_factor" in names or "im_shape" in names:
|
|
if "image" in names:
|
|
feed["image"] = image
|
|
if "im_shape" in names:
|
|
feed["im_shape"] = im_shape
|
|
if "scale_factor" in names:
|
|
feed["scale_factor"] = scale_factor
|
|
# Fill any remaining unknown inputs best-effort
|
|
for inp in inputs:
|
|
if inp.name in feed:
|
|
continue
|
|
shape = inp.shape
|
|
if len(shape) == 4:
|
|
if shape[-1] == 3:
|
|
feed[inp.name] = np.transpose(image, (0, 2, 3, 1))
|
|
else:
|
|
feed[inp.name] = image
|
|
elif len(shape) == 2:
|
|
feed[inp.name] = im_shape
|
|
return feed
|
|
|
|
# Single (or primary) image tensor models
|
|
if not inputs:
|
|
return feed
|
|
name = inputs[0].name
|
|
shape = inputs[0].shape
|
|
if len(shape) == 4 and shape[-1] == 3:
|
|
feed[name] = np.transpose(image, (0, 2, 3, 1))
|
|
else:
|
|
feed[name] = image
|
|
return feed
|
|
|
|
|
|
def _parse_pp_doclayout_boxes(boxes: np.ndarray, geom: PageGeom, min_score: float) -> list[Region]:
|
|
"""Parse PP-DocLayoutV3 rows: [cls, score, x1, y1, x2, y2, read_order] (pixel space)."""
|
|
import numpy as np
|
|
|
|
regions: list[Region] = []
|
|
a = np.asarray(boxes)
|
|
if a.ndim != 2 or a.shape[-1] < 6:
|
|
return regions
|
|
for row in a:
|
|
row = [float(x) for x in row]
|
|
n = len(row)
|
|
# 7-col: cls, score, xyxy, order | 8-col: img_idx, cls, score, xyxy, order
|
|
if n >= 8 and row[2] <= 1.0 and row[1] < 50:
|
|
cls, sc, x1, y1, x2, y2 = int(row[1]), row[2], row[3], row[4], row[5], row[6]
|
|
order = int(row[7]) if n > 7 else len(regions)
|
|
elif n >= 7 and row[1] <= 1.5:
|
|
cls, sc, x1, y1, x2, y2 = int(row[0]), row[1], row[2], row[3], row[4], row[5]
|
|
order = int(row[6]) if n > 6 else len(regions)
|
|
elif n >= 6 and row[4] <= 1.0 and row[5] < 50:
|
|
x1, y1, x2, y2, sc, cls = row[0], row[1], row[2], row[3], row[4], int(row[5])
|
|
order = len(regions)
|
|
else:
|
|
continue
|
|
if sc < min_score:
|
|
continue
|
|
if x2 < x1 or y2 < y1:
|
|
x2, y2 = x1 + max(0.0, x2), y1 + max(0.0, y2)
|
|
if max(x1, y1, x2, y2) <= 1.5:
|
|
x1 *= geom.pixel_w
|
|
x2 *= geom.pixel_w
|
|
y1 *= geom.pixel_h
|
|
y2 *= geom.pixel_h
|
|
label = _map_label(cls)
|
|
bbox = pixel_bbox_to_pdf(float(x1), float(y1), float(x2), float(y2), geom)
|
|
regions.append(Region(label=label, bbox_pdf=bbox, score=float(sc), reading_index=order))
|
|
regions.sort(key=lambda r: (r.reading_index, -(r.bbox_pdf.y + r.bbox_pdf.h), r.bbox_pdf.x))
|
|
for i, r in enumerate(regions):
|
|
r.reading_index = i
|
|
return regions
|
|
|
|
|
|
def _parse_outputs(outs: list, geom: PageGeom, min_score: float) -> list[Region]:
|
|
"""Best-effort parse of common detection outputs (boxes, scores, labels)."""
|
|
import numpy as np
|
|
|
|
# Prefer PP-DocLayout-style (N, 6|7|8) float matrix
|
|
for arr in outs:
|
|
a = np.asarray(arr)
|
|
if a.ndim == 2 and a.shape[-1] in (6, 7, 8) and a.dtype.kind == "f":
|
|
parsed = _parse_pp_doclayout_boxes(a, geom, min_score)
|
|
if parsed:
|
|
return parsed
|
|
|
|
boxes = scores = labels = None
|
|
for arr in outs:
|
|
a = np.asarray(arr)
|
|
if a.ndim >= 2 and a.shape[-1] >= 4 and boxes is None and a.dtype.kind == "f":
|
|
boxes = a.reshape(-1, a.shape[-1])
|
|
elif a.ndim == 1 and a.dtype != object and scores is None and a.size > 0:
|
|
if a.dtype.kind == "f":
|
|
scores = a.reshape(-1)
|
|
else:
|
|
labels = a.reshape(-1)
|
|
elif a.ndim == 2 and a.shape[0] == 1 and a.shape[1] > 4 and a.dtype.kind == "f":
|
|
boxes = a.reshape(-1, a.shape[-1])
|
|
|
|
if boxes is None:
|
|
return []
|
|
|
|
# Reuse PP parser when shape matches; else legacy path
|
|
if boxes.shape[-1] in (6, 7, 8):
|
|
return _parse_pp_doclayout_boxes(boxes, geom, min_score)
|
|
|
|
regions: list[Region] = []
|
|
for i, row in enumerate(boxes):
|
|
row = list(row)
|
|
if len(row) < 4:
|
|
continue
|
|
if len(row) >= 6:
|
|
if row[4] <= 1.0 and row[5] < 20:
|
|
x1, y1, x2, y2, sc, cls = row[0], row[1], row[2], row[3], row[4], int(row[5])
|
|
else:
|
|
sc, cls, x1, y1, x2, y2 = row[0], int(row[1]), row[2], row[3], row[4], row[5]
|
|
else:
|
|
x1, y1, x2, y2 = row[0], row[1], row[2], row[3]
|
|
sc = float(scores[i]) if scores is not None and i < len(scores) else 1.0
|
|
cls = int(labels[i]) if labels is not None and i < len(labels) else 0
|
|
|
|
if sc < min_score:
|
|
continue
|
|
if x2 < x1 or y2 < y1:
|
|
x2, y2 = x1 + max(0.0, x2), y1 + max(0.0, y2)
|
|
if max(x1, y1, x2, y2) <= 1.5:
|
|
x1 *= geom.pixel_w
|
|
x2 *= geom.pixel_w
|
|
y1 *= geom.pixel_h
|
|
y2 *= geom.pixel_h
|
|
label = _map_label(cls)
|
|
bbox = pixel_bbox_to_pdf(float(x1), float(y1), float(x2), float(y2), geom)
|
|
regions.append(Region(label=label, bbox_pdf=bbox, score=float(sc), reading_index=len(regions)))
|
|
regions.sort(key=lambda r: (-(r.bbox_pdf.y + r.bbox_pdf.h), r.bbox_pdf.x))
|
|
for i, r in enumerate(regions):
|
|
r.reading_index = i
|
|
return regions
|
|
|
|
|
|
def detect_regions(
|
|
png_bytes: bytes,
|
|
page_width: float,
|
|
page_height: float,
|
|
*,
|
|
dpi: int | None = None,
|
|
) -> tuple[list[Region], str | None]:
|
|
"""
|
|
Returns (regions, error_or_none).
|
|
Empty regions + error ⇒ caller must fail-open to heuristics.
|
|
"""
|
|
from app.services.convert.layout.page_raster import layout_ml_dpi
|
|
|
|
d = dpi if dpi is not None else layout_ml_dpi()
|
|
geom = PageGeom(page_width=page_width, page_height=page_height, dpi=float(d))
|
|
sess = _get_session()
|
|
if sess is None:
|
|
return [], _load_error or "layout ML session unavailable"
|
|
|
|
try:
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
|
arr = np.asarray(img)
|
|
feed = build_onnx_feed(sess, arr)
|
|
|
|
def _run():
|
|
with _lock:
|
|
return sess.run(None, feed)
|
|
|
|
outs = run_with_timeout(_run, _page_timeout(), label="layout ML infer")
|
|
regions = _parse_outputs(list(outs), geom, _min_score())
|
|
return regions, None
|
|
except Exception as exc:
|
|
return [], str(exc)
|