Yolo output to array ?

[Samuel-Bachorik]

Hello I am using this code for detection on image
`
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
img = 'C:/Users/Samuel/PycharmProjects/Condapytorch/image for yolo/KKT.jpg'

results = model(img)
results.print() # or .show(), .save()
`

I want to ask you is it possible to get this result as array in python ? I mean output image with detected objects as array on tensor.

Or my question is, is there any other way to get array as output ? I want to sent image to yolo and I want it to return output image as array.

Thank you!

[glenn-jocher]

@Samuel-Bachorik yes of course, see the PyTorch Hub Tutorial for examples of returning arrays:

Tutorials

• Train Custom Data  🚀 RECOMMENDED
• Weights & Biases Logging  🌟 NEW
• Multi-GPU Training
• PyTorch Hub  ⭐ NEW
• ONNX and TorchScript Export
• Test-Time Augmentation (TTA)
• Model Ensembling
• Model Pruning/Sparsity
• Hyperparameter Evolution
• Transfer Learning with Frozen Layers  ⭐ NEW
• TensorRT Deployment

Detailed Example

To load YOLOv5 from PyTorch Hub for batched inference with PIL, OpenCV, Numpy or PyTorch inputs, with results plotted, saved and printed to screen:

import cv2
import torch
from PIL import Image

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)  # for file/URI/PIL/cv2/np inputs and NMS

# Images
for f in ['zidane.jpg', 'bus.jpg']:  # download 2 images
    print(f'Downloading {f}...')
    torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/releases/download/v1.0/' + f, f)
img1 = Image.open('zidane.jpg')  # PIL image
img2 = cv2.imread('bus.jpg')[:, :, ::-1]  # OpenCV image (BGR to RGB)
imgs = [img1, img2]  # batched list of images

# Inference
results = model(imgs, size=640)  # includes NMS

# Results
results.print()  
results.save()  # or .show()

# Data
print(results.xyxy[0])  # print img1 predictions (pixels)
#                   x1           y1           x2           y2   confidence        class
# tensor([[7.50637e+02, 4.37279e+01, 1.15887e+03, 7.08682e+02, 8.18137e-01, 0.00000e+00],
#         [9.33597e+01, 2.07387e+02, 1.04737e+03, 7.10224e+02, 5.78011e-01, 0.00000e+00],
#         [4.24503e+02, 4.29092e+02, 5.16300e+02, 7.16425e+02, 5.68713e-01, 2.70000e+01]])

For all inference options see YOLOv5 autoShape() forward method:

yolov5/models/common.py

Lines 182 to 191 in 3551b07

	def forward(self, imgs, size=640, augment=False, profile=False):
	# Inference from various sources. For height=720, width=1280, RGB images example inputs are:
	# filename: imgs = 'data/samples/zidane.jpg'
	# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
	# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(720,1280,3)
	# PIL: = Image.open('image.jpg') # HWC x(720,1280,3)
	# numpy: = np.zeros((720,1280,3)) # HWC
	# torch: = torch.zeros(16,3,720,1280) # BCHW
	# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images

Inference Settings

NMS IoU threshold and confidence threshold are model attributes, and can be modified by:

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

results = model(imgs, size=320)  # custom inference size

[pravastacaraka]

Can I use PyTorch Hub with a custom model?

[glenn-jocher]

@pravastacaraka yes of course. See PyTorch Hub Tutorial for details:

Tutorials

• Train Custom Data  🚀 RECOMMENDED
• Weights & Biases Logging  🌟 NEW
• Multi-GPU Training
• PyTorch Hub  ⭐ NEW
• ONNX and TorchScript Export
• Test-Time Augmentation (TTA)
• Model Ensembling
• Model Pruning/Sparsity
• Hyperparameter Evolution
• Transfer Learning with Frozen Layers  ⭐ NEW
• TensorRT Deployment

[Samuel-Bachorik]

@glenn-jocher Hello thank you for your reply. I understand basics of pytorch that's no problem. But the problem is I need output image as array. When I use results.save() I want this image that is saved as array. Is this possible ? results.save() saves me output image to my PC but I want this image to array.

[glenn-jocher]

@Samuel-Bachorik the results object contains the original images as results.imgs. You can overlay results on the original images with the results.render() method. See detections class for full details:

yolov5/models/common.py

Lines 232 to 247 in 886f1c0

	class Detections:
	# detections class for YOLOv5 inference results
	def __init__(self, imgs, pred, files, names=None):
	super(Detections, self).__init__()
	d = pred[0].device # device
	gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
	self.imgs = imgs # list of images as numpy arrays
	self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
	self.names = names # class names
	self.files = files # image filenames
	self.xyxy = pred # xyxy pixels
	self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
	self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
	self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
	self.n = len(self.pred)