Upload ONNX export

pp_doclayout_v3_onnx.py

@@ -0,0 +1,382 @@
+"""PP-DocLayoutV3 inference on ONNX Runtime — no torch, no transformers.
+
+Pre- and post-processing are ported from
+transformers/models/pp_doclayout_v3/image_processing_pp_doclayout_v3.py
+so results match the PyTorch pipeline (boxes, labels, reading order, polygons).
+
+    from pp_doclayout_v3_onnx import PPDocLayoutV3ONNX
+
+    det = PPDocLayoutV3ONNX("pp_doclayoutv3.onnx", device="cuda")
+    for r in det.predict("page.jpg"):
+        print(r["order"], r["label"], r["score"], r["box"])
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+INPUT_SIZE = 800          # image processor resizes to a fixed 800x800 square
+MASK_STRIDE = 4           # mask head resolution is 800/4 = 200
+RESCALE_FACTOR = 1.0 / 255.0
+# image_mean = [0,0,0], image_std = [1,1,1] -> normalisation is a no-op
+
+ID2LABEL = {
+    0: "abstract", 1: "algorithm", 2: "aside_text", 3: "chart", 4: "content",
+    5: "formula", 6: "doc_title", 7: "figure_title", 8: "footer", 9: "footer",
+    10: "footnote", 11: "formula_number", 12: "header", 13: "header", 14: "image",
+    15: "formula", 16: "number", 17: "paragraph_title", 18: "reference",
+    19: "reference_content", 20: "seal", 21: "table", 22: "text", 23: "text",
+    24: "vision_footnote",
+}
+
+
+@dataclass
+class Timings:
+    preprocess: float = 0.0
+    inference: float = 0.0
+    postprocess: float = 0.0
+
+    @property
+    def total(self) -> float:
+        return self.preprocess + self.inference + self.postprocess
+
+    def __str__(self) -> str:
+        return (
+            f"pre={self.preprocess * 1e3:.1f}ms infer={self.inference * 1e3:.1f}ms "
+            f"post={self.postprocess * 1e3:.1f}ms total={self.total * 1e3:.1f}ms"
+        )
+
+
+# --------------------------------------------------------------------------- #
+# Preprocessing
+# --------------------------------------------------------------------------- #
+def load_image_rgb(image: Any) -> np.ndarray:
+    """Accept a path, a PIL image, or an HWC array. Returns RGB uint8."""
+    if isinstance(image, (str, Path)):
+        bgr = cv2.imread(str(image), cv2.IMREAD_COLOR)
+        if bgr is None:
+            raise FileNotFoundError(f"Could not read image: {image}")
+        return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+    if isinstance(image, np.ndarray):
+        if image.ndim == 2:
+            return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+        if image.shape[2] == 4:
+            return cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+        return image
+    return np.asarray(image.convert("RGB"))  # PIL
+
+
+def preprocess(images: Sequence[np.ndarray]) -> tuple[np.ndarray, list[tuple[int, int]]]:
+    """Resize to 800x800 (bicubic, no antialias) and scale to [0, 1] NCHW float32.
+
+    The HF processor uses torchvision resize with antialias=False specifically to
+    approximate cv2.resize, so cv2 INTER_CUBIC is the reference behaviour here.
+    """
+    batch = np.empty((len(images), 3, INPUT_SIZE, INPUT_SIZE), dtype=np.float32)
+    target_sizes: list[tuple[int, int]] = []
+    for i, img in enumerate(images):
+        h, w = img.shape[:2]
+        target_sizes.append((h, w))
+        resized = cv2.resize(img, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_CUBIC)
+        batch[i] = resized.astype(np.float32).transpose(2, 0, 1) * RESCALE_FACTOR
+    return batch, target_sizes
+
+
+# --------------------------------------------------------------------------- #
+# Post-processing (ported 1:1 from PPDocLayoutV3ImageProcessor)
+# --------------------------------------------------------------------------- #
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    return 1.0 / (1.0 + np.exp(-x, dtype=np.float64)).astype(np.float32)
+
+
+def get_order_seqs(order_logits: np.ndarray) -> np.ndarray:
+    """(B, Q, Q) pointer logits -> (B, Q) reading-order rank per query."""
+    scores = _sigmoid(order_logits)
+    batch_size, seq_len, _ = scores.shape
+
+    votes = np.triu(scores, 1).sum(axis=1) + np.tril(
+        1.0 - scores.transpose(0, 2, 1), -1
+    ).sum(axis=1)
+
+    pointers = np.argsort(votes, axis=1, kind="stable")
+    order_seq = np.empty_like(pointers)
+    ranks = np.broadcast_to(np.arange(seq_len), (batch_size, seq_len))
+    np.put_along_axis(order_seq, pointers, ranks, axis=1)
+    return order_seq
+
+
+def _extract_custom_vertices(polygon: np.ndarray, sharp_angle_thresh: float = 45) -> list[tuple]:
+    poly = np.array(polygon)
+    n = len(poly)
+    res = []
+    for i in range(n):
+        previous_point = poly[(i - 1) % n]
+        current_point = poly[i]
+        next_point = poly[(i + 1) % n]
+        v1 = previous_point - current_point
+        v2 = next_point - current_point
+        cross = (v1[1] * v2[0]) - (v1[0] * v2[1])
+        if cross < 0:
+            n1, n2 = np.linalg.norm(v1), np.linalg.norm(v2)
+            if n1 == 0 or n2 == 0:
+                res.append(tuple(current_point))
+                continue
+            angle = np.degrees(np.arccos(np.clip((v1 @ v2) / (n1 * n2), -1.0, 1.0)))
+            if abs(angle - sharp_angle_thresh) < 1:
+                direction = v1 / n1 + v2 / n2
+                direction = direction / np.linalg.norm(direction)
+                step = (n1 + n2) / 2
+                res.append(tuple(current_point + direction * step))
+            else:
+                res.append(tuple(current_point))
+    return res
+
+
+def _mask2polygon(mask: np.ndarray, epsilon_ratio: float = 0.004):
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return None
+    contour = max(contours, key=cv2.contourArea)
+    epsilon = epsilon_ratio * cv2.arcLength(contour, True)
+    approx = cv2.approxPolyDP(contour, epsilon, True)
+    points = np.atleast_2d(approx.squeeze())
+    return _extract_custom_vertices(points)
+
+
+def _extract_polygons(boxes: np.ndarray, masks: np.ndarray, scale_ratio) -> list:
+    scale_w, scale_h = scale_ratio[0] / MASK_STRIDE, scale_ratio[1] / MASK_STRIDE
+    mask_h, mask_w = masks.shape[1:]
+    polygons = []
+
+    for i in range(len(boxes)):
+        x_min, y_min, x_max, y_max = boxes[i].astype(np.int32)
+        box_w, box_h = int(x_max - x_min), int(y_max - y_min)
+        rect = np.array(
+            [[x_min, y_min], [x_max, y_min], [x_max, y_max], [x_min, y_max]], dtype=np.float32
+        )
+        if box_w <= 0 or box_h <= 0:
+            polygons.append(rect)
+            continue
+
+        x_start, x_end = np.clip(
+            [int(round(float(x_min * scale_w))), int(round(float(x_max * scale_w)))], 0, mask_w
+        )
+        y_start, y_end = np.clip(
+            [int(round(float(y_min * scale_h))), int(round(float(y_max * scale_h)))], 0, mask_h
+        )
+        cropped = masks[i, y_start:y_end, x_start:x_end]
+        if cropped.size == 0 or cropped.sum() == 0:
+            polygons.append(rect)
+            continue
+
+        resized = cv2.resize(cropped.astype(np.uint8), (box_w, box_h), interpolation=cv2.INTER_NEAREST)
+        polygon = _mask2polygon(resized)
+        if polygon is None or len(polygon) < 4:
+            polygons.append(rect)
+            continue
+        polygons.append(np.array(polygon, dtype=np.float32) + np.array([x_min, y_min]))
+    return polygons
+
+
+def postprocess(
+    logits: np.ndarray,
+    pred_boxes: np.ndarray,
+    order_logits: np.ndarray,
+    out_masks: np.ndarray | None,
+    target_sizes: Sequence[tuple[int, int]],
+    threshold: float = 0.5,
+) -> list[list[dict]]:
+    """Returns one list of detections per image, already sorted by reading order."""
+    order_seqs = get_order_seqs(order_logits)
+
+    # cxcywh (normalised) -> xyxy (absolute)
+    centers, dims = pred_boxes[..., :2], pred_boxes[..., 2:]
+    boxes = np.concatenate([centers - 0.5 * dims, centers + 0.5 * dims], axis=-1)
+    sizes = np.asarray(target_sizes, dtype=np.float32)          # (B, 2) as (h, w)
+    scale = np.stack([sizes[:, 1], sizes[:, 0], sizes[:, 1], sizes[:, 0]], axis=1)
+    boxes = boxes * scale[:, None, :]
+
+    batch_size, num_queries, num_classes = logits.shape
+    scores_all = _sigmoid(logits)
+
+    results: list[list[dict]] = []
+    for b in range(batch_size):
+        flat = scores_all[b].reshape(-1)
+        # torch.topk(k=num_queries) over the flattened (query, class) grid
+        top = np.argpartition(-flat, num_queries - 1)[:num_queries]
+        top = top[np.argsort(-flat[top], kind="stable")]
+
+        scores = flat[top]
+        labels = top % num_classes
+        query_idx = top // num_classes
+
+        keep = scores >= threshold
+        scores, labels, query_idx = scores[keep], labels[keep], query_idx[keep]
+
+        order = order_seqs[b][query_idx]
+        srt = np.argsort(order, kind="stable")
+        scores, labels, query_idx, order = scores[srt], labels[srt], query_idx[srt], order[srt]
+        sel_boxes = boxes[b][query_idx]
+
+        if out_masks is not None and len(sel_boxes):
+            masks = (_sigmoid(out_masks[b][query_idx]) > threshold).astype(np.uint8)
+            h, w = target_sizes[b]
+            polygons = _extract_polygons(sel_boxes, masks, [INPUT_SIZE / w, INPUT_SIZE / h])
+        else:
+            polygons = [
+                np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype=np.float32)
+                for x0, y0, x1, y1 in sel_boxes
+            ]
+
+        results.append(
+            [
+                {
+                    "order": int(o),
+                    "label_id": int(l),
+                    "label": ID2LABEL.get(int(l), str(l)),
+                    "score": float(s),
+                    "box": [round(float(v), 2) for v in box],
+                    "polygon": poly,
+                }
+                for o, l, s, box, poly in zip(order, labels, scores, sel_boxes, polygons)
+            ]
+        )
+    return results
+
+
+# --------------------------------------------------------------------------- #
+# Engine
+# --------------------------------------------------------------------------- #
+class PPDocLayoutV3ONNX:
+    """ONNX Runtime session, created once and reused."""
+
+    def __init__(
+        self,
+        onnx_path: str | Path,
+        *,
+        device: str = "cpu",
+        device_id: int = 0,
+        intra_op_num_threads: int | None = None,
+        threshold: float = 0.5,
+        trt_cache: str | None = None,
+        warmup: bool = True,
+    ) -> None:
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        if intra_op_num_threads:
+            so.intra_op_num_threads = intra_op_num_threads
+
+        providers: list[Any] = []
+        if device == "tensorrt":
+            providers.append((
+                "TensorrtExecutionProvider",
+                {
+                    "device_id": device_id,
+                    "trt_fp16_enable": True,
+                    "trt_engine_cache_enable": bool(trt_cache),
+                    "trt_engine_cache_path": trt_cache or "",
+                },
+            ))
+        if device in ("cuda", "tensorrt"):
+            providers.append(("CUDAExecutionProvider", {"device_id": device_id}))
+        providers.append("CPUExecutionProvider")
+
+        self.session = ort.InferenceSession(str(onnx_path), sess_options=so, providers=providers)
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_names = [o.name for o in self.session.get_outputs()]
+        self.has_masks = "out_masks" in self.output_names
+        self.threshold = threshold
+        self.last_timings = Timings()
+
+        if warmup:
+            self.session.run(
+                None, {self.input_name: np.zeros((1, 3, INPUT_SIZE, INPUT_SIZE), dtype=np.float32)}
+            )
+
+    @property
+    def providers(self) -> list[str]:
+        return self.session.get_providers()
+
+    def predict(
+        self, images: Any, threshold: float | None = None
+    ) -> list[dict] | list[list[dict]]:
+        """One image -> list of detections. A list of images -> list of those lists."""
+        single = not isinstance(images, (list, tuple))
+        image_list = [images] if single else list(images)
+        threshold = self.threshold if threshold is None else threshold
+
+        t0 = time.perf_counter()
+        rgb = [load_image_rgb(im) for im in image_list]
+        batch, target_sizes = preprocess(rgb)
+
+        t1 = time.perf_counter()
+        outputs = self.session.run(None, {self.input_name: batch})
+
+        t2 = time.perf_counter()
+        named = dict(zip(self.output_names, outputs))
+        results = postprocess(
+            named["logits"],
+            named["pred_boxes"],
+            named["order_logits"],
+            named.get("out_masks"),
+            target_sizes,
+            threshold=threshold,
+        )
+        t3 = time.perf_counter()
+        self.last_timings = Timings(t1 - t0, t2 - t1, t3 - t2)
+
+        return results[0] if single else results
+
+
+def draw(image_path: str | Path, detections: list[dict], out_path: str | Path) -> None:
+    """Quick visual sanity check: polygons + reading-order index."""
+    img = cv2.imread(str(image_path))
+    for det in detections:
+        poly = np.asarray(det["polygon"], dtype=np.int32).reshape(-1, 1, 2)
+        cv2.polylines(img, [poly], True, (0, 165, 255), 2)
+        x0, y0 = int(det["box"][0]), int(det["box"][1])
+        cv2.putText(
+            img, f"{det['order']}:{det['label']} {det['score']:.2f}",
+            (x0, max(y0 - 5, 12)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2,
+        )
+    cv2.imwrite(str(out_path), img)
+
+
+if __name__ == "__main__":
+    import argparse
+    import json
+
+    p = argparse.ArgumentParser(description="PP-DocLayoutV3 ONNX Runtime inference")
+    p.add_argument("--onnx", required=True)
+    p.add_argument("--image", required=True, nargs="+")
+    p.add_argument("--device", default="cpu", choices=["cpu", "cuda", "tensorrt"])
+    p.add_argument("--threshold", type=float, default=0.5)
+    p.add_argument("--threads", type=int, default=None)
+    p.add_argument("--draw", default=None, help="write an annotated copy of the first image")
+    args = p.parse_args()
+
+    det = PPDocLayoutV3ONNX(
+        args.onnx, device=args.device, intra_op_num_threads=args.threads, threshold=args.threshold
+    )
+    print(f"providers: {det.providers}")
+
+    results = det.predict(args.image)
+    if not isinstance(results[0], list):
+        results = [results]
+
+    for path, dets in zip(args.image, results):
+        print(f"\n=== {path} === ({det.last_timings})")
+        for d in dets:
+            print(f"  Order {d['order'] + 1}: {d['label']} {d['score']:.2f} {d['box']}")
+
+    if args.draw:
+        draw(args.image[0], results[0], args.draw)
+        print(f"\nannotated -> {args.draw}")
+    print(json.dumps({"count": [len(r) for r in results]}))