| import os |
| import cv2 |
| import math |
| import numpy as np |
| from time import time |
| import argparse |
| import logging |
| import axengine as ort |
|
|
| logging.basicConfig( |
| level=logging.DEBUG, |
| format='[%(name)s] [%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s', |
| datefmt='%H:%M:%S' |
| ) |
| logger = logging.getLogger("YOLO26-OBB") |
|
|
| 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 infer_hw_layout(shape): |
| """Model input (NCHW or NHWC) -> H, W, layout name.""" |
| shape = list(shape) |
| if len(shape) == 4 and shape[-1] == 3: |
| h = int(shape[1] or 1024) |
| w = int(shape[2] or 1024) |
| return h, w, "NHWC" |
| if len(shape) == 4 and shape[1] == 3: |
| h = int(shape[2] or 1024) |
| w = int(shape[3] or 1024) |
| return h, w, "NCHW" |
| return 1024, 1024, "NCHW" |
|
|
|
|
| def preprocess_image(image, input_size=(1024, 1024), layout="NCHW", padding_value=114): |
| """LetterBox + BGR→RGB; uint8 NHWC or NCHW (same geometry as onnx_infer letterbox).""" |
| 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 |
| dh = new_h - new_unpad_h |
| dw /= 2.0 |
| dh /= 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) |
| if layout == "NHWC": |
| tensor = rgb[None, ...].astype(np.uint8) |
| else: |
| tensor = np.transpose(rgb, (2, 0, 1))[None, ...].astype(np.uint8) |
| return tensor, ratio_pad, (orig_h, orig_w) |
|
|
|
|
| def softmax(x, axis=-1): |
| e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) |
| return e_x / np.sum(e_x, 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=None): |
| """dist2rbox -> xywhr in pixels.""" |
| 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): |
| """Fast NMS with probiou.""" |
| 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] = boxes[:, 0] + offset |
| boxes[:, 1] = 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): |
| """Letterbox inverse for xywhr centers.""" |
| 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): |
| """Theta in [0, pi/2); swap w/h when needed.""" |
| 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='YOLO26-OBB Inference (AXERARuntime)') |
| ap.add_argument('--model-path', type=str, default='yolo26n-obb_1024x1024.axmodel') |
| ap.add_argument('--test-img', type=str, default='boats.jpg') |
| ap.add_argument('--img-save-path', type=str, default='result_yolo26_obb.jpg') |
| ap.add_argument('--score-thres', type=float, default=0.25) |
| ap.add_argument('--nms-thres', type=float, default=0.45) |
| ap.add_argument('--num-classes', type=int, default=15) |
| ap.add_argument('--max-det', type=int, default=300) |
| ap.add_argument('--agnostic-nms', action='store_true', help='Class-agnostic NMS.') |
| ap.add_argument('--providers', type=str, default='AxEngineExecutionProvider') |
| opt = ap.parse_args() |
|
|
| if not os.path.exists(opt.model_path): |
| logger.error(f"Model not found: {opt.model_path}") |
| return |
| if not os.path.exists(opt.test_img): |
| logger.error(f"Image not found: {opt.test_img}") |
| return |
|
|
| t0 = time() |
| providers = [p.strip() for p in opt.providers.split(",") if p.strip()] or None |
| sess = ort.InferenceSession(opt.model_path, providers=providers) |
| logger.debug(f"\033[1;31mLoad model time = {(time() - t0) * 1000:.2f} ms\033[0m") |
|
|
| inp = sess.get_inputs()[0] |
| input_name = inp.name |
| m_h, m_w, layout = infer_hw_layout(inp.shape) |
|
|
| img0 = cv2.imread(opt.test_img) |
| if img0 is None: |
| logger.error(f"Cannot read image: {opt.test_img}") |
| return |
|
|
| t0 = time() |
| input_tensor, ratio_pad, orig_shape = preprocess_image( |
| img0.copy(), input_size=(m_h, m_w), layout=layout |
| ) |
| logger.debug(f"\033[1;31mPre-process time = {(time() - t0) * 1000:.2f} ms\033[0m") |
|
|
| t0 = time() |
| output_names = [o.name for o in sess.get_outputs()] |
| outputs = sess.run(output_names, {input_name: input_tensor}) |
| logger.debug(f"\033[1;31mForward time = {(time() - t0) * 1000:.2f} ms\033[0m") |
|
|
| t0 = time() |
| 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_idx = scale_idx * 3 |
| cls_idx = scale_idx * 3 + 1 |
| ang_idx = scale_idx * 3 + 2 |
| if ang_idx >= len(outputs): |
| continue |
|
|
| box_data = outputs[box_idx] |
| cls_data = outputs[cls_idx] |
| ang_data = outputs[ang_idx] |
|
|
| 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 = ang_data[valid] |
| v_score = 1.0 / (1.0 + np.exp(-cls_logits[valid])) |
| v_id = cls_ids[valid] |
|
|
| 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 len(rboxes_all) == 0: |
| logger.info("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] |
| logger.debug(f"\033[1;31mPost-process time = {(time() - t0) * 1000:.2f} ms\033[0m") |
|
|
| if len(keep) == 0: |
| logger.info("No detections after NMS.") |
| cv2.imwrite(opt.img_save_path, img0) |
| return |
|
|
| final_rboxes = rboxes_all[keep] |
| final_scores = scores_all[keep] |
| final_classes = classes_all[keep] |
|
|
| final_rboxes = scale_rboxes_lefttop(final_rboxes, ratio_pad, orig_shape) |
| final_rboxes = regularize_rbox(final_rboxes) |
|
|
| logger.info(f"\033[1;32mDraw Results ({len(final_rboxes)} oriented objects):\033[0m") |
| 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)] |
| logger.info( |
| 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) |
| logger.info(f"Saved to {opt.img_save_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
