import os import torch from yolov5.models.common import DetectMultiBackend from yolov5.utils.torch_utils import select_device from yolov5.utils.dataloaders import LoadImages from yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_coords, xyxy2xywh from deep_sort.deep_sort import DeepSort class VideoTracker(object): def __init__(self, weights_pt, data, video_path, save_path, idx_to_class): self.video_path = video_path self.save_path = save_path self.idx_to_class = idx_to_class # 选择设备（CPU 或 GPU） device = select_device('cpu') self.vdo = cv2.VideoCapture() self.detector = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False) self.deepsort = DeepSort( model_path="deep_sort/deep/checkpoint/ckpt.t7", # ReID 模型路径 max_dist=0.2, # 外观特征匹配阈值（越小越严格） max_iou_distance=0.7, # 最大IoU距离阈值 max_age=70, # 目标最大存活帧数（未匹配时保留的帧数） n_init=3 # 初始确认帧数（连续匹配到n_init次后确认跟踪） ) self.class_names = self.detector.class_names def __enter__(self): self.vdo.open(self.video_path) self.im_width = int(self.vdo.get(cv2.CAP_PROP_FRAME_WIDTH)) self.im_height = int(self.vdo.get(cv2.CAP_PROP_FRAME_HEIGHT)) assert self.vdo.isOpened() if self.save_path: os.makedirs(self.args.save_path, exist_ok=True) # path of saved video and results self.save_video_path = os.path.join(self.save_path, "results.avi") self.save_results_path = os.path.join(self.save_path, "results.txt") # create video writer fourcc = cv2.VideoWriter_fourcc(*'MJPG') self.writer = cv2.VideoWriter(self.save_video_path, fourcc, 20, (self.im_width, self.im_height)) return self def __exit__(self, exc_type, exc_value, exc_traceback): if exc_type: print(exc_type, exc_value, exc_traceback) def run(self): stride, names, pt = self.model.stride, self.model.names, self.model.pt imgsz = check_img_size((640, 640), s=stride) # check image size dataset = LoadImages(self.video_path, img_size=imgsz, stride=stride, auto=pt, vid_stride=1) bs = len(dataset) self.model.warmup(imgsz=(1 if pt or self.model.triton else bs, 3, *imgsz)) for path, im, im0s, vid_cap, s in dataset: im = torch.from_numpy(im).to(self.model.device) im = im.half() if self.model.fp16 else im.float() # uint8 to fp16/32 im /= 255 # 0 - 255 to 0.0 - 1.0 if len(im.shape) == 3: im = im[None] # expand for batch dim if self.model.xml and im.shape[0] > 1: ims = torch.chunk(im, im.shape[0], 0) # Inference if self.model.xml and im.shape[0] > 1: pred = None for image in ims: if pred is None: pred = self.model(image, augment=False, visualize=False).unsqueeze(0) else: pred = torch.cat( (pred, self.model(image, augment=False, visualize=False).unsqueeze(0)), dim=0 ) pred = [pred, None] else: pred = self.model(im, augment=False, visualize=False) # NMS pred = non_max_suppression(pred, 0.40, 0.45, None, False, max_det=1000)[0] image = im0s[0] pred[:, :4] = scale_coords(im.shape[2:], pred[:, :4], image.shape).round() # 使用YOLOv5进行检测后得到的pred bbox_xywh, cls_conf, cls_ids = yolov5_to_deepsort_format(pred) # select person class mask = cls_ids == 0 bbox_xywh = bbox_xywh[mask] # bbox dilation just in case bbox too small, delete this line if using a better pedestrian detector bbox_xywh[:, 2:] *= 1.2 cls_conf = cls_conf[mask] cls_ids = cls_ids[mask] # 调用Deep SORT更新方法 outputs, _ = self.deepsort.update(bbox_xywh, cls_conf, cls_ids, image) count_result = {} for key in self.idx_to_class.keys(): count_result[key] = set() # draw boxes for visualization if len(outputs) > 0: bbox_xyxy = outputs[:, :4] # 这个是检测所在框的坐标的数组 identities = outputs[:, -1] # 这个是每个元素的计数的数组 cls = outputs[:, -2] # 这个是标签数组id的数组 names = [self.idx_to_class[str(label)] for label in cls] image = draw_boxes(image, bbox_xyxy, names, identities) for i in range(len(cls)): count_result[str(cls[i])].add(identities[i]) def draw_boxes(img, bbox, names=None, identities=None, offset=(0, 0)): for i, box in enumerate(bbox): x1, y1, x2, y2 = [int(i) for i in box] x1 += offset[0] x2 += offset[0] y1 += offset[1] y2 += offset[1] # box text and bar id = int(identities[i]) if identities is not None else 0 color = compute_color_for_labels(id) label = '{:}{:d}'.format(names[i], id) t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2, 2)[0] cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) cv2.rectangle(img, (x1, y1), (x1 + t_size[0] + 3, y1 + t_size[1] + 4), color, -1) cv2.putText(img, label, (x1, y1 + t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 2, [255, 255, 255], 2) return img def compute_color_for_labels(label): """ Simple function that adds fixed color depending on the class """ color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette] return tuple(color) def yolov5_to_deepsort_format(pred): """ 将YOLOv5的预测结果转换为Deep SORT所需的格式 :param pred: YOLOv5的预测结果 :return: 转换后的bbox_xywh, confs, class_ids """ pred[:, :4] = xyxy2xywh(pred[:, :4]) xywh = pred[:, :4].cpu().numpy() conf = pred[:, 4].cpu().numpy() cls = pred[:, 5].cpu().numpy() return xywh, conf, cls if __name__ == "__main__": args = parse_args() cfg = get_config() if args.segment: cfg.USE_SEGMENT = True else: cfg.USE_SEGMENT = False if args.mmdet: cfg.merge_from_file(args.config_mmdetection) cfg.USE_MMDET = True else: cfg.merge_from_file(args.config_detection) cfg.USE_MMDET = False cfg.merge_from_file(args.config_deepsort) if args.fastreid: cfg.merge_from_file(args.config_fastreid) cfg.USE_FASTREID = True else: cfg.USE_FASTREID = False with VideoTracker(cfg, args, video_path=args.VIDEO_PATH) as vdo_trk: vdo_trk.run()