src/perception/ocr_extractor.py

"""OCR 文本提取：PP-OCRv5 ONNX 推理 + PP-DocLayout-S sim-cut ONNX 版面分析"""
from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
import math
from pathlib import Path
from threading import RLock
from typing import List, Optional, Sequence, Tuple

import cv2
import numpy as np
from PIL import Image, ImageEnhance

from ..preprocess.scene_classifier import Scene, SceneResult

PERCEPTION_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = PERCEPTION_DIR.parents[1]

AXMODEL_DIR = PROJECT_ROOT / "axmodel"
DET_MODEL_PATH = AXMODEL_DIR / "ppocrv5" / "det_npu1.axmodel"
CLS_MODEL_PATH = AXMODEL_DIR / "ppocrv5" / "cls_npu1.axmodel"
REC_MODEL_PATH = AXMODEL_DIR / "ppocrv5" / "rec_npu1.axmodel"
DICT_PATH = PERCEPTION_DIR / "dict" / "ppocrv5_dict.txt"
LAYOUT_MODEL_PATH = AXMODEL_DIR / "ppstructurev3" / "ppstructure_npu1.axmodel"


DET_INPUT_SIZE = 960
CLS_INPUT_SIZE = (80, 160)  # h, w
REC_INPUT_SIZE = (48, 320)  # h, w
LAYOUT_INPUT_SIZE = (480, 480)  # h, w
DET_MEAN = np.array([123.675, 116.28, 103.53], dtype=np.float32)
DET_STD = np.array([58.395, 57.12, 57.375], dtype=np.float32)
CLS_MEAN = np.array([127.5, 127.5, 127.5], dtype=np.float32)
CLS_STD = np.array([127.5, 127.5, 127.5], dtype=np.float32)
REC_MEAN = np.array([127.5, 127.5, 127.5], dtype=np.float32)
REC_STD = np.array([127.5, 127.5, 127.5], dtype=np.float32)
LAYOUT_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
LAYOUT_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

DET_BIN_THRESHOLD = 0.30
DET_MIN_COMPONENT_AREA = 16
DET_SCORE_THRESHOLD = 0.35
CLS_ROTATE_THRESHOLD = 0.90
LAYOUT_SCORE_THRESHOLD = 0.35
LAYOUT_DEDUP_IOU = 0.85

LAYOUT_LABELS = [
    "paragraph_title",
    "image",
    "text",
    "number",
    "abstract",
    "content",
    "figure_title",
    "formula",
    "table",
    "table_title",
    "reference",
    "doc_title",
    "footnote",
    "header",
    "algorithm",
    "footer",
    "seal",
    "chart_title",
    "chart",
    "formula_number",
    "header_image",
    "footer_image",
    "aside_text",
]

LAYOUT_TYPE_MAP = {
    "paragraph_title": "title",
    "doc_title": "title",
    "text": "text",
    "content": "text",
    "reference": "text",
    "abstract": "text",
    "number": "text",
    "aside_text": "text",
    "figure_title": "figure_caption",
    "table_title": "figure_caption",
    "chart_title": "figure_caption",
    "formula_number": "figure_caption",
    "image": "figure",
    "chart": "figure",
    "header_image": "figure",
    "footer_image": "figure",
    "table": "table",
    "formula": "figure",
    "header": "unknown",
    "footer": "unknown",
    "seal": "unknown",
    "algorithm": "text",
    "footnote": "text",
}

LAYOUT_ENABLED_LABELS = {
    "paragraph_title",
    "image",
    "text",
    "number",
    "abstract",
    "content",
    "figure_title",
    "formula",
    "table",
    "table_title",
    "reference",
    "doc_title",
    "footnote",
    "algorithm",
    "chart_title",
    "chart",
}

_ort = None
_runtime_lock = RLock()
_det_sess = None
_cls_sess = None
_rec_sess = None
_layout_sess = None
_char_list = None
_pyclipper = None


def _get_ort():
    global _ort
    if _ort is None:
        try:
            # import onnxruntime as ort  # type: ignore
            import axengine as ort  # type: ignore
        except ImportError as exc:  # pragma: no cover - environment issue
            raise RuntimeError(
                "onnxruntime is required for PP-OCRv5 ONNX inference. "
                "Please install it in the runtime environment."
            ) from exc
        _ort = ort
    return _ort


def _get_session(path: Path):
    with _runtime_lock:
        ort = _get_ort()
        return ort.InferenceSession(str(path), providers=["AxEngineExecutionProvider"])


def _get_det_sess():
    global _det_sess
    if _det_sess is None:
        _det_sess = _get_session(DET_MODEL_PATH)
    return _det_sess


def _get_cls_sess():
    global _cls_sess
    if _cls_sess is None:
        _cls_sess = _get_session(CLS_MODEL_PATH)
    return _cls_sess


def _get_rec_sess():
    global _rec_sess
    if _rec_sess is None:
        _rec_sess = _get_session(REC_MODEL_PATH)
    return _rec_sess


def _get_layout_sess():
    global _layout_sess
    if _layout_sess is None:
        _layout_sess = _get_session(LAYOUT_MODEL_PATH)
    return _layout_sess


def _get_char_list() -> List[str]:
    global _char_list
    if _char_list is None:
        chars = DICT_PATH.read_text(encoding="utf-8").splitlines()
        # PP-OCRv5 rec output class count = blank + dict + ascii space.
        _char_list = ["blank"] + chars + [" "]
    return _char_list


def _get_pyclipper():
    global _pyclipper
    if _pyclipper is None:
        try:
            import pyclipper  # type: ignore
        except ImportError:
            _pyclipper = False
        else:
            _pyclipper = pyclipper
    return _pyclipper


@dataclass
class TextBlock:
    text: str
    score: float
    bbox: List[int]  # [x1, y1, x2, y2]
    region_type: str = "unknown"  # text / title / figure_caption / unknown


class OCRTextState(str, Enum):
    OK = "ok"
    LOW_QUALITY = "low_quality"
    MISSING_EXPECTED_TEXT = "missing_expected_text"
    MISSING_UNCERTAIN = "missing_uncertain"


@dataclass
class OCRResult:
    blocks: List[TextBlock] = field(default_factory=list)
    avg_score: float = 0.0
    low_confidence: bool = False
    text_state: OCRTextState = OCRTextState.OK
    figure_images: List[Image.Image] = field(default_factory=list)
    figure_bboxes: List[List[int]] = field(default_factory=list)


def _resize_no_pad(arr_rgb: np.ndarray, target_h: int, target_w: int, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
    resized = np.array(Image.fromarray(arr_rgb).resize((target_w, target_h), Image.BICUBIC)).astype(np.float32).transpose(2, 0, 1)[np.newaxis]
    # resized = resized / 255.0
    # resized = ((resized - mean) / std).transpose(2, 0, 1)[np.newaxis]
    return resized


def _get_mini_boxes(contour: np.ndarray) -> Tuple[List[np.ndarray], float]:
    bounding_box = cv2.minAreaRect(contour)
    points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])

    if points[1][1] > points[0][1]:
        index_1, index_4 = 0, 1
    else:
        index_1, index_4 = 1, 0

    if points[3][1] > points[2][1]:
        index_2, index_3 = 2, 3
    else:
        index_2, index_3 = 3, 2

    box = [points[index_1], points[index_2], points[index_3], points[index_4]]
    return box, min(bounding_box[1])


def _box_score_fast(bitmap: np.ndarray, box: np.ndarray) -> float:
    h, w = bitmap.shape[:2]
    scored_box = box.copy()
    xmin = np.clip(np.floor(scored_box[:, 0].min()).astype("int32"), 0, w - 1)
    xmax = np.clip(np.ceil(scored_box[:, 0].max()).astype("int32"), 0, w - 1)
    ymin = np.clip(np.floor(scored_box[:, 1].min()).astype("int32"), 0, h - 1)
    ymax = np.clip(np.ceil(scored_box[:, 1].max()).astype("int32"), 0, h - 1)

    mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
    scored_box[:, 0] = scored_box[:, 0] - xmin
    scored_box[:, 1] = scored_box[:, 1] - ymin
    cv2.fillPoly(mask, scored_box.reshape(1, -1, 2).astype("int32"), 1)
    return float(cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0])


def _unclip(box: np.ndarray, unclip_ratio: float) -> np.ndarray:
    area = float(abs(cv2.contourArea(box.astype(np.float32))))
    length = float(cv2.arcLength(box.astype(np.float32), True))
    if area <= 1e-6 or length <= 1e-6:
        return np.array([box], dtype=np.float32)

    pyclipper = _get_pyclipper()
    if not pyclipper:
        return np.array([box], dtype=np.float32)

    distance = area * unclip_ratio / length
    offset = pyclipper.PyclipperOffset()
    offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
    expanded = offset.Execute(distance)
    if not expanded:
        return np.array([box], dtype=np.float32)
    return np.array(expanded)


def _boxes_from_bitmap(pred: np.ndarray, bitmap: np.ndarray, dest_width: int, dest_height: int) -> Tuple[np.ndarray, List[float]]:
    box_thresh = 0.6
    max_candidates = 1000
    unclip_ratio = 1.5
    min_size = 3

    height, width = bitmap.shape
    outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    if len(outs) == 3:
        contours = outs[1]
    else:
        contours = outs[0]

    num_contours = min(len(contours), max_candidates)
    boxes: List[np.ndarray] = []
    scores: List[float] = []
    for index in range(num_contours):
        contour = contours[index]
        points, sside = _get_mini_boxes(contour)
        if sside < min_size:
            continue
        points = np.array(points)
        score = _box_score_fast(pred, points.reshape(-1, 2))
        if score < box_thresh:
            continue

        box = _unclip(points, unclip_ratio).reshape(-1, 1, 2)
        box, sside = _get_mini_boxes(box)
        if sside < min_size + 2:
            continue
        box = np.array(box)
        box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
        box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
        boxes.append(box.astype("int32"))
        scores.append(float(score))
    if not boxes:
        return np.zeros((0, 4, 2), dtype="int32"), scores
    return np.array(boxes, dtype="int32"), scores


def _order_points_clockwise(points: np.ndarray) -> np.ndarray:
    rect = np.zeros((4, 2), dtype="float32")
    point_sums = points.sum(axis=1)
    rect[0] = points[np.argmin(point_sums)]
    rect[2] = points[np.argmax(point_sums)]
    remaining = np.delete(points, (np.argmin(point_sums), np.argmax(point_sums)), axis=0)
    diffs = np.diff(np.array(remaining), axis=1)
    rect[1] = remaining[np.argmin(diffs)]
    rect[3] = remaining[np.argmax(diffs)]
    return rect


def _clip_det_res(points: np.ndarray, img_height: int, img_width: int) -> np.ndarray:
    for index in range(points.shape[0]):
        points[index, 0] = int(min(max(points[index, 0], 0), img_width - 1))
        points[index, 1] = int(min(max(points[index, 1], 0), img_height - 1))
    return points


def _filter_tag_det_res(dt_boxes: np.ndarray, image_shape: Sequence[int]) -> np.ndarray:
    img_height, img_width = image_shape[0:2]
    filtered_boxes = []
    for box in dt_boxes:
        current_box = np.array(box) if isinstance(box, list) else box.copy()
        current_box = _order_points_clockwise(current_box)
        current_box = _clip_det_res(current_box, img_height, img_width)
        rect_width = int(np.linalg.norm(current_box[0] - current_box[1]))
        rect_height = int(np.linalg.norm(current_box[0] - current_box[3]))
        if rect_width <= 3 or rect_height <= 3:
            continue
        filtered_boxes.append(current_box)
    if not filtered_boxes:
        return np.zeros((0, 4, 2), dtype=np.float32)
    return np.array(filtered_boxes)


def _sorted_boxes(dt_boxes: np.ndarray) -> List[np.ndarray]:
    num_boxes = dt_boxes.shape[0]
    ordered = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
    boxes = list(ordered)
    for i in range(num_boxes - 1):
        for j in range(i, -1, -1):
            if abs(boxes[j + 1][0][1] - boxes[j][0][1]) < 10 and boxes[j + 1][0][0] < boxes[j][0][0]:
                boxes[j], boxes[j + 1] = boxes[j + 1], boxes[j]
            else:
                break
    return boxes


def _det_postprocess(outs_dict: dict, shape_list: np.ndarray) -> List[dict]:
    pred = outs_dict["maps"]
    pred = pred[:, 0, :, :]
    segmentation = pred > DET_BIN_THRESHOLD
    boxes_batch = []
    for batch_index in range(pred.shape[0]):
        src_h, src_w = int(shape_list[batch_index][0]), int(shape_list[batch_index][1])
        mask = segmentation[batch_index]
        boxes, _ = _boxes_from_bitmap(pred[batch_index], mask, src_w, src_h)
        boxes_batch.append({"points": boxes})
    return boxes_batch


def _detect_boxes(arr_bgr: np.ndarray) -> List[np.ndarray]:
    orig_h, orig_w = arr_bgr.shape[:2]
    det_input = cv2.resize(arr_bgr, (DET_INPUT_SIZE, DET_INPUT_SIZE)).astype(np.float32)
    det_input = det_input.transpose(2, 0, 1)[np.newaxis]
    shape_list = np.array(
        [[orig_h, orig_w, DET_INPUT_SIZE / max(orig_h, 1), DET_INPUT_SIZE / max(orig_w, 1)]],
        dtype=np.float32,
    )
    sess = _get_det_sess()
    det_out = sess.run(None, {sess.get_inputs()[0].name: det_input})
    post_result = _det_postprocess({"maps": det_out[0]}, shape_list)
    dt_boxes = _filter_tag_det_res(post_result[0]["points"], arr_bgr.shape)
    if dt_boxes.size == 0:
        return []
    return _sorted_boxes(dt_boxes)


def _get_rotate_crop_image(img: np.ndarray, points: np.ndarray) -> np.ndarray:
    if len(points) != 4:
        raise ValueError("shape of points must be 4*2")
    img_crop_width = int(max(np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])))
    img_crop_height = int(max(np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])))
    pts_std = np.float32(
        [
            [0, 0],
            [img_crop_width, 0],
            [img_crop_width, img_crop_height],
            [0, img_crop_height],
        ]
    )
    transform = cv2.getPerspectiveTransform(points.astype(np.float32), pts_std)
    dst_img = cv2.warpPerspective(
        img,
        transform,
        (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE,
        flags=cv2.INTER_CUBIC,
    )
    dst_img_height, dst_img_width = dst_img.shape[:2]
    if dst_img_width == 0 or dst_img_height == 0:
        return dst_img
    if dst_img_height * 1.0 / dst_img_width >= 1.5:
        dst_img = np.rot90(dst_img)
    return dst_img


def _resize_norm_img(img: np.ndarray, shape: Tuple[int, int, int]) -> np.ndarray:
    img_h, img_w = img.shape[:2]
    img_c, target_h, target_w = shape
    ratio = img_w / float(max(img_h, 1))
    if math.ceil(target_h * ratio) > target_w:
        resized_w = target_w
    else:
        resized_w = int(math.ceil(target_h * ratio))
    resized_image = cv2.resize(img, (resized_w, target_h)).astype(np.float32)
    resized_image = resized_image.transpose((2, 0, 1))
    padding_im = np.zeros((img_c, target_h, target_w), dtype=np.float32)
    padding_im[:, :, 0:resized_w] = resized_image
    return padding_im


def _classify_orientation(crop_bgr: np.ndarray) -> np.ndarray:
    if crop_bgr.size == 0:
        return crop_bgr
    cls_input = _resize_norm_img(crop_bgr, (3, CLS_INPUT_SIZE[0], CLS_INPUT_SIZE[1]))[np.newaxis].copy()
    sess = _get_cls_sess()
    probs = sess.run(None, {sess.get_inputs()[0].name: cls_input})[0]
    pred_idx = int(np.argmax(probs[0]))
    pred_score = float(probs[0, pred_idx])
    if pred_idx == 1 and pred_score > CLS_ROTATE_THRESHOLD:
        return cv2.rotate(crop_bgr, cv2.ROTATE_180)
    return crop_bgr


def _preprocess_rec(crop_bgr: np.ndarray) -> np.ndarray:
    if crop_bgr.size == 0:
        return np.zeros((1, 3, REC_INPUT_SIZE[0], REC_INPUT_SIZE[1]), dtype=np.float32)
    return _resize_norm_img(crop_bgr, (3, REC_INPUT_SIZE[0], REC_INPUT_SIZE[1]))[np.newaxis].copy()


def _decode_text_indices(
    chars: List[str],
    text_index: np.ndarray,
    text_prob: np.ndarray,
    is_remove_duplicate: bool = False,
) -> List[Tuple[str, float]]:
    result_list: List[Tuple[str, float]] = []
    ignored_tokens = [0]
    batch_size = len(text_index)
    for batch_idx in range(batch_size):
        selection = np.ones(len(text_index[batch_idx]), dtype=bool)
        if is_remove_duplicate:
            selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
        for ignored_token in ignored_tokens:
            selection &= text_index[batch_idx] != ignored_token
        char_list = [chars[text_id] for text_id in text_index[batch_idx][selection] if text_id < len(chars)]
        conf_list = text_prob[batch_idx][selection]
        if len(conf_list) == 0:
            conf_list = np.array([0.0], dtype=np.float32)
        result_list.append(("".join(char_list), float(np.mean(conf_list).tolist())))
    return result_list


def _decode_rec(logits: np.ndarray) -> Tuple[str, float]:
    chars = _get_char_list()
    preds = logits[-1] if isinstance(logits, (tuple, list)) else logits
    preds = preds.astype(np.float32)
    preds_idx = preds.argmax(axis=2)
    preds_prob = preds.max(axis=2)
    return _decode_text_indices(chars, preds_idx, preds_prob, is_remove_duplicate=True)[0]


def _recognize_crop(crop_bgr: np.ndarray) -> Tuple[str, float]:
    oriented = _classify_orientation(crop_bgr)
    rec_input = _preprocess_rec(oriented)
    sess = _get_rec_sess()
    logits = sess.run(None, {sess.get_inputs()[0].name: rec_input})[0]
    return _decode_rec(logits)


def _enhance_crop(crop_bgr: np.ndarray) -> np.ndarray:
    img = Image.fromarray(crop_bgr[:, :, ::-1])
    img = ImageEnhance.Contrast(img).enhance(1.5)
    return np.array(img)[:, :, ::-1]


def _quad_to_bbox(points: np.ndarray, width: int, height: int) -> List[int]:
    min_x = max(0, int(np.floor(np.min(points[:, 0]))))
    min_y = max(0, int(np.floor(np.min(points[:, 1]))))
    max_x = min(width, int(np.ceil(np.max(points[:, 0]))))
    max_y = min(height, int(np.ceil(np.max(points[:, 1]))))
    return [min_x, min_y, max_x, max_y]


def _iou(a: List[int], b: List[int]) -> float:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    if inter == 0:
        return 0.0
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    return inter / min(area_a, area_b)


def _clip_bbox(bbox: Sequence[float], width: int, height: int) -> List[int]:
    x1, y1, x2, y2 = bbox
    clipped = [
        int(round(max(0.0, min(float(width), x1)))),
        int(round(max(0.0, min(float(height), y1)))),
        int(round(max(0.0, min(float(width), x2)))),
        int(round(max(0.0, min(float(height), y2)))),
    ]
    return clipped


def _dedupe_layout_regions(regions: List[dict]) -> List[dict]:
    kept: List[dict] = []
    for region in sorted(regions, key=lambda item: item["score"], reverse=True):
        duplicate = False
        for existing in kept:
            same_type = existing.get("type") == region.get("type")
            same_label = existing.get("raw_type") == region.get("raw_type")
            if not same_type:
                continue
            if region.get("type") != "figure" and not same_label:
                continue
            if _iou(existing["bbox"], region["bbox"]) >= LAYOUT_DEDUP_IOU:
                duplicate = True
                break
        if not duplicate:
            kept.append(region)
    kept.sort(key=lambda item: (item["bbox"][1], item["bbox"][0]))
    return kept


def _run_ocr_full(arr_rgb: np.ndarray, threshold: float) -> List[TextBlock]:
    blocks = []
    arr_bgr = arr_rgb[:, :, ::-1]
    height, width = arr_rgb.shape[:2]
    for dt_box in _detect_boxes(arr_bgr):
        crop = _get_rotate_crop_image(arr_bgr, np.asarray(dt_box, dtype=np.float32))
        if crop.size == 0:
            continue
        text, score = _recognize_crop(crop)
        if score < threshold:
            retry_text, retry_score = _recognize_crop(_enhance_crop(crop))
            if retry_score > score:
                text, score = retry_text, retry_score
        bbox = _quad_to_bbox(np.asarray(dt_box, dtype=np.float32), width, height)
        blocks.append(TextBlock(text=text, score=score, bbox=bbox))
    return blocks


def _assign_region(block: TextBlock, regions: list) -> tuple:
    """返回 (region_type, region_index)"""
    best_type, best_iou, best_idx = "unknown", 0.0, len(regions)
    for i, region in enumerate(regions):
        bbox = region.get("bbox")
        if bbox is None:
            continue
        score = _iou(block.bbox, bbox)
        if score > best_iou:
            best_iou = score
            best_type = region.get("type", "unknown")
            best_idx = i
            if best_iou >= 0.9:
                return best_type, best_idx
    return (best_type if best_iou > 0.1 else "unknown"), best_idx


def _run_structure(arr_rgb: np.ndarray):
    sess = _get_layout_sess()
    h, w = arr_rgb.shape[:2]
    image_input = _resize_no_pad(arr_rgb, LAYOUT_INPUT_SIZE[0], LAYOUT_INPUT_SIZE[1], LAYOUT_MEAN, LAYOUT_STD)

    def _append_region(regions: List[dict], cls_id: int, score: float, bbox_values: Sequence[float], scale_to_image: bool):
        if score < LAYOUT_SCORE_THRESHOLD:
            return
        raw_label = LAYOUT_LABELS[cls_id] if 0 <= cls_id < len(LAYOUT_LABELS) else str(cls_id)
        if raw_label not in LAYOUT_ENABLED_LABELS:
            return
        region_type = LAYOUT_TYPE_MAP.get(raw_label, "unknown")
        x1, y1, x2, y2 = map(float, bbox_values)
        if x2 < x1:
            x1, x2 = x2, x1
        if y2 < y1:
            y1, y2 = y2, y1
        if scale_to_image:
            x_scale = w / float(LAYOUT_INPUT_SIZE[1])
            y_scale = h / float(LAYOUT_INPUT_SIZE[0])
            bbox = _clip_bbox([x1 * x_scale, y1 * y_scale, x2 * x_scale, y2 * y_scale], w, h)
        else:
            bbox = _clip_bbox([x1, y1, x2, y2], w, h)
        if bbox[2] - bbox[0] < 4 or bbox[3] - bbox[1] < 4:
            return
        region = {"type": region_type, "raw_type": raw_label, "score": score, "bbox": bbox}
        if region_type == "figure":
            x1c, y1c, x2c, y2c = bbox
            crop = arr_rgb[max(0, y1c):max(0, y2c), max(0, x1c):max(0, x2c)]
            if crop.size > 0:
                region["img"] = crop
        regions.append(region)

    regions = []
    inputs = sess.get_inputs()
    if len(inputs) >= 2:
        # 兼容原版 PP-DocLayout-S：输出为 [cls_id, score, x1, y1, x2, y2] + num
        scale_factor = np.array([[LAYOUT_INPUT_SIZE[0] / max(h, 1), LAYOUT_INPUT_SIZE[1] / max(w, 1)]], dtype=np.float32)
        outputs = sess.run(None, {inputs[0].name: image_input, inputs[1].name: scale_factor})
        dets, num = outputs
        num = int(num[0])
        for row in dets[:num]:
            _append_region(regions, int(row[0]), float(row[1]), row[2:].tolist(), scale_to_image=False)
    else:
        # 兼容 sim-cut 导出：输出为 boxes[N,4] + cls[C,N]（均在 480x480 输入坐标系）
        outputs = sess.run(None, {inputs[0].name: image_input})
        if len(outputs) != 2:
            raise RuntimeError(f"Unexpected layout outputs: {len(outputs)}")
        out_names = [o.name for o in sess.get_outputs()]
        out_map = {name: value for name, value in zip(out_names, outputs)}

        if "Mul.109" in out_map and "Concat.9" in out_map:
            boxes = out_map["Mul.109"][0]
            cls_blob = out_map["Concat.9"]
        else:
            first, second = outputs
            if first.ndim == 3 and first.shape[-1] == 4:
                boxes, cls_blob = first[0], second
            elif second.ndim == 3 and second.shape[-1] == 4:
                boxes, cls_blob = second[0], first
            else:
                raise RuntimeError(f"Unexpected sim-cut output shapes: {first.shape}, {second.shape}")
        if cls_blob.ndim != 3 or cls_blob.shape[0] != 1:
            raise RuntimeError(f"Unexpected cls output shape: {cls_blob.shape}")
        if cls_blob.shape[1] == len(LAYOUT_LABELS) and cls_blob.shape[2] == boxes.shape[0]:
            cls_scores = cls_blob[0].transpose(1, 0)  # [N, C]
        elif cls_blob.shape[2] == len(LAYOUT_LABELS) and cls_blob.shape[1] == boxes.shape[0]:
            cls_scores = cls_blob[0]  # [N, C]
        else:
            raise RuntimeError(f"Unexpected cls/box alignment: cls={cls_blob.shape}, boxes={boxes.shape}")

        best_cls = np.argmax(cls_scores, axis=1)
        best_score = np.max(cls_scores, axis=1)
        for i in range(boxes.shape[0]):
            _append_region(
                regions,
                int(best_cls[i]),
                float(best_score[i]),
                boxes[i].tolist(),
                scale_to_image=True,
            )
    return _dedupe_layout_regions(regions)


def extract(img: Image.Image, scene_result: SceneResult, debug_prefix: Optional[str] = None) -> OCRResult:
    arr_rgb = np.array(img.convert("RGB"))
    threshold = scene_result.ocr_threshold

    if scene_result.scene == Scene.SCREENSHOT:
        blocks = _run_ocr_full(arr_rgb, threshold)
        figure_images, figure_bboxes = [], []
    else:
        structure_result = _run_structure(arr_rgb)
        blocks = _run_ocr_full(arr_rgb, threshold)

        figure_images, figure_bboxes = [], []
        for region in structure_result:
            if region.get("type") == "figure":
                sub = region.get("img")
                bbox = region.get("bbox")
                if sub is not None and bbox is not None:
                    figure_images.append(Image.fromarray(sub))
                    figure_bboxes.append(bbox)

        region_groups = {}
        unknown_blocks = []
        for block in blocks:
            rtype, ridx = _assign_region(block, structure_result)
            if rtype != "unknown":
                if ridx not in region_groups:
                    region_groups[ridx] = (rtype, [])
                region_groups[ridx][1].append(block)
            else:
                unknown_blocks.append(block)

        region_blocks = []
        for ridx, (rtype, grp) in sorted(region_groups.items()):
            grp.sort(key=lambda b: (int((b.bbox[1] + b.bbox[3]) / 40), b.bbox[0]))
            text = " ".join(b.text for b in grp if b.text)
            score = sum(b.score for b in grp) / len(grp)
            bbox = [
                min(b.bbox[0] for b in grp),
                min(b.bbox[1] for b in grp),
                max(b.bbox[2] for b in grp),
                max(b.bbox[3] for b in grp),
            ]
            region_blocks.append((ridx, TextBlock(text=text, score=score, bbox=bbox, region_type=rtype)))

        def _row_x(b: TextBlock):
            return (int((b.bbox[1] + b.bbox[3]) / 40), b.bbox[0])

        region_blocks.sort(key=lambda t: (t[0], *_row_x(t[1])))
        sorted_region = [b for _, b in region_blocks]
        unknown_blocks.sort(key=_row_x)

        import bisect

        rx_keys = [_row_x(b) for b in sorted_region]
        for ub in unknown_blocks:
            pos = bisect.bisect_left(rx_keys, _row_x(ub))
            sorted_region.insert(pos, ub)
            rx_keys.insert(pos, _row_x(ub))
        blocks = sorted_region

    blocks = [b for b in blocks if len(b.text.strip()) >= 2]
    if not blocks:
        text_state = (
            OCRTextState.MISSING_EXPECTED_TEXT
            if scene_result.scene in {Scene.SCREENSHOT, Scene.DOCUMENT}
            else OCRTextState.MISSING_UNCERTAIN
        )
        return OCRResult(
            blocks=[],
            avg_score=0.0,
            low_confidence=False,
            text_state=text_state,
            figure_images=figure_images,
            figure_bboxes=figure_bboxes,
        )

    avg_score = sum(b.score for b in blocks) / len(blocks)
    low_confidence = avg_score < 0.65 and len(blocks) < 3
    text_state = OCRTextState.LOW_QUALITY if low_confidence else OCRTextState.OK
    return OCRResult(
        blocks=blocks,
        avg_score=avg_score,
        low_confidence=low_confidence,
        text_state=text_state,
        figure_images=figure_images,
        figure_bboxes=figure_bboxes,
    )