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!
From what I see you can't load the module named 'model'. Maybe you need to install some dependencies