Search code examples
opencvpytorchyolo

How to get a predicted image of YOLOv5 model?

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.

Solution

  • 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.