1 prediction per image in a batch

[ahmadmustafaanis]

Search before asking

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

Question

Let's say I have a batch of images. i.e (8, 20,20,3) and I want to get predictions on it using Yolov5. Now I want it to predict exactly 1 bbx for each image in the batch. i.e model(batch) should output 8 tensors. If there is no object found, it should return empty tensor for that else it should return the highest confidence object. I was looking in the articles/documentation but can not find the solution to it.

Additional

No response

[glenn-jocher]

@ahmadmustafaanis 👋 Hello! Thanks for asking about handling inference results. YOLOv5 🚀 PyTorch Hub models allow for simple model loading and inference in a pure python environment without using detect.py.

Simple Inference Example

This example loads a pretrained YOLOv5s model from PyTorch Hub as model and passes an image for inference. 'yolov5s' is the YOLOv5 'small' model. For details on all available models please see the README. Custom models can also be loaded, including custom trained PyTorch models and their exported variants, i.e. ONNX, TensorRT, TensorFlow, OpenVINO YOLOv5 models.

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, etc.
# model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/best.pt')  # custom trained model

# Images
im = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, URL, PIL, OpenCV, numpy, list

# Inference
results = model(im)

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

results.xyxy[0]  # im predictions (tensor)
results.pandas().xyxy[0]  # im predictions (pandas)
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

See YOLOv5 PyTorch Hub Tutorial for details.

Good luck 🍀 and let us know if you have any other questions!

[ahmadmustafaanis]

How should I limit it to only 1 prediction per image? let's say I have a batch of 8 images. I only want 1 prediction per image in all the images in the batch.

[adrianholovaty]

@ahmadmustafaanis Sounds like you might have a classification problem instead of an object-detection problem?

[ahmadmustafaanis]

@adrianholovaty no it is not a classification problem. This is obj detection problem. I just want to make sure model predicts only 1 bbx per 1 image in the batch. so if the batch has 3 images i.e (3,20,20,3) then the output of the model I want is 3 tensors, each tensor belonging to each image in the batch. Even if there is no detection in that specific image in batch, it should give empty tensor for that batch.