Search code examples
pythonobject-detectionyolov5

ModuleNotFound Error in Yolov5 using custom weights


I'm currently learning object detection using a custom trained module in YoloV5. I trained the dataset locally and that operation ran just fine but when i tried to implement the custom weights it gave me this error:

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

This is the code i'm using to test it out on a live webcam:

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

Am i doing anything wrong? Suggestions are much appreciated!


Solution

  • From what I see you can't load the module named 'model'. Maybe you need to install some dependencies