inference.py

"""Standalone Inference Script for DocLayout-YOLO ONNX."""
from __future__ import annotations

import argparse
import ast
import os
from pathlib import Path
from typing import Any, Dict, List, Tuple

import cv2
import numpy as np
import onnx
import onnxruntime as ort


def load_model(model_path: str = "doclayout_yolo_docstructbench_imgsz1024.onnx"):
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")

    model_proto = onnx.load(model_path)
    metadata = {d.key: d.value for d in model_proto.metadata_props}
    stride = ast.literal_eval(metadata.get("stride", "32"))
    names = ast.literal_eval(metadata.get("names", "{}"))

    providers = ort.get_available_providers()
    print(f"Available ONNX Providers: {providers}")

    session = ort.InferenceSession(model_path, providers=providers)
    print(f"Active Provider: {session.get_providers()[0]}")
    return session, stride, names


def resize_and_pad(image: np.ndarray, target_size: int = 1024, stride: int = 32):
    h, w = image.shape[:2]
    r = min(target_size / h, target_size / w)
    new_h, new_w = int(round(h * r)), int(round(w * r))
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    pad_w = (target_size - new_w) % stride
    pad_h = (target_size - new_h) % stride
    dw, dh = pad_w / 2, pad_h / 2
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))

    padded = cv2.copyMakeBorder(
        resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)
    )
    return padded, r, (left, top)


def predict(session, stride, names, image_path: str, conf_threshold: float = 0.25):
    img_bgr = cv2.imread(image_path)
    if img_bgr is None:
        raise ValueError(f"Could not read image from {image_path}")

    orig_h, orig_w = img_bgr.shape[:2]
    padded, ratio, (pad_x, pad_y) = resize_and_pad(img_bgr, target_size=1024, stride=stride)

    blob = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB).transpose(2, 0, 1)
    blob = np.expand_dims(blob, axis=0).astype(np.float32) / 255.0

    input_name = session.get_inputs()[0].name
    preds = session.run(None, {input_name: blob})[0]
    if len(preds.shape) == 3:
        preds = preds[0]

    detections = []
    for row in preds:
        conf = float(row[4])
        if conf < conf_threshold:
            continue
        cls_id = int(row[5])
        class_name = names.get(cls_id, str(cls_id))

        x0 = max(0.0, (float(row[0]) - pad_x) / ratio)
        y0 = max(0.0, (float(row[1]) - pad_y) / ratio)
        x1 = min(float(orig_w), (float(row[2]) - pad_x) / ratio)
        y1 = min(float(orig_h), (float(row[3]) - pad_y) / ratio)

        detections.append({
            "box": [round(x0, 2), round(y0, 2), round(x1, 2), round(y1, 2)],
            "label": class_name,
            "confidence": round(conf, 4),
        })

    return detections


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="DocLayout-YOLO Inference")
    parser.add_argument("--image", required=True, help="Path to input image (PNG/JPG)")
    parser.add_argument("--model", default="doclayout_yolo_docstructbench_imgsz1024.onnx", help="Path to ONNX model")
    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
    args = parser.parse_args()

    session, stride, names = load_model(args.model)
    results = predict(session, stride, names, args.image, conf_threshold=args.conf)
    print(f"\nDetected {len(results)} layout elements:")
    for r in results:
        print(f"  - [{r['label'].upper()}] (conf: {r['confidence']}): {r['box']}")