人工智能基于深度学习YOLOv5行人车辆跟踪检测识别计数系统源码.zip

# 注意： ## 本项目使用Yolov5 3.0版本，最新版本5.0请移步： [https://github.com/Sharpiless/yolov5-deepsort](https://github.com/Sharpiless/yolov5-deepsort) ## 注：新版本添加了类别显示功能 # 项目简介： 使用YOLOv5+Deepsort实现车辆行人追踪和计数，代码封装成一个Detector类，更容易嵌入到自己的项目中。 ```python class Detector(baseDet): def __init__(self): super(Detector, self).__init__() self.init_model() self.build_config() def init_model(self): self.weights = 'weights/yolov5m.pt' self.device = '0' if torch.cuda.is_available() else 'cpu' self.device = select_device(self.device) model = attempt_load(self.weights, map_location=self.device) model.to(self.device).eval() model.half() # torch.save(model, 'test.pt') self.m = model self.names = model.module.names if hasattr( model, 'module') else model.names def preprocess(self, img): img0 = img.copy() img = letterbox(img, new_shape=self.img_size)[0] img = img[:, :, ::-1].transpose(2, 0, 1) img = np.ascontiguousarray(img) img = torch.from_numpy(img).to(self.device) img = img.half() # 半精度 img /= 255.0 # 图像归一化 if img.ndimension() == 3: img = img.unsqueeze(0) return img0, img def detect(self, im): im0, img = self.preprocess(im) pred = self.m(img, augment=False)[0] pred = pred.float() pred = non_max_suppression(pred, self.threshold, 0.4) pred_boxes = [] for det in pred: if det is not None and len(det): det[:, :4] = scale_coords( img.shape[2:], det[:, :4], im0.shape).round() for *x, conf, cls_id in det: lbl = self.names[int(cls_id)] if not lbl in ['person', 'car', 'truck']: continue x1, y1 = int(x[0]), int(x[1]) x2, y2 = int(x[2]), int(x[3]) pred_boxes.append( (x1, y1, x2, y2, lbl, conf)) return im, pred_boxes ``` 调用 self.detect 方法返回图像和预测结果 # DeepSort追踪器： ```python deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, 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) ``` 调用 self.update 方法更新追踪结果 # 运行demo： ```bash python demo.py ``` # 训练自己的模型： 训练好后放到 weights 文件夹下 # 调用接口： ## 创建检测器： ```python from AIDetector_pytorch import Detector det = Detector() ``` ## 调用检测接口： ```python func_status = {} func_status['headpose'] = None result = det.feedCap(im, func_status) ``` 其中 im 为 BGR 图像 返回的 result 是字典，result['frame'] 返回可视化后的图像