Result of onnxruntime inference

[bruceYangxianglin]

Search before asking

• [X] I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

When I use 'torch.hub.load()' to load custom trained model, the result of model output is [ xmin ymin xmax ymax confidence class]. But If I use 'onnxruntime.InferenceSession' to load my model, the shape of result is (1, 25200, 6). I do not know how can I convert these result to binding box. I notice that whether I fix the image size or not, the resolution of saved result image is the same. How can I achieve the same result when I use onnxruntime?

Additional

No response

[github-actions[bot]]

👋 Hello @bruceYangxianglin, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

If this is a 🐛 Bug Report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online W&B logging if available.

For business inquiries or professional support requests please visit https://ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.

Requirements

Python>=3.6.0 with all requirements.txt installed including PyTorch>=1.7. To get started:

$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ pip install -r requirements.txt

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab and Kaggle notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), validation (val.py), inference (detect.py) and export (export.py) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.

[glenn-jocher]

@bruceYangxianglin YOLOv5 PyTorch Hub models are AutoShape() instances. ONNX workflows are not enabled for AutoShape() instances, but you can use detect.py as an example of ONNX inference:

python detect.py --weights yolov5s.onnx

[cipher982]

@bruceYangxianglin I spent the past couple of days working on converting my current project from PyTorch to ONNX inference and I think I finally achieved the results I was looking for after a lot of looking through the current code repository. I have found the minimum needed functions needed to detect and frame an image are as follows:

from utils.augmentations import letterbox
from utils.general import non_max_suppression, scale_coords, xywh2xyxy
from utils.plots import Annotator, colors

I then used a select few portions from detect.py to create the following end-to-end example below:

# Setup imports
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
from pathlib import Path

import cv2
import numpy as np
import onnxruntime
import onnx
from matplotlib import pyplot as plt
import torch

from utils.general import non_max_suppression, scale_coords, xywh2xyxy
from utils.plots import Annotator, colors
from utils.augmentations import letterbox

%matplotlib inline

# Load class names
class_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

# Load the ONNX model
ONNX_PATH = "./models/yolov5s.onnx"
IMG_PATH = "./classifier/images/bird_0295_1637697315497.jpg"
HALF = False

# Load a sample image
# If you already have an image as an array skip to letterbox step
def read_img(fpath, half):
    img0 = cv2.imread(fpath)
    img = letterbox(img0, new_shape=(960, 960))[0]
    img = np.moveaxis(img, -1, 0)
    img = torch.from_numpy(img).to("cpu")
    img = img.half() if half else img.float() 
    img /= 255  # 0 - 255 to 0.0 - 1.0
    if len(img.shape) == 3:
        img = img[None]
    return img.numpy(), img0

img, img0 = read_img(IMG_PATH, HALF)

# Run inference with ONNX
ort_session = onnxruntime.InferenceSession(ONNX_PATH)
y = ort_session.run([ort_session.get_outputs()[0].name], {ort_session.get_inputs()[0].name: img})[0]
y = torch.tensor(y) if isinstance(y, np.ndarray) else y

# Run ONNX specific modifications to the model output
pred = non_max_suppression(y)

# Interpret predictions and show image
def show_img(pred, im_path, img0, 
             img, names, view_img=True,
             hide_labels=False, hide_conf=False,
            ):
    s = f'image {0}/{0} {im_path}: '
    for i, det in enumerate(pred):  # per image
        print(det)
        p, im0, frame = im_path, img0.copy(), img
        #p = Path(p)  # to Path
        #save_path = str(save_dir / p.name)  # im.jpg
        s += '%gx%g ' % img.shape[2:]  # print string

        gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
        imc = im0  # for save_crop
        annotator = Annotator(im0, line_width=3, example=str(names))

        if len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
            # Print results
            for c in det[:, -1].unique():
                n = (det[:, -1] == c).sum()  # detections per class
                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
                print(s)
            # Write results
            for *xyxy, conf, cls in reversed(det):
                if view_img:  # Add bbox to image
                    c = int(cls)  # integer class
                    label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                    annotator.box_label(xyxy, label, color=colors(c, True))

        im0 = annotator.result()
        plt.imshow(im0)

show_img(pred, Path(IMG_PATH), img0, img, class_names)

[bruceYangxianglin]

@cipher982 I think we have spent the same effort to achieve the same purpose with python. And I have to finally preprocess images and process the result with PLC language but not python. That would be a hard task.

[github-actions[bot]]

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

• Wiki – https://github.com/ultralytics/yolov5/wiki
• Tutorials – https://docs.ultralytics.com/yolov5
• Docs – https://docs.ultralytics.com

Access additional Ultralytics ⚡ resources:

• Ultralytics HUB – https://ultralytics.com/hub
• Vision API – https://ultralytics.com/yolov5
• About Us – https://ultralytics.com/about
• Join Our Team – https://ultralytics.com/work
• Contact Us – https://ultralytics.com/contact

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!