onnx_infer.py

#!/usr/bin/env python3
"""
ONNX inference for the TRIMMED YOLO11-OBB model.

Differences from YOLO26-OBB post-process:
  - reg_max = 16 (DFL enabled) -> auto-detected by box channel count
  - angle = (sigmoid(raw) - 0.25) * pi   <-- YOLO11 specific
  - cls scores = sigmoid(raw_logits)     <-- same as YOLO26
  - box is ltrb distance in feature-map units, multiplied by stride

Expected model outputs (9 tensors, NHWC):
    [box0, cls0, ang0, box1, cls1, ang1, box2, cls2, ang2]
with strides [8, 16, 32].

Usage:
    python3 onnx_infer.py \
        --model yolo11n-obb_1024x1024_trim.onnx \
        --img   boats.jpg \
        --output result_yolo11_obb.jpg
"""

import argparse
import math
import os

import cv2
import numpy as np
import onnxruntime as ort


DOTA_CLASSES = [
    "plane", "ship", "storage tank", "baseball diamond", "tennis court",
    "basketball court", "ground track field", "harbor", "bridge",
    "large vehicle", "small vehicle", "helicopter", "roundabout",
    "soccer ball field", "swimming pool",
]

DOTA_COLORS = [
    (255, 56, 56),   (255, 159, 56),  (255, 207, 56),  (180, 255, 56),
    (102, 255, 56),  (56, 255, 122),  (56, 255, 207),  (56, 207, 255),
    (56, 122, 255),  (102, 56, 255),  (180, 56, 255),  (255, 56, 207),
    (255, 56, 122),  (200, 200, 200), (128, 128, 255),
]


def preprocess_image(image, input_size=(1024, 1024), padding_value=114):
    """LetterBox + BGR->RGB + /255 -> float32 NCHW."""
    orig_h, orig_w = image.shape[:2]
    new_h, new_w = input_size
    r = min(new_h / orig_h, new_w / orig_w)
    new_unpad_w = round(orig_w * r)
    new_unpad_h = round(orig_h * r)
    dw = (new_w - new_unpad_w) / 2.0
    dh = (new_h - new_unpad_h) / 2.0
    if (orig_w, orig_h) != (new_unpad_w, new_unpad_h):
        image = cv2.resize(image, (new_unpad_w, new_unpad_h), interpolation=cv2.INTER_LINEAR)
    top = round(dh - 0.1)
    bottom = round(dh + 0.1)
    left = round(dw - 0.1)
    right = round(dw + 0.1)
    padded = cv2.copyMakeBorder(image, top, bottom, left, right,
                                cv2.BORDER_CONSTANT, value=(padding_value,) * 3)
    ratio_pad = (r, (left, top))
    rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    tensor = np.transpose(rgb, (2, 0, 1))[None, ...]
    return tensor, ratio_pad, (orig_h, orig_w)


def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)


def dfl_decode(box_pred, reg_max):
    n = box_pred.shape[0]
    box_pred = box_pred.reshape(n, 4, reg_max)
    box_pred = softmax(box_pred, axis=-1)
    proj = np.arange(reg_max, dtype=np.float32)
    return np.sum(box_pred * proj, axis=-1)


def decode_obb(box_preds, angle_preds, anchors, stride, reg_max):
    """dist2rbox -> xywhr in pixels. angle_preds here already processed (radians)."""
    if reg_max is not None and box_preds.shape[-1] == 4 * reg_max and reg_max > 1:
        box_preds = dfl_decode(box_preds, reg_max)

    angle = angle_preds.reshape(-1)
    cos_a = np.cos(angle)
    sin_a = np.sin(angle)

    lt = box_preds[:, :2]
    rb = box_preds[:, 2:]
    xf = (rb[:, 0] - lt[:, 0]) * 0.5
    yf = (rb[:, 1] - lt[:, 1]) * 0.5

    cx = xf * cos_a - yf * sin_a + anchors[:, 0]
    cy = xf * sin_a + yf * cos_a + anchors[:, 1]
    w = lt[:, 0] + rb[:, 0]
    h = lt[:, 1] + rb[:, 1]
    return np.stack([cx * stride, cy * stride, w * stride, h * stride, angle], axis=1)


def _get_covariance_matrix(boxes):
    a = (boxes[:, 2] ** 2) / 12.0
    b = (boxes[:, 3] ** 2) / 12.0
    c = boxes[:, 4]
    cos = np.cos(c)
    sin = np.sin(c)
    cos2 = cos * cos
    sin2 = sin * sin
    return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin


def batch_probiou(obb1, obb2, eps=1e-7):
    x1 = obb1[:, 0:1]
    y1 = obb1[:, 1:2]
    x2 = obb2[:, 0][None, :]
    y2 = obb2[:, 1][None, :]
    a1, b1, c1 = (v[:, None] for v in _get_covariance_matrix(obb1))
    a2_full, b2_full, c2_full = _get_covariance_matrix(obb2)
    a2 = a2_full[None, :]
    b2 = b2_full[None, :]
    c2 = c2_full[None, :]

    sum_ab = (a1 + a2) * (b1 + b2) - (c1 + c2) ** 2
    t1 = ((a1 + a2) * (y1 - y2) ** 2 + (b1 + b2) * (x1 - x2) ** 2) / (sum_ab + eps) * 0.25
    t2 = ((c1 + c2) * (x2 - x1) * (y1 - y2)) / (sum_ab + eps) * 0.5
    inner = ((a1 * b1 - c1 ** 2).clip(min=0.0) * (a2 * b2 - c2 ** 2).clip(min=0.0))
    t3 = np.log(sum_ab / (4.0 * np.sqrt(inner) + eps) + eps) * 0.5
    bd = np.clip(t1 + t2 + t3, eps, 100.0)
    hd = np.sqrt(1.0 - np.exp(-bd) + eps)
    return 1.0 - hd


def nms_rotated_probiou(rboxes, scores, classes, iou_thres, max_wh=7680.0, agnostic=False):
    if rboxes.size == 0:
        return np.empty((0,), dtype=np.int64)
    boxes = rboxes.copy()
    if not agnostic:
        offset = classes.astype(np.float32) * float(max_wh)
        boxes[:, 0] += offset
        boxes[:, 1] += offset
    order = np.argsort(-scores)
    sorted_boxes = boxes[order]
    ious = batch_probiou(sorted_boxes, sorted_boxes)
    n = sorted_boxes.shape[0]
    triu = np.triu(np.ones((n, n), dtype=bool), k=1)
    ious = ious * triu
    keep_mask = (ious >= iou_thres).sum(axis=0) <= 0
    return order[keep_mask]


def scale_rboxes_lefttop(rboxes, ratio_pad, orig_shape):
    rboxes = rboxes.copy()
    gain, (pad_x, pad_y) = ratio_pad
    rboxes[:, 0] -= pad_x
    rboxes[:, 1] -= pad_y
    rboxes[:, :4] /= gain
    rboxes[:, 0] = np.clip(rboxes[:, 0], 0, orig_shape[1])
    rboxes[:, 1] = np.clip(rboxes[:, 1], 0, orig_shape[0])
    return rboxes


def regularize_rbox(rboxes):
    rboxes = rboxes.copy()
    t_mod = np.mod(rboxes[:, 4], np.pi)
    swap = t_mod >= (np.pi / 2.0)
    if np.any(swap):
        w_old = rboxes[swap, 2].copy()
        rboxes[swap, 2] = rboxes[swap, 3]
        rboxes[swap, 3] = w_old
    rboxes[:, 4] = np.mod(rboxes[:, 4], np.pi / 2.0)
    return rboxes


def rbox_to_corners(rbox):
    cx, cy, w, h, ag = rbox
    cos_a, sin_a = math.cos(ag), math.sin(ag)
    wx, wy = w / 2.0 * cos_a, w / 2.0 * sin_a
    hx, hy = -h / 2.0 * sin_a, h / 2.0 * cos_a
    return np.array([
        [cx - wx - hx, cy - wy - hy],
        [cx + wx - hx, cy + wy - hy],
        [cx + wx + hx, cy + wy + hy],
        [cx - wx + hx, cy - wy + hy],
    ], dtype=np.float32)


def main():
    ap = argparse.ArgumentParser(description="YOLO11-OBB Trimmed ONNX Inference")
    ap.add_argument("-m", "--model", default="yolo11n-obb_1024x1024_trim.onnx",
                    dest="model_path")
    ap.add_argument("-i", "--img", default="boats.jpg", dest="test_img")
    ap.add_argument("-o", "--output", default="result_yolo11_obb.jpg",
                    dest="img_save_path")
    ap.add_argument("--score-thres", type=float, default=0.25)
    ap.add_argument("--nms-thres", type=float, default=0.45)
    ap.add_argument("--max-det", type=int, default=300)
    ap.add_argument("--agnostic-nms", action="store_true")
    opt = ap.parse_args()

    if not os.path.exists(opt.model_path):
        print(f"Model not found: {opt.model_path}")
        return
    if not os.path.exists(opt.test_img):
        print(f"Image not found: {opt.test_img}")
        return

    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    try:
        sess = ort.InferenceSession(opt.model_path, providers=providers)
    except Exception:
        sess = ort.InferenceSession(opt.model_path, providers=["CPUExecutionProvider"])

    input_name = sess.get_inputs()[0].name
    input_shape = sess.get_inputs()[0].shape
    imgsz = (int(input_shape[2]), int(input_shape[3]))
    output_names = [o.name for o in sess.get_outputs()]

    # The 9 outputs come out in graph order. We must sort them into
    # [box0, cls0, ang0, box1, cls1, ang1, box2, cls2, ang2] by their names
    # which the trim script writes as "{kind}_scale{i}_stride{s}".
    def key(name):
        kind_order = {"box": 0, "cls": 1, "ang": 2}
        parts = name.split("_")
        kind = parts[0]
        scale_idx = int(parts[1].replace("scale", ""))
        return (scale_idx, kind_order.get(kind, 9))

    sorted_names = sorted(output_names, key=key)

    img0 = cv2.imread(opt.test_img)
    if img0 is None:
        print(f"Cannot read image: {opt.test_img}")
        return

    img, ratio_pad, orig_shape = preprocess_image(img0.copy(), imgsz)
    raw_outputs = sess.run(sorted_names, {input_name: img.astype(np.float32)})

    strides = [8, 16, 32]
    conf_raw = -math.log(1.0 / opt.score_thres - 1.0)
    rboxes_all, scores_all, classes_all = [], [], []

    for scale_idx, stride in enumerate(strides):
        box_data = raw_outputs[scale_idx * 3 + 0]
        cls_data = raw_outputs[scale_idx * 3 + 1]
        ang_data = raw_outputs[scale_idx * 3 + 2]

        h, w = box_data.shape[1:3]
        box_channels = box_data.shape[-1]
        reg_max = None
        if box_channels >= 4 and box_channels % 4 == 0:
            reg_max = box_channels // 4

        box_data = box_data[0].reshape(-1, box_channels)
        cls_data = cls_data[0].reshape(-1, cls_data.shape[-1])
        ang_data = ang_data[0].reshape(-1, ang_data.shape[-1])

        if cls_data.shape[-1] == 1:
            cls_logits = cls_data[:, 0]
            cls_ids = np.zeros(len(cls_logits), dtype=np.int32)
        else:
            cls_logits = np.max(cls_data, axis=1)
            cls_ids = np.argmax(cls_data, axis=1)

        valid = cls_logits >= conf_raw
        if not np.any(valid):
            continue

        v_box = box_data[valid]
        v_ang_raw = ang_data[valid]
        v_score = 1.0 / (1.0 + np.exp(-cls_logits[valid]))
        v_id = cls_ids[valid]

        # YOLO11-OBB angle: (sigmoid(raw) - 0.25) * pi
        v_ang = (1.0 / (1.0 + np.exp(-v_ang_raw)) - 0.25) * np.pi

        gy, gx = np.indices((h, w))
        anchors = np.stack((gx.ravel(), gy.ravel()), axis=-1).astype(np.float32) + 0.5
        anchors = anchors[valid]

        rboxes = decode_obb(v_box, v_ang, anchors, stride, reg_max)
        rboxes_all.append(rboxes)
        scores_all.append(v_score)
        classes_all.append(v_id)

    if not rboxes_all:
        print("No detections found.")
        cv2.imwrite(opt.img_save_path, img0)
        return

    rboxes_all = np.concatenate(rboxes_all, axis=0).astype(np.float32)
    scores_all = np.concatenate(scores_all, axis=0).astype(np.float32)
    classes_all = np.concatenate(classes_all, axis=0).astype(np.int32)

    keep = nms_rotated_probiou(
        rboxes_all, scores_all, classes_all,
        iou_thres=opt.nms_thres, agnostic=opt.agnostic_nms,
    )
    keep = keep[: opt.max_det]
    if len(keep) == 0:
        print("No detections after NMS.")
        cv2.imwrite(opt.img_save_path, img0)
        return

    final_rboxes = scale_rboxes_lefttop(rboxes_all[keep], ratio_pad, orig_shape)
    final_rboxes = regularize_rbox(final_rboxes)
    final_scores = scores_all[keep]
    final_classes = classes_all[keep]

    print(f"Done! Found {len(final_rboxes)} oriented objects.")
    for i in range(len(final_rboxes)):
        cx, cy, w, h, theta = final_rboxes[i]
        conf = float(final_scores[i])
        cid = int(final_classes[i])
        name = DOTA_CLASSES[cid] if cid < len(DOTA_CLASSES) else f"cls{cid}"
        color = DOTA_COLORS[cid % len(DOTA_COLORS)]
        print(
            f"  {name:20s} conf={conf:.2f} cx={cx:.1f} cy={cy:.1f} "
            f"w={w:.1f} h={h:.1f} theta={math.degrees(theta):+.1f} deg"
        )
        corners = rbox_to_corners((cx, cy, w, h, theta)).astype(np.int32)
        cv2.polylines(img0, [corners], isClosed=True, color=color,
                      thickness=2, lineType=cv2.LINE_AA)
        label = f"{name} {conf:.2f}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        x_text, y_text = int(corners[0][0]), max(0, int(corners[0][1]) - 5)
        cv2.rectangle(img0, (x_text, y_text - th - 2),
                      (x_text + tw + 2, y_text + 2), color, -1)
        cv2.putText(img0, label, (x_text + 1, y_text - 1),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)

    cv2.imwrite(opt.img_save_path, img0)
    print(f"Result saved to {opt.img_save_path}")


if __name__ == "__main__":
    main()