from operator import index import sys from numpy.core.fromnumeric import size from numpy.lib.nanfunctions import nanstd sys.path.insert(0, './yolov5') from yolov5.utils.datasets import LoadImages, LoadStreams # from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords, xyxy2xywh from yolov5.utils.plots import Annotator, colors from yolov5.utils.torch_utils import select_device, time_synchronized from deep_sort_pytorch.utils.parser import get_config from deep_sort_pytorch.deep_sort import DeepSort import argparse import os import platform import shutil import time from pathlib import Path import cv2 import torch import torch.backends.cudnn as cudnn from yolov5.utils.Point2GPS import * import ec.countunit from ec.rtsppushserver import RtstPushServer palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) def detect(opt, save_img=False): out, source, weights, view_img, save_txt, imgsz = \ opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size webcam = source == '0' or source.startswith( 'rtsp') or source.startswith('rtmp') or source.startswith('http') or source.endswith('.txt') # initialize deepsort cfg = get_config() cfg.merge_from_file(opt.config_deepsort) deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, use_cuda=True) # Initialize device = select_device(opt.device) if os.path.exists(out): shutil.rmtree(out) # delete output folder os.makedirs(out) # make new output folder half = device.type != 'cpu' # half precision only supported on CUDA # Load model model = torch.load(weights, map_location=device)[ 'model'].float() # load to FP32 model.to(device).eval() if half: model.half() # to FP16 # Set Dataloader vid_path, vid_writer = None, None if webcam: view_img = True cudnn.benchmark = True # set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz) else: view_img = True save_img = True dataset = LoadImages(source, img_size=imgsz) # Get names and colors names = model.module.names if hasattr(model, 'module') else model.names # Run inference t0 = time.time() img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img # run once _ = model(img.half() if half else img) if device.type != 'cpu' else None save_path = str(Path(out)) txt_path = str(Path(out)) + '/results.txt' rtstPushServer = RtstPushServer(9,opt.view_img_size[0],opt.view_img_size[1],'rtsp://192.168.1.182:8554/video') for frame_idx, (path, img, im0s, vid_cap) in enumerate(dataset): draw_box=[] # if(framenumber==0): # continue img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 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=opt.augment)[0] # Apply NMS pred = non_max_suppression( pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) t2 = time_synchronized() # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() else: p, s, im0 = path, '', im0s # 设置显示窗体大小 im0= cv2.resize(im0,opt.view_img_size) s += '%gx%g ' % img.shape[2:] # print string save_path = str(Path(out) / Path(p).name) if det is not None and len(det): 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' * (n > 1)}, " # add to string xywhs = xyxy2xywh(det[:, 0:4]) confs = det[:, 4] clss = det[:, 5] outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0) if outputs is not None and len(outputs) > 0: for j, (output, conf) in enumerate(zip(outputs, confs)): bboxes = output[0:4] id = output[4] cls = output[5] c = int(cls) # integer class # label = f'{names[c]} {conf:.2f}' label = f'{id} {names[c]} {conf:.2f}' color=colors(c, True) draw_box.append([bboxes,label,color]) # draw_boxes(im0, bboxes, id,label) # ori_im = draw_boxes2(im0, bbox_xyxy, identities, clsTracking, trace, h_inv) # Write MOT compliant results to file if save_txt and len(outputs) != 0: for j, output in enumerate(outputs): bbox_left = output[0] bbox_top = output[1] bbox_w = output[2] bbox_h = output[3] identity = output[-1] with open(txt_path, 'a') as f: f.write(('%g ' * 10 + '\n') % (frame_idx, identity, bbox_left, bbox_top, bbox_w, bbox_h, -1, -1, -1, -1)) # label format ec.countunit.addcount(outputs,im0) else: deepsort.increment_ages() # Print time (inference + NMS) # print('%sDone. (%.3fs)' % (s, t2 - t1)) # Stream results if view_img: im0= ec.countunit.addtitlerectangle(im0) annotator = Annotator(im0, line_width=2, pil=True) for bboxes,label,color in draw_box: annotator.box_label(bboxes, label, color) annotator.drawcount( ec.countunit.countlist, ec.countunit.clsnames) im0=annotator.cvimg() im0= ec.countunit.addcolorimg(im0) rtstPushServer.write(im0) cv2.imshow(p, im0) if cv2.waitKey(1) == ord('q'): # q to quit raise StopIteration if save_img: print('saving video!') if vid_path != save_path: # new video vid_path = save_path if isinstance(vid_writer, cv2.VideoWriter): vid_writer.release() # release previous video writer fps =20# vid_cap.get(cv2.CAP_PROP_FPS) # w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) # h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) vid_writer = cv2.VideoWriter( save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps,opt.view_img_size) vid_writer.write(im0) # Save results (image with detections) # if save_img: # print('saving img!') # if dataset.mode == 'images': # cv2.imwrite(save_path, im0) # else: # print('saving video!') # if vid_path != save_path: # new video # vid_path = save_path # if isinstance(vid_writer, cv2.VideoWriter): # vid_writer.release() # release previous video writer # fps = vid_cap.get(cv2.CAP_PROP_FPS) # w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) # h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # vid_writer = cv2.VideoWriter( # save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h)) # vid_writer.write(im0) # if save_txt or save_img: # print('Results saved to %s' % os.getcwd() + os.sep + out) # if platform == 'darwin': # MacOS # os.system('open ' + save_path) # print('Done. (%.3fs)' % (time.time() - t0)) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='yolov5/weights/yolov5s.pt', help='model.pt path') # file/folder, 0 for webcam parser.add_argument('--source', type=str, default='rtsp://admin:hk123456@192.168.1.65:554', help='source') # default='F:/AIData/video/hqxjy.MP4', help='source') # default='rtsp://admin:hk123456@120.240.37.42:554', help='source') parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder parser.add_argument('--img-size', type=int, default=960, help='inference size (pixels)') # conf 相似度 parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold') # 重叠去除 针对可能是 两种东西 parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--view-img', action='store_true',default=True, help='display results') #显示窗体大小--- 注意 countunit.py img_size =(1920,1080)也要设置 parser.add_argument('--view-img-size', type=int, default=(1920,1080), help='inference size (pixels)') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') # class 0 is person 跟踪类别 parser.add_argument('--classes', nargs='+', type=int, default=[0,1,2,3,5,7,16], help='filter by class') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument("--config_deepsort", type=str, default="deep_sort_pytorch/configs/deep_sort.yaml") args = parser.parse_args() args.img_size = check_img_size(args.img_size) print(args) with torch.no_grad(): print(torch.__version__) #注意是双下划线 detect(args)