Output data extraction yolov5

47 Views Asked by At

I have to extract the boxes from an inference made with yolo5, I have the model weights in ONXX format. The inference is done onnxruntime.

In output I get the following values:

[[[4.95401955e+00 6.82735062e+00 1.05338869e+01 ... 1.51931345e-02
   6.55480444e-01 2.38993138e-01]
  [8.96611691e+00 6.16843033e+00 1.58229980e+01 ... 1.78753138e-02
   4.12897855e-01 4.19612855e-01]
  [1.64256668e+01 5.33465385e+00 2.31931877e+01 ... 2.36142874e-02
   2.44807631e-01 2.74172604e-01]
  ...
  [5.58593933e+02 6.20841675e+02 1.38702347e+02 ... 9.27328765e-02
   1.33206725e-01 2.77576238e-01]
  [5.95784424e+02 6.15207581e+02 1.13587395e+02 ... 4.44734097e-02
   1.17133498e-01 2.47565389e-01]
  [6.20036377e+02 6.13586914e+02 1.13036095e+02 ... 7.49872029e-02
   1.10100895e-01 2.02615947e-01]]]`

How can I extract the confidences and coordinates of the boxes efficiently?

I'm running the inference with this code:

import onnx
import onnxruntime as ort
import cv2
import numpy as np

def load_image(image_path, target_size=(640, 640)):
    # Carica l'immagine
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Ridimensiona l'immagine alla dimensione desiderata
    image = cv2.resize(image, target_size)

    # Trasponi e normalizza l'immagine
    image = np.transpose(image, (2, 0, 1))
    image = image.astype(np.float32) / 255.0
    image = np.expand_dims(image, axis=0)

    return image

def main():
    # Percorso del modello ONNX
    onnx_model_path = r"C:\Users\Giustino\Desktop\DPI_WORK\Small\weights\best.onnx"

    # Carica il modello ONNX
    onnx_model = onnx.load(onnx_model_path)

    # Carica un'immagine di input
    image_path = r"C:\Users\Giustino\Desktop\DPI_WORK\DPI\data\image0.jpg"
    input_image = load_image(image_path)

    # Crea un'istanza del runtime ONNX
    ort_session = ort.InferenceSession(onnx_model_path)

    # Esegui l'inferenza sull'immagine di input
    input_name = ort_session.get_inputs()[0].name
    output_name = ort_session.get_outputs()[0].name
    result = ort_session.run([output_name], {input_name: input_image})

    # Il risultato è disponibile in result[0]
    print("Risultato dell'inferenza:", result[0])

    outputs = result[0]  # Assuming result is a list and the detections are in the first element

if __name__ == "__main__":
    main()
0

There are 0 best solutions below