1

I'm tryting to recognize lego bricks from video cam using opencv. It performs extremely bad comparing with just running detect.py in Yolov5. Thus I made some experiments about just recognizing images, and I found using openCV still performs dramatically bad as well, is there any clue? Here are the experiments I did.

This is the result from detect.py by just running

python detect.py --weights runs/train/yolo/weights/best.pt --source legos.jpg

enter image description here

This is the result from openCV by implementing this

import torch
import cv2
import numpy as np

model = torch.hub.load('.', 'custom', path='runs/train/yolo/weights/last.pt', source='local')

cap = cv2.VideoCapture('legos.jpg')

while cap.isOpened():
    ret, frame = cap.read()

    # Make detections
    results = model(frame)

    cv2.imshow('YOLO', np.squeeze(results.render()))

    if cv2.waitKey(0) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

enter image description here

If I simply do this, it gives a pretty good result

import torch
results = model('legos.jpg')
results.show()

enter image description here

Any genious ideas?

Markus
  • 5,976
  • 5
  • 6
  • 21
Haotian Ma
  • 13
  • 3

1 Answers1

0

Probably your model is trained with RGB images while opencv is using BGR format. Please try to convert the colour space accordingly. Example:

import torch
import cv2

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# read image and convert to RGB
img = cv2.imread('zidane.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# make detections
results = model(img)

# render results and convert back to BGR
results.render()
out = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

cv2.imshow('YOLO', out)

cv2.waitKey(-1)
cv2.destroyAllWindows()
Markus
  • 5,976
  • 5
  • 6
  • 21