pythonpytorchyolov5roboflowobject-tracking

YOLOv5 get boxes, scores, classes, nums

im trying to bind the Object Tracking with Deep Sort in my Project and i need to get the boxes, scores, classes, nums.

Loading Pretrained Yolov5 model:

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

Getting the Prediction:

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')

so now my question is how do i get the boxes, scores, classes, nums in each variables? I need that for the Object Tracking

I tried it once with the example on Pytorch Documentation: result.xyxy[0]

but in my Case I get an Error:

Tensor has no attribute xyxy
Solution

  • The output from the model is a torch tensor and has no xyxy method. You need to extract the values manually. Either you can go through each detection one by one:

    import torch

det = torch.rand(8, 6)

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

    or you can slice the detections tensor by:

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

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