How to set and get confidence threshold from custom YOLOv5 model?

Question

I am trying to perform inference on my custom YOLOv5 model. The official documentation uses the default detect.py script for inference.

Example: python detect.py --source data/images --weights yolov5s.pt --conf 0.25

I have written my own python script but I can neither set the confidence threshold during initialisation nor retrieve it from the predictions of the model. I am only able to get the labels and bounding box coordinates. Here is my code:

import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt') 
results = model("my_image.png")
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()

Answer 1

It works for me:

model.conf = 0.25  # confidence threshold (0-1)
model.iou = 0.45  # NMS IoU threshold (0-1)  

More information:

https://github.com/ultralytics/yolov5/issues/36

Answer 2

Check the detection.y file for an explicit explanation on own to use the model. You have to do the score thresholding manually after the model inference in your code.

This means that detection scores are ranging from 0.0 to 1.0 and you must filter them with your confidence threshold.

Answer 3

To set the confidence threshold of the custom-trained YOLOv5 model, Use the following:

import torch
model = torch.hub.load('ultralytics/yolov5', 'custom',
                       path='absolute/path/to/.pt/file', source='local')
model.conf = 0.25

To retrieve the inference image and the other values, use

# Inference Result
img = "my_image.png"
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.