在这段基于 pytorch 的对象检测代码中,我哪里出错了?

What and where am I going wrong in this code for pytorch based object detection?

我在这个项目中使用 Yolov5

这是我的代码

import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn
from models.experimental import attempt_load
from utils.general import non_max_suppression

weights = '/Users/nidhi/Desktop/yolov5/best.pt'
device = torch.device('cpu')

model = attempt_load(weights, map_location=device)  # load FP32 model
stride = int(model.stride.max())  # model stride
cudnn.benchmark = True

# Capture with opencv and detect object
cap = cv2.VideoCapture('Pothole testing.mp4')
width, height = (352, 352) # quality 
cap.set(3, width) # width
cap.set(4, height) # height

while(cap.isOpened()):
    time.sleep(0.2) # wait for 0.2 second 
    ret, frame = cap.read()
    if ret ==True:
        now = time.time()
        img = torch.from_numpy(frame).float().to(device).permute(2, 0, 1)
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        pred = model(img, augment=False)[0]
        pred = non_max_suppression(pred, 0.39, 0.45, classes=0, agnostic=True) # img, conf, iou, classes, ...
        print('time -> ', time.time()-now)
     else:
         break

cap.release()

我得到的错误:

  File "run.py", line 38, in <module>
    pred = model(img, augment=False)[0]
  File "/Users/nidhi/Library/Python/3.8/lib/python/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/nidhi/Desktop/yolov5/models/yolo.py", line 118, in forward
    return self.forward_once(x, profile)  # single-scale inference, train
  File "/Users/nidhi/Desktop/yolov5/models/yolo.py", line 134, in forward_once
    x = m(x)  # run
  File "/Users/nidhi/Library/Python/3.8/lib/python/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/nidhi/Desktop/yolov5/models/common.py", line 152, in forward
    return torch.cat(x, self.d)
RuntimeError: Sizes of tensors must match except in dimension 1. Got 108 and 107 in dimension 3 (The offending index is 1)

操作系统:macOS Big Sur 11.2.3

Python版本:3.8.2

模型使用best.pt,我在Google Colab上训练过,我使用yolov5l模型训练数据集。

您是否在以下行中收到错误?

pred = model(img, augment=False)[0]

可能是因为YOLO期望输入的图像大小是32的倍数。所以320×320、352×352等。但是你是352x288。您要么必须调整它的大小,要么用 white/black 像素填充 288 尺寸使其成为 352。

如果您不确定哪里出现错误,可以附上整个错误吗?

在此处获取解决方案 https://www.youtube.com/watch?v=_gQ2Xzld0m4 它对我的工作与 model=yolov5s.pt

的想法完全相同