Yolov5 使用自定义权重时出现 ModuleNotFound 错误

Question

我目前正在使用 YoloV5 中的自定义训练模块学习目标检测。我在本地训练了数据集，该操作运行得很好，但是当我尝试实现自定义权重时，它给了我这个错误：

Traceback (most recent call last):
  File "C:\xampp\htdocs\Online-Examination-System\object-detection-test.py", line 5, in <module>
    from models.experimental import attempt_load
ModuleNotFoundError: No module named 'models

这是我用来在实时网络摄像头上测试它的代码：

import cv2
import torch
from torchvision import transforms
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.datasets import letterbox


model = attempt_load("yolov5\runs\train\exp3\weights\best.pt", map_location=torch.device('cpu')).autoshape()

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

conf_threshold = 0.5
iou_threshold = 0.4

class_labels = ['0', '1', '2'] 

cap = cv2.VideoCapture(0)

cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

def detect_objects():
    while True:
        ret, frame = cap.read()

        img = letterbox(frame, new_shape=640)[0]
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = torch.from_numpy(img).unsqueeze(0).to(device).float() / 255.0

       
        pred = model(img)[0]
        pred = non_max_suppression(pred, conf_threshold, iou_threshold)

        for det in pred:
            if det is not None and len(det):
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()

                
                for *xyxy, conf, cls in reversed(det):
                    label = f'{class_labels[int(cls)]} {conf:.2f}'
                    cv2.rectangle(frame, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (255, 0, 0), 2)
                    cv2.putText(frame, label, (int(xyxy[0]), int(xyxy[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

       
        cv2.imshow('Object Detection', frame)

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

    cap.release()
    cv2.destroyAllWindows()

detect_objects()

我做错了什么吗？非常感谢您的建议！

Answer 1

据我所知，您无法加载名为“model”的模块。也许你需要安装一些依赖项