yolov5测试
import argparse import time from pathlib import Path import cv2 import torch import torch.backends.cudnn as cudnn from numpy import random import numpy as np from models.experimental import attempt_load from utils.datasets import LoadStreams,LoadStreams2, LoadImages,LoadWebcam,letterbox from utils.general import check_img_size, check_requirements, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path from utils.plots import plot_one_box from utils.torch_utils import select_device, load_classifier, time_synchronized device = select_device('') augment = False conf_thres=0.55 iou_thres=0.45 model = attempt_load('yolov5s.pt', map_location=device) img_size = 640 names = model.module.names if hasattr(model, 'module') else model.names colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] def detectionObjectFunction(): vc = cv2.VideoCapture(2) #rval, frame = vc.read() while True: rval, cameraImg = vc.read() img = letterbox(cameraImg, new_shape=img_size)[0] # Convert img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 img = np.ascontiguousarray(img) #################################################### img = torch.from_numpy(img).to(device) #img = img.half() if half else img.float() # uint8 to fp16/32 im0 = cameraImg.copy() img = img.half() img = img.float() img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference t1 = time_synchronized() pred = model(img, augment=augment)[0] #pred = model(img, augment=opt.augment)[0] #print('thres:%d '%conf_thres) # Apply NMS pred = non_max_suppression(pred, conf_thres, iou_thres) #def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()): t2 = time_synchronized() # Apply Classifier # Process detections for i, det in enumerate(pred): # detections per image # batch_size >= 1 #if webcam: # p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count #else: # p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0) # #p = Path(p) # to Path #save_path = str(save_dir / p.name) # img.jpg #txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt #s += '%gx%g ' % img.shape[2:] # print string # normalization gain whwh #gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] if len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() # Print results for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class #s += f'{n} {names[int(c)]}s, ' # add to string # Write results for *xyxy, conf, cls in reversed(det): label = f'{names[int(cls)]} {conf:.2f}' #plot_one_box2(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=2) #plot_one_box2(xyxy, im0, label=label, color=(0,255,0), line_thickness=2) #plot_one_box(xyxy, im0, label=label, color=(0,255,0), line_thickness=2) plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=2) # Print time (inference + NMS) print(f'detection time. ({t2 - t1:.3f}s)') # Stream results #if view_img: cv2.imshow("win1", im0) #img2 = im0.copy() #################################################### #pass if cv2.waitKey(10) == 27: break detectionObjectFunction()
