YOLOv8 获取预测边界框

Question

我想将OpenCV与来自“ultralytics”的YOLOv8集成，因此我想从模型预测中获取边界框坐标。我该如何做到这一点？

from ultralytics import YOLO
import cv2

model = YOLO('yolov8n.pt')
cap = cv2.VideoCapture(0)
cap.set(3, 640)
cap.set(4, 480)

while True:
    _, frame = cap.read()
    
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    results = model.predict(img)

    for r in results:
        for c in r.boxes.cls:
            print(model.names[int(c)])

    cv2.imshow('YOLO V8 Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord(' '):
        break

cap.release()
cv2.destroyAllWindows()

我想在OpenCV中显示YOLO注释图像。我知道我可以在model.predict(source='0', show=True)中使用流参数。但是我想要连续监视预测的类名，同时显示图像输出。

- Louis

Answer 1

这将循环遍历视频中的每一帧，并使用内置的ultralytics注释器绘制其对应的边界框。

from ultralytics import YOLO
import cv2
from ultralytics.utils.plotting import Annotator  # ultralytics.yolo.utils.plotting is deprecated

model = YOLO('yolov8n.pt')
cap = cv2.VideoCapture(0)
cap.set(3, 640)
cap.set(4, 480)

while True:
    _, img = cap.read()
    
    # BGR to RGB conversion is performed under the hood
    # see: https://github.com/ultralytics/ultralytics/issues/2575
    results = model.predict(img)

    for r in results:
        
        annotator = Annotator(img)
        
        boxes = r.boxes
        for box in boxes:
            
            b = box.xyxy[0]  # get box coordinates in (top, left, bottom, right) format
            c = box.cls
            annotator.box_label(b, model.names[int(c)])
          
    img = annotator.result()  
    cv2.imshow('YOLO V8 Detection', img)     
    if cv2.waitKey(1) & 0xFF == ord(' '):
        break

cap.release()
cv2.destroyAllWindows()

- Mike B

---

谢谢..@Mike B，你知道怎样关闭model.predict的输出吗？ - Louis

model.predict(img, verbose=False) @Louis模型.预测(img, verbose=False) @Louis - Mike B

我认为BGR2RGB转换是一个错误，不应该进行，详见https://github.com/ultralytics/ultralytics/issues/2575。 - undefined

还有，r.plot()或results.plot()呢？根据文档，它们应该可以工作，但实际上并没有。 - undefined

我认为BGR2RGB转换是一个错误，不应该进行。谢谢你的澄清。我会更新我的答案！ - undefined

Answer 2

您可以使用以下代码获取所有信息：

for result in results:
    # detection
    result.boxes.xyxy   # box with xyxy format, (N, 4)
    result.boxes.xywh   # box with xywh format, (N, 4)
    result.boxes.xyxyn  # box with xyxy format but normalized, (N, 4)
    result.boxes.xywhn  # box with xywh format but normalized, (N, 4)
    result.boxes.conf   # confidence score, (N, 1)
    result.boxes.cls    # cls, (N, 1)

    # segmentation
    result.masks.masks     # masks, (N, H, W)
    result.masks.segments  # bounding coordinates of masks, List[segment] * N

    # classification
    result.probs     # cls prob, (num_class, )

您可以在文档中进一步阅读相关的IT技术内容。

- Roy Amoyal

Answer 3

这个错误在哪里？

from ultralytics import YOLO
import cv2
img_source = cv2.imread('dog.jpeg')
model = YOLO('yolov8n.pt')
img = cv2.cvtColor(img_source, cv2.COLOR_BGR2RGB)
results = model.predict(source=img)

当源不是路径时，无法获得结果。我如何对通过opencv加载的图像进行预测？

在Windows 10上使用Anaconda。

- julio

Answer 4

请找到一种方法来获取坐标。boxe对象使用torch tensor。可以使用torch.Tensor.tolist来获取坐标。

from ultralytics import YOLO
import cv2

im1 = cv2.imread('/dir/im1.jpg')
im2 = cv2.imread('/dir/im2.jpg')

model = YOLO('yolov8n.pt')
results = model.predict(source=[im1, im2])

fig, axs = plt.subplots(1,2, figsize=(10, 6))
axs = axs.ravel()
plt.subplots_adjust(left=0.1,bottom=0.1, 
                    right=0.9, top=0.9, 
                    wspace=0.2, hspace=0.4)

fig.suptitle("images", fontsize=18, y=0.95)

for i, (r, im) in enumerate(zip(results, images)):

    image = cv2.imread('/dir/' + im)

    c = r.boxes.xywh.tolist()[0] # To get the coordinates.
    x, y, w, h = c[0], c[1], c[2], c[3] # x, y are the center coordinates.
    
    axs[i].imshow(image)
    axs[i].add_patch(Rectangle((x-w/2, y-h/2), w, h,
                     edgecolor='blue', facecolor='none',
                     lw=3))

- Jon