在 openVINO 运行时上运行 YOLOv8

Question

我使用推荐的yolov8n.pt 库训练了一个检测模型。但它非常慢，推理时间太长，这使得该项目不切实际，但我知道在像 nvidia 这样的好 GPU 上它会快得多，但我的系统使用英特尔 CPU。使 intel 的 openvino 运行时成为优化选项。 我使用 openvino mo 命令将“onnx”模型转换为 IR 文件，然后继续加载权重以在 openvino 运行时上运行。

我的问题

输出形状似乎不规则，我尝试操作和查看 yolov8n.yaml 文件并将其与 IR 文件进行比较以理解输出。但我无法理解任何意义。 我尝试使用 yolov8s.pt 和不同的数据集，但 Openvino 优化后的输出形状保持不变

输入形状：[1,3,640,640]

输出形状：[1,6,8400]

import cv2
import matplotlib.pyplot as plt
import numpy as np
from openvino.runtime import Core
from ultralytics import YOLO
ie = Core()

model = ie.read_model(model=r"\models\IR\model_fit.xml")
compiled_model = ie.compile_model(model=model, device_name="CPU")

input_layer_ir = compiled_model.input(0)
output_layer_ir = compiled_model.output()


image = cv2.imread(r'\Resources\test_image.jpg')
N, C, H, W = input_layer_ir.shape
resized_image = cv2.resize(image, (W, H))
input_image = np.expand_dims(resized_image.transpose(2, 0, 1), 0)
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB));
boxes = compiled_model([input_image])[output_layer_ir]

print(boxes.shape)

<Output: names[output0] shape[1,6,8400] type: f32>

有人经历过这种情况吗？我该如何获取我的盒子坐标以进行检测？ 或者一些计算机视觉专家可以尝试这样做吗？注意我是计算机视觉的初学者

Answer 1

根据你的描述，表明输出形状是固定的，这不是你的预期，对吧？

我想你可以检查从 PyTorch 到 ONNX 的转换。在此阶段，您应该保持形状动态而不是静态。

如果可以的话，也许你可以分享一下原始的pytorch模型。

Answer 2

有点晚了，但我也偶然发现了这个问题，我相信从输出形状 [1,6,8400] 你看到你可以采取第一个值并转置为 [8400, 6] 给你 8400格式为 [x, y, 宽度, 高度, 概率类 0, 概率类 1] 的行，假设您有 2 个类。因此，处理代码将如下所示：

import cv2
import matplotlib.pyplot as plt
import numpy as np
from openvino.runtime import Core

ie = Core()

model = ie.read_model(model=r"\models\IR\model_fit.xml")
compiled_model = ie.compile_model(model=model, device_name="CPU")

input_layer_ir = compiled_model.input(0)
output_layer_ir = compiled_model.output()

image = cv2.imread(r'\Resources\test_image.jpg')
N, C, H, W = input_layer_ir.shape
resized_image = cv2.resize(image, (W, H))
input_image = np.expand_dims(resized_image.transpose(2, 0, 1), 0)
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB));
output = compiled_model([input_image])[output_layer_ir]

output = cv2.transpose(output[0])
print("Response shape", output.shape)

boxes = []
scores = []
class_ids = []

for row in output:
    # Each row is [x, y, width, height, probability class 0, probability class 1, ...]
    classes_scores = row[4:]
    (_min_score, max_score, _min_class_loc, (_x, max_class_index)) = cv2.minMaxLoc(classes_scores)
    if max_score >= 0.25:
        box = [row[0] - (0.5 * row[2]), row[1] - (0.5 * row[3]), row[2], row[3]]
        boxes.append(box)
        scores.append(max_score)
        class_ids.append(max_class_index)

[height, width, _] = image.shape
length = max((height, width))
scale = length/480

# Apply NMS (Non-maximum suppression)
RESULT_BOXES = cv2.dnn.NMSBoxes(boxes, scores, 0.5, 0.6, 0.5)
for index in RESULT_BOXES:
    box = boxes[index]
    x, y = round(box[0] * scale), round(box[1] * scale)
    x_plus_w, y_plus_h = round((box[0] + box[2]) * scale), round((box[1] + box[3]) * scale)
    draw_image = cv2.rectangle(cv2.UMat(image), (x, y), (x_plus_w, y_plus_h), (0, 0, 255), 4)

cv2.imwrite(r'\Resources\test_image_result.jpg', image)

这里有更多详细信息：https://github.com/ultralytics/ultralytics/blob/main/examples/YOLOv8-OpenCV-ONNX-Python/main.py

希望有帮助，如果有的话你能告诉我你是如何从 YoloV8 导出模型的吗？我已经做到了，但使用 IE Core 给我带来了奇怪的结果，而不是使用 YOLO 加载和运行