#!/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()