File size: 3,683 Bytes
aa9fdf6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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']}")
|