yolov5 custom trained weights converted to ONNX showing wrong labels

user_15854271 0
Viewed 581

After converting custom trained Yolov5 weights (.pt) to ONNX and running inference on the ONNX file using:

https://github.com/BlueMirrors/Yolov5-ONNX.git

the detection works well but my image labels/classes are using COCO labels (ie person, airplane, etc) instead of my labels. How can I change the labels to my own. I'm unsure of the formatting of the json (or yaml?) file. Thanks!

def detect_image(device, weight, image_path, output_image):
    # load model
    model = Yolov5Onnx(classes="coco",
                       backend="onnx",
                       weight=weight,
                       device=device)

    # read image
    image = cv2.imread(image_path)

    # inference
    preds = model(image)
    print(preds)

    # draw image
    preds.draw(image)

    # write image
    cv2.imwrite(output_image, image)
yolov5 onnx python
2 Answers

If you would like this is a workaround solution using ultralytics's detect.py (source: https://github.com/ultralytics/yolov5/blob/master/detect.py)

Modify line 93

device = select_device(device)
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
model.names = ['your_class_name', 'your_class_name2'] #Add this line
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride)  # check image size
user_125513950

You can see here: https://github.com/BlueMirrors/cvu/blob/3efb3758cdc0bdc73fe321bc7f566c0f383d49a6/cvu/detector/yolov5/core.py#L35, that the classes argument in Yolov5Onnx takes in a list. So you can just pass a list containing you class names. My model predict 12 different classes, so I passed a 12 element long string list to this argument:

from cvu.detector.yolov5 import Yolov5 as Yolov5Onnx
from vidsz.opencv import Reader, Writer
import numpy as np
import cv2


def detect_image(device, weight, input_frame, output_image):
    # load model
    class_names = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11']
    model = Yolov5Onnx(classes=class_names,
                       backend="onnx",
                       weight=weight,
                       device=device)

    # read image
    image = cv2.imread(image_path)
    resized = cv2.resize(image, (640, 640), interpolation = cv2.INTER_AREA)

    # inference
    preds = model(resized)
    print(preds)

    # draw image
    preds.draw(resized)

    # write image
    cv2.imwrite(output_image, resized)

device = 'cpu'
weight = '/path/to/best.onnx'
image_path = '/path/to/img.jpg'
output_image = './out_frame.jpg'
detect_image(device, weight, image_path, output_image)

Running this with my model gives the following output:

(YDSO) ➜  cvu git:(master) ✗ python3 test.py
id:0    class=10;       conf=0.89;      top-left=(459.0, 1.0  );        bottom-right=(640.0, 531.0)
id:1    class=10;       conf=0.81;      top-left=(0.0, 2.0  );  bottom-right=(472.0, 560.0)
user_95067730
Related