File size: 4,723 Bytes
103e15d | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | #!/usr/bin/env python3
# =============================================================================
# YOLO26-OBB Export ONNX Model Script
#
# Copyright (c) 2025, AXERA Semiconductor Co., Ltd. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use
# this file except in compliance with the License. You may obtain a copy of
# the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
# =============================================================================
#
# Author: GUOFANGMING
#
import torch
from ultralytics import YOLO
from ultralytics.nn.modules.head import OBB26
import os
import shutil
def npu_obb26_forward(self, x):
"""
YOLO26 OBB26 Head Modified for NPU.
OBB26 structure:
- cv2 / one2one_cv2: box regression layers (output 4*reg_max channels, reg_max=1 -> 4)
- cv3 / one2one_cv3: classification layers (output nc channels, nc=15 for DOTAv1)
- cv4 / one2one_cv4: angle regression layers (output ne channels, ne=1)
Note:
OBB26 outputs raw angle predictions (in radians) without sigmoid transformation,
which is different from the original OBB head where the angle is wrapped via
``(angle.sigmoid() - 0.25) * pi`` in ``forward_head``. The post-process therefore
consumes the raw angle directly via ``cos(angle)`` / ``sin(angle)``.
Output:
List of Tensors (9 items: 3 scales * 3 outputs), Layout: NHWC for all branches.
For each scale:
- Box_Raw (B, H, W, 4*reg_max), <-- bbox regression (ltrb distance, reg_max=1)
- Cls_Raw (B, H, W, nc), <-- class scores (raw logits)
- Angle_Raw (B, H, W, ne), <-- rotation angle (raw radians, ne=1)
"""
if not isinstance(x, (list, tuple)):
x = [x]
res = []
# Box / Cls use one2one branch (end2end mode) when available
if hasattr(self, 'one2one_cv2') and hasattr(self, 'one2one_cv3'):
box_layers = self.one2one_cv2
cls_layers = self.one2one_cv3
else:
box_layers = self.cv2
cls_layers = self.cv3
# Angle uses one2one_cv4 in end2end mode
if hasattr(self, 'one2one_cv4'):
angle_layers = self.one2one_cv4
else:
angle_layers = self.cv4
for i in range(self.nl):
# 1. Box branch - NHWC (4 * reg_max channels, reg_max=1 -> 4)
bboxes = box_layers[i](x[i]).permute(0, 2, 3, 1)
# 2. Cls branch - NHWC (nc channels, raw logits)
scores = cls_layers[i](x[i]).permute(0, 2, 3, 1)
# 3. Angle branch - NHWC (ne channels, raw radians)
angle = angle_layers[i](x[i]).permute(0, 2, 3, 1)
res.append(bboxes)
res.append(scores)
res.append(angle)
return res
def batch_export_yolo26_obb():
variants = ['n', 's', 'm', 'l', 'x']
imgsz = 1024 # YOLO26-OBB officially uses 1024x1024 (DOTAv1 pretrained)
# Execute Monkey Patch
OBB26.forward = npu_obb26_forward
print("Monkey patch applied for OBB26: Output Layout forced to NHWC (Box, Cls, Angle for each scale).")
for v in variants:
model_name = f"yolo26{v}-obb"
pt_path = f"{model_name}.pt"
onnx_final_name = f"{model_name}_{imgsz}x{imgsz}.onnx"
print(f"\n--- Processing {model_name} ---")
try:
# Load model
model = YOLO(pt_path)
# Reapply monkey patch
OBB26.forward = npu_obb26_forward
# Ensure the model's head also uses the new forward
if hasattr(model.model, 'model') and len(model.model.model) > 0:
head = model.model.model[-1]
if isinstance(head, OBB26):
head.forward = lambda x: npu_obb26_forward(head, x)
# Execute export
exported_path = model.export(
format="onnx",
imgsz=imgsz,
dynamic=False,
opset=11,
simplify=True,
nms=False
)
# Move and rename
if exported_path:
shutil.move(exported_path, onnx_final_name)
print(f"Success: {onnx_final_name}")
except Exception as e:
print(f"Failed to export {model_name}: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
batch_export_yolo26_obb()
|