YOLOv5 得到boxes, scores, 类, nums

Question

我正在尝试在我的项目中将对象跟踪与深度排序绑定，我需要获取框、分数、类、nums。

正在加载预训练的 Yolov5 模型：

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

获得预测：

result = model(img)
print(result.shape)
print(result)
torch.Size([8, 6])
tensor([[277.50000, 379.25000, 410.50000, 478.75000,   0.90625,   2.00000],
        [404.00000, 205.12500, 498.50000, 296.00000,   0.88623,   2.00000],
        [262.50000, 247.75000, 359.50000, 350.25000,   0.88281,   2.00000],
        [210.50000, 177.75000, 295.00000, 261.75000,   0.83154,   2.00000],
        [195.50000, 152.50000, 257.75000, 226.00000,   0.78223,   2.00000],
        [137.00000, 146.75000, 168.00000, 162.00000,   0.55713,   2.00000],
        [ 96.00000, 130.12500, 132.50000, 161.12500,   0.54199,   2.00000],
        [ 43.56250,  89.56250,  87.68750, 161.50000,   0.50146,   5.00000]], device='cuda:0')
tensor([[277.50000, 379.25000, 410.50000, 478.75000,   0.90625,   2.00000],
        [404.00000, 205.12500, 498.50000, 296.00000,   0.88623,   2.00000],
        [262.50000, 247.75000, 359.50000, 350.25000,   0.88281,   2.00000],
        [210.50000, 177.75000, 295.00000, 261.75000,   0.83154,   2.00000],
        [195.50000, 152.50000, 257.75000, 226.00000,   0.78223,   2.00000],
        [137.00000, 146.75000, 168.00000, 162.00000,   0.55713,   2.00000],
        [ 96.00000, 130.12500, 132.50000, 161.12500,   0.54199,   2.00000],
        [ 43.56250,  89.56250,  87.68750, 161.50000,   0.50146,   5.00000]], device='cuda:0')

所以现在我的问题是如何获取每个变量中的方框、分数、类、数值？我需要它用于对象跟踪

我用 Pytorch 文档上的例子试了一次： result.xyxy[0]

但在我的案例中我得到了一个错误：

Tensor has no attribute xyxy

Answer 1

模型的输出是火炬张量，没有 xyxy 方法。您需要手动提取值。要么你一个一个过一遍每一个检测：

import torch

det = torch.rand(8, 6)

for *xyxy, conf, cls in det:
    print(*xyxy)
    print(conf)
    print(cls)

或者您可以通过以下方式对检测张量进行切片：

xyxy = det[:, 0:4]
conf = det[:, 4]
cls = det[:, 5]

print(xyxy)
print(conf)
print(cls)

YOLOv5 得到boxes, scores, 类, nums

YOLOv5 get boxes, scores, classes, nums

pytorch

yolov5

roboflow

object-tracking