CODEWITHSUNDEEP
pythonpytorchinferenceyolov5
TechXpert
TechXpert

Reputation: 175

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

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()

Upvotes: 5

Views: 19809

Answers (3)

rs_punia
rs_punia

Reputation: 449

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.

Upvotes: 0

Jordi
Jordi

Reputation: 71

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

Upvotes: 7

Louis Lac
Louis Lac

Reputation: 6406

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.

Upvotes: 0

Related Questions