How to get a predicted image of YOLOv5 model?

Question

I'm trying to load YOLOv5 model and using it to predict specific image. My problem is I want to show predicted image with bounding box into my application so I need to get it directly from the predict method of PyTorch to show in my application.

    model = torch.hub.load('yolov5', 'custom', path=model_name, force_reload=True, 
    source='local')
    pred = model(image)
    pred.show() #show image but can't assign to a variable
    pred.save() #save image to runs\detect\exp

I want something like:

    predict_image = model(image)
    cv2.imshow('Predict', predict_image)

Thank you.

Answer 1

Because the Answer from @mak13 did not work for me and i wanted to understand this behaviour, here is my solution:

Simple get the image as variable

import torch
import pathlib

img_path = pathlib.Path("test_img.jpg")

model = torch.hub.load('ultralytics/yolov5', 'yolov5n')
results = model(img_path)
r_img = results.render() # returns a list with the images as np.array
img_with_boxes = r_img[0] # image with boxes as np.array

Existing infos for this topic at GitHub

Looking into the official Pytorch Hub Wiki from yolov5 in the section Base64 Results we find info on how to use render and also some handy info for yolov5 and REST API's and why this was implemented.

Looking into the code to understand this behaviour

If take a look into the type of results type(results) we get <class 'models.common.Detections'> which we can find here at yolov5's GitHub page.

This class has this (render, show, etc.) methods that run the _run method with different parameters. The render method will cause the _run method to override all images with the detected ones:

if render:
    self.ims[i] = np.asarray(im)