2

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.

Minh Ngôn Đặng
  • 35
  • 1
  • 7
  • `pred = model(image)` -- this does not predict an image, but a list of box candidates. you have to filter those by prob, apply nms, then you can *draw your own image* – berak Apr 25 '22 at 07:28
  • you might need to do something similar to https://stackoverflow.com/questions/60272086/after-finetuning-faster-rcnn-object-detection-model-how-to-visualize-bbox-predi?rq=1 – Jeru Luke Apr 26 '22 at 14:15

2 Answers2

5

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)
perperam
  • 432
  • 3
  • 10
  • This is the correct answer because it seems the resulting image data is not updated instantly and render() function enforces update. – Nikolay Frick Dec 14 '22 at 01:04
2

A quick workaround will be to use "imgs" object from "pred" like the following:

predict_image = model(image)
im_rgb = cv2.cvtColor(predict_image.imgs[0], cv2.COLOR_BGR2RGB) # Because of OpenCV reading images as BGR
cv2_imshow(im_rgb)

I wish this will help you, have a good day.

mak13
  • 146
  • 4