Как использовать ONNX для обнаружения объектов с помощью пользовательской модели yolov5s? ⇐ Python

[Anonymous]

Используя репозиторий yolov5 ultralytics, я обучил модель yolov5s обнаруживать определенные объекты, используя набор данных, аннотированный в Roboflow с одним классом объектов.
Эта модель отлично работает в репозитории Python при загрузке через концентратор горелки. Однако Torch был болезненным для упаковки зависимостей, поэтому я решил переключиться на onnx-runtime.
Я экспортировал свои веса yolo в onnx с помощью команды python3 Export.py --weights run/train/ X_v3/weights/best.pt --img 640 --batch 1 --include onnx.
Все еще в репозитории yolo я попробовал запустить python3 Detect.py --weights run/train/X_v3/weights/best.onnx --source test_image.jpg и все работает отлично. Для ясности: это test_image.jpg — одно из моих тестовых изображений из Roboflow (то есть оно было растянуто до 640 x 640, без какой-либо другой специальной предварительной обработки).
Затем я написал скрипт (см. ниже) для непосредственного тестирования onnx. Это произошло после того, как я поместил файл весов в netron.app, чтобы увидеть формат ввода и вывода.
Ограничивающая рамка вывода, нарисованная сценарием, находится непривычно далеко от реального объекта (довольно мала по сравнению с к объекту), а изображение onnx_test.png — то же самое, что я использовал в репозитории yolov5 для обнаружения. Что я упускаю?
Координаты кажутся правдоподобными. Я напечатал максимальные значения для x0, y0, x1, y1, и они составляют около 643 (что странно - yolo когда-нибудь дает координаты, выходящие за пределы размера изображения?). Я перевернул их при рисовании прямоугольника, но предполагаю, что именно так onnx выводит координаты.

Код: Выделить всё

import onnxruntime
import numpy as np
from PIL import Image, ImageDraw

# onnx expects `tensor: float32[1,3,640,640]`
image = Image.open('src/onnx_test.png')
image = image.resize((640, 640)) # convert to 640 x 640
image = image.convert('RGB') # 3 channel RGB
image_array = np.array(image).astype(np.float32)
image_array = np.transpose(image_array, (2, 0, 1)) # (height: 0, width: 1, channels: 2) -> (channels: 2, height: 0, width: 1)
image_array = np.expand_dims(image_array, axis=0) # add batch dim (1, ...)
print(image_array.shape)

session = onnxruntime.InferenceSession('src/processing/weights/v3.onnx')
input_name = session.get_inputs()[0].name
label_name = session.get_outputs()[0].name
pred_onx = session.run(
[label_name], {input_name: image_array})[0]
print(pred_onx.shape) # tensor: float32[1,25200,6]
print(np.round(
pred_onx[0][np.argmax(pred_onx[0][:, 4])],
2).tolist()) # get highest confidence inference

# OUTPUT: [357.510009765625, 247.05999755859375, 51.18000030517578, 32.560001373291016, 0.9599999785423279, 1.0]

image_draw = ImageDraw.Draw(image)
image_draw.rectangle((51, 32, 357, 247), outline ="red")
image.show()

Я ожидал увидеть ту же ограничивающую рамку, которую вижу при запуске define.py в репозитории yolov5. Однако результат намного меньше объекта и висит в верхнем левом углу изображения.

Подробнее здесь: https://stackoverflow.com/questions/793 ... ov5s-model