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 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = attempt_load('yolov5s.pt', map_location=device) # model weight here, replace yolov5s.pt stride = int(model.stride.max()) cudnn.benchmark = True names = model.module.names if hasattr(model, 'module') else model.names cap = cv2.VideoCapture(0) # source: replace the 0 for other source. while(cap.isOpened()): ret, frame = cap.read() if ret == True: img = torch.from_numpy(frame) img = img.permute(2, 0, 1 ).float().to(device) 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.28, 0.45) # img, conf, iou for i, det in enumerate(pred): if len(det): for d in det: # d = (x1, y1, x2, y2, conf, cls) x1 = int(d[0].item()) y1 = int(d[1].item()) x2 = int(d[2].item()) y2 = int(d[3].item()) conf = round(d[4].item(), 2) c = int(d[5].item()) frame = cv2.rectangle(frame, (x1, y1), (x2, y2), (255,0,0), 2) # box frame = cv2.putText(frame, names[c], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA) # object name print(f'{x1} {x2} {y1} {y2} {conf} {c}') # print (x1, y1, x2, y2, conf, cls) cv2.imshow('frame',frame) if cv2.waitKey(1) & 0xFF == ord('q'): break else: break cap.release() cv2.destroyAllWindows()