yolov8-obb旋转框目标检测onnxruntime部署
发布时间
阅读量：
阅读量
yolov8-obb使用onnxruntime部署
一、部署代码大致流程
大致流程为读取图片–>图像预处理–>模型推理–>图像后处理–>可视化。
二、部署过程详解
1、模型导出
首先将模型文件导出为onnx格式
from ultralytics import YOLO
model=YOLO("") #引号中间填入模型路径
model.export(format='onnx')
AI写代码2、模型文件分析
使用Netron来查看模型的输入、输出以及类别：
可以看到我的模型输入是[1,3,1024,1024]，输出为[1,6,21504]，输出中的6代表：
6=边界框四个坐标（x,y,w,h）+每个类别的置信度（scores）+角度(angle)
网络会将输入图像划分为3个尺度（8、16、32）的网格，每个尺度下的网格数量为128 128、64 64、32*32，加在一起为21504个网格对象：
128*128+64*64+32*32=21504
3、代码实现
首先对读入的图片进行预处理，包括letterbox、归一化等操作，letterbox操作目的就是将原图的尺寸转换成网络的输入尺寸，采用的方式是等比例缩放的方式，先找出长边将其缩放成1024，按照长边的缩放比例同时给短边进行缩放，然后把短边补充灰边至目标尺寸。
def letterbox(img, input_shape, color=(114, 114, 114)): #letterbox将原图尺寸转换成模型输入尺寸
shape = img.shape[:2] #获取图片宽高
r = min(input_shape[0] / shape[0], input_shape[1] / shape[1]) #计算得到缩放比例
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) #按照比例缩放
dw, dh = (input_shape[1] - new_unpad[0]) / 2, (input_shape[0] - new_unpad[1]) / 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))
if shape[::-1] != new_unpad: #如果shape和按照比例缩放后的new_unpad不相等，直接对img进行resize
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) #给图片添加边框
return img
AI写代码预处理函数：
def Preprocess(image): #图像预处理
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) #BGR->RGB
image = letterbox(image,input_shape).astype(np.float32)
image = image / 255.0 #对图片进行归一化
image = np.transpose(image, (2, 0, 1)) #将(w,h,c)变为(c,w,h)
input_tensor = np.expand_dims(image, axis=0) #增加一个维度变为(b,c,w,h)
return input_tensor
AI写代码推理模块：
def Inference(input_tensor, onnx_path): #推理过程
# sess_options = ort.SessionOptions()
# sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
# sess_options.intra_op_num_threads = 4
session_model = ort.InferenceSession( #官方封装好的推理函数，拿来直接用
path_or_bytes=onnx_path, #传入模型路径
providers = ['CPUExecutionProvider'], #我这里部署采用CPU
# sess_options=sess_options
)
inputs = {session_model.get_inputs()[0].name: input_tensor}
outputs = session_model.run(None, inputs)
return outputs
AI写代码后处理模块：图片经过模型推理之后得到预测信息，需要将该信息进行后处理之后才能得到最后的目标检测框信息，主要包括NMS和Scale_boxes模块。下面为调用官方的NMSBoxesRotated函数实现，也可以使用手搓NMS实现，只需将代码中的xywh2xyxy和nms解除注释，并将官方NMS注释掉即可调用。
def filter_box(outputs):
outputs = np.squeeze(outputs) #去掉大小为1的维度
detected_boxes = []
rotated_boxes = [] #用来存放预测框信息
scores = [] #存放置信度信息
class_ids = [] #存放索引
classes_scores = outputs[4:(4 + len(class_names)), ...] #所有预测框的置信度
angles = outputs[-1, ...] #所有预测框的角度信息
for i in range(outputs.shape[1]): #遍历所有预测框
class_id = np.argmax(classes_scores[..., i]) #获取该预测框置信度然后获取其索引
score = classes_scores[class_id][i] #根据索引获取置信度
angle = angles[i] #根据索引获取该预测框角度
if 0.5 * math.pi <= angle <= 0.75 * math.pi: #处理角度
angle -= math.pi
if score > score_threshold: #处理置信度大于置信度阈值的预测框，置信度低于阈值的预测框直接舍弃
rotated_boxes.append(np.concatenate([outputs[:4, i], np.array([score, class_id, angle * 180 / math.pi])])) #将该预测框的坐标、置信度、索引、角度进行拼接后放入rotated_boxes
scores.append(score) #将其置信度放入scores
class_ids.append(class_id)
rotated_boxes = np.array(rotated_boxes) #将list转为array
# boxes = xywh2xyxy(rotated_boxes) #将坐标的中心点和宽高表达形式转为左上角右下角形式
scores = np.array(scores) #将list转为array
# indices = nms(boxes, scores, nms_threshold) #手搓NMS过滤预测框
for boxes in rotated_boxes: #先把要传入的boxes信息处理成官方NMS数据接受的格式
detected_boxes.append(((boxes[0],boxes[1]),(boxes[2],boxes[3]),boxes[4]))
indices = cv2.dnn.NMSBoxesRotated(detected_boxes, scores, score_threshold, nms_threshold) #调用官方的旋转矩形NMS，返回索引
output = rotated_boxes[indices] #根据NMS返回的索引保留预测框信息
return output
AI写代码 def nms(boxes, scores, nms_threshold):
x1 = boxes[:, 0] #分别获取剩余预测框的左上角和右下角坐标信息
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1) #根据坐标计算预测框面积
keep = [] #用来存放索引
index = scores.argsort()[::-1] #将剩余预测框的置信度从大到小排列并获取它们的索引
while index.size > 0:
i = index[0] #第一个元素即为置信度最高的预测框
keep.append(i) #存入keep中
x11 = np.maximum(x1[i], x1[index[1:]]) #获取两两预测框相交部分的左上角和右下角坐标信息
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1) #根据计算的到的坐标得到相交部分的宽高
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h #计算相交部分的面积
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps) #计算iou，即相交面积除以合并面积
idx = np.where(ious <= nms_threshold)[0] #获取iou小于NMS阈值的索引，大于该阈值的预测框直接舍弃
index = index[idx + 1] #将其索引全部+1，因为还得算上最大置信度的那个预测框
return keep #经过循环筛选iou值之后返回最有可能为目标的预测框索引
def xywh2xyxy(x):
y = np.copy(x) #拷贝一份预测框信息
y[:, 0] = x[:, 0] - x[:, 2] / 2 #根据中心点和宽高信息分别计算得到左上角和右下角坐标信息
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y
AI写代码scale_boxes模块的作用是是将预测结果映射回到原始输入图片尺寸。
def scale_boxes(boxes, shape):
gain = min(input_shape[0] / shape[0], input_shape[1] / shape[1]) #获取缩放比例
pad = (input_shape[1] - shape[1] * gain) / 2, (input_shape[0] - shape[0] * gain) / 2 #高度宽度上的填充大小
boxes[..., [0, 1]] -= pad # xy padding #减去填充大小
boxes[..., :4] /= gain #除缩放比例恢复原始尺寸
return boxes
AI写代码最后将预测结果可视化。
def draw(image, box_data):
box_data = scale_boxes(box_data, image.shape)
boxes = box_data[..., :4] #获取预测框坐标信息
scores = box_data[..., 4] #获取置信度
classes = box_data[..., 5].astype(np.int32) #获取索引
angles = box_data[..., 6] #获取角度
for box, score, cls, angle in zip(boxes, scores, classes, angles):
rotate_box = ((box[0], box[1]), (box[2], box[3]), angle) #接下来就是获取坐标和角度把预测框画上去了
points = cv2.boxPoints(rotate_box)
points = np.int32(points)
cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=2)
cv2.putText(image, '{0} {1:.2f}'.format(class_names[cls], score), (points[0][0], points[0][1]),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
AI写代码推理时需要用到的参数和主函数：
class_names = ["1"] #检测的所有类别
score_threshold = 0.5 #置信度阈值,置信度低于此阈值的预测框将被忽略
nms_threshold = 0.4 #nms阈值，高于此阈值的预测框将被抑制
input_shape = (1024, 1024) #输入尺寸
onnx_path = "best2.onnx" #模型路径
image = cv2.imread("test.png") #输入图片
input_tensor = Preprocess(image) #将图片放入预处理模块
outputs = Inference(input_tensor, onnx_path)
boxes = filter_box(outputs)
draw(image, boxes)
cv2.imwrite('result.png', image)
AI写代码推理前：
推理结果可视化后：
三、完整代码
完整代码下载链接：yolov8-obb旋转框目标检测onnxruntime部署完整代码
全部评论 (0)
还没有任何评论哟~