layout.py

"""PP-DocLayoutV3 the way the OCR SDK runs it.

The picture is resized to 800x800 in BGR and scaled to 0-1. Scores, boxes, and
reading order come from the raw model outputs. Same-class boxes are suppressed,
and a box that sits inside a larger one is dropped.
"""

from __future__ import annotations

from typing import List

import numpy as np
import torch

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",
]
# Keep the large region when one of these classes contains another box.
_KEEP_LARGE = {3, 5, 6, 15, 17}


def _sigmoid(values: np.ndarray) -> np.ndarray:
    values = np.clip(values, -50, 50)
    return 1.0 / (1.0 + np.exp(-values))


def _reading_order(order_logits: np.ndarray) -> np.ndarray:
    scores = _sigmoid(order_logits.astype(np.float32))
    upper = np.triu(scores, 1)
    lower = np.tril(np.float32(1.0) - scores.T, -1)
    votes = upper.sum(axis=0) + lower.sum(axis=0)
    sequence = np.zeros(order_logits.shape[0], dtype=np.int64)
    sequence[np.argsort(votes)] = np.arange(order_logits.shape[0])
    return sequence


def _boxes_7col(logits, pred_boxes, order_logits, width: int, height: int) -> np.ndarray:
    num_classes = logits.shape[1]
    flat = _sigmoid(logits).reshape(-1)
    top = min(300, flat.size)
    chosen = np.argpartition(-flat, top - 1)[:top]
    chosen = chosen[np.argsort(-flat[chosen])]
    labels = (chosen % num_classes).astype(np.float32)
    query = chosen // num_classes
    centers, sizes = pred_boxes[:, :2], pred_boxes[:, 2:]
    corners = np.concatenate([centers - 0.5 * sizes, centers + 0.5 * sizes], axis=-1)
    scaled = corners * np.array([width, height, width, height], dtype=np.float32)
    order = _reading_order(order_logits)[query].astype(np.float32)
    picked = scaled[query]
    return np.column_stack([labels, flat[chosen], picked, order])


def _iou(left: np.ndarray, right: np.ndarray) -> float:
    x1 = max(left[2], right[2])
    y1 = max(left[3], right[3])
    x2 = min(left[4], right[4])
    y2 = min(left[5], right[5])
    intersection_w = x2 - x1 + 1
    intersection_h = y2 - y1 + 1
    if intersection_w <= 0 or intersection_h <= 0:
        return 0.0
    intersection = intersection_w * intersection_h
    area_left = (left[4] - left[2] + 1) * (left[5] - left[3] + 1)
    area_right = (right[4] - right[2] + 1) * (right[5] - right[3] + 1)
    return float(intersection / (area_left + area_right - intersection))


def _nms(boxes: np.ndarray) -> List[int]:
    remaining = list(np.argsort(-boxes[:, 1]))
    kept: List[int] = []
    while remaining:
        current = remaining.pop(0)
        kept.append(current)
        current_box = boxes[current]
        still = []
        for other in remaining:
            other_box = boxes[other]
            limit = 0.6 if current_box[0] == other_box[0] else 0.98
            if _iou(current_box, other_box) < limit:
                still.append(other)
        remaining = still
    return kept


def _contained(inner: np.ndarray, outer: np.ndarray) -> bool:
    x1 = max(inner[2], outer[2])
    y1 = max(inner[3], outer[3])
    x2 = min(inner[4], outer[4])
    y2 = min(inner[5], outer[5])
    area = max(0.0, x2 - x1) * max(0.0, y2 - y1)
    inner_area = (inner[4] - inner[2]) * (inner[5] - inner[3])
    return bool(inner_area > 0 and area / inner_area >= 0.9)


def _drop_contained(boxes: np.ndarray) -> np.ndarray:
    contained = np.zeros(len(boxes), dtype=bool)
    for class_id in _KEEP_LARGE:
        for index, box in enumerate(boxes):
            for other_index, other in enumerate(boxes):
                if other_index == index or other[0] != class_id:
                    continue
                if _contained(box, other):
                    contained[index] = True
                    break
    return boxes[~contained]


def _drop_overlapping(regions: List[dict]) -> List[dict]:
    regions = [region for region in regions if region["label"] != "reference"]
    if not regions:
        return regions
    coords = np.array([region["bbox"] for region in regions], dtype=np.float64)
    widths = coords[:, 2] - coords[:, 0]
    heights = coords[:, 3] - coords[:, 1]
    areas = np.maximum(widths, 0) * np.maximum(heights, 0)
    dropped = set()
    for index, region in enumerate(regions):
        if widths[index] < 6 or heights[index] < 6:
            dropped.add(index)
    for left in range(len(regions)):
        for right in range(left + 1, len(regions)):
            if left in dropped or right in dropped:
                continue
            x1 = max(coords[left, 0], coords[right, 0])
            y1 = max(coords[left, 1], coords[right, 1])
            x2 = min(coords[left, 2], coords[right, 2])
            y2 = min(coords[left, 3], coords[right, 3])
            intersection = max(0.0, x2 - x1) * max(0.0, y2 - y1)
            smaller = min(areas[left], areas[right])
            ratio = intersection / smaller if smaller > 0 else 0.0
            labels = {regions[left]["label"], regions[right]["label"]}
            if "inline_formula" in labels and ratio > 0.5:
                if regions[left]["label"] == "inline_formula":
                    dropped.add(left)
                if regions[right]["label"] == "inline_formula":
                    dropped.add(right)
                continue
            if ratio > 0.7:
                dropped.add(right if areas[left] >= areas[right] else left)
    return [region for index, region in enumerate(regions) if index not in dropped]


def detect(model, image) -> List[tuple]:
    import cv2

    rgb = np.asarray(image.convert("RGB"))
    height, width = rgb.shape[:2]
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    resized = cv2.resize(bgr, (800, 800), interpolation=cv2.INTER_CUBIC)
    channels = [(channel.astype(np.float32) / 255.0) for channel in cv2.split(resized)]
    array = cv2.merge(channels).transpose(2, 0, 1)
    pixel_values = torch.from_numpy(array).unsqueeze(0).to("cuda")
    with torch.inference_mode():
        outputs = model(pixel_values=pixel_values)
    logits = outputs.logits.detach().float().cpu().numpy()[0]
    pred_boxes = outputs.pred_boxes.detach().float().cpu().numpy()[0]
    order_logits = outputs.order_logits.detach().float().cpu().numpy()[0]
    boxes = _boxes_7col(logits, pred_boxes, order_logits, width, height)
    boxes[:, 2:6] = np.round(boxes[:, 2:6])
    boxes = boxes[(boxes[:, 1] > 0.3) & (boxes[:, 0] > -1)]
    if len(boxes) == 0:
        return []
    boxes = boxes[_nms(boxes[:, :6])]
    if len(boxes) > 1:
        image_id = LABELS.index("image")
        area_limit = 0.82 if width > height else 0.93
        page_area = width * height
        kept = []
        for box in boxes:
            if box[0] != image_id:
                kept.append(box)
                continue
            x1, y1 = max(0, int(box[2])), max(0, int(box[3]))
            x2, y2 = min(width, int(box[4])), min(height, int(box[5]))
            if (x2 - x1) * (y2 - y1) <= area_limit * page_area:
                kept.append(box)
        boxes = np.array(kept, dtype=boxes.dtype) if kept else boxes[:0]
    if len(boxes) >= 2:
        boxes = _drop_contained(boxes)
    if len(boxes) == 0:
        return []
    boxes = boxes[np.argsort(boxes[:, 6])]
    regions = []
    for box in boxes:
        x1 = int(max(0, box[2]))
        y1 = int(max(0, box[3]))
        x2 = int(min(width, box[4]))
        y2 = int(min(height, box[5]))
        if x2 <= x1 or y2 <= y1:
            continue
        class_id = int(box[0])
        regions.append({"label": LABELS[class_id], "bbox": [x1, y1, x2, y2]})
    return [(region["label"], region["bbox"]) for region in _drop_overlapping(regions)]