├── AIDetector_pytorch.py ├── LICENSE ├── README.md ├── deep_sort ├── configs │ └── deep_sort.yaml ├── deep_sort │ ├── README.md │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-37.pyc │ │ ├── deep_sort.cpython-36.pyc │ │ └── deep_sort.cpython-37.pyc │ ├── deep │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── __init__.cpython-37.pyc │ │ │ ├── feature_extractor.cpython-36.pyc │ │ │ ├── feature_extractor.cpython-37.pyc │ │ │ ├── model.cpython-36.pyc │ │ │ └── model.cpython-37.pyc │ │ ├── checkpoint │ │ │ ├── .gitkeep │ │ │ └── ckpt.t7 │ │ ├── evaluate.py │ │ ├── feature_extractor.py │ │ ├── model.py │ │ ├── original_model.py │ │ ├── test.py │ │ ├── train.jpg │ │ └── train.py │ ├── deep_sort.py │ └── sort │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-37.pyc │ │ ├── detection.cpython-36.pyc │ │ ├── detection.cpython-37.pyc │ │ ├── iou_matching.cpython-36.pyc │ │ ├── iou_matching.cpython-37.pyc │ │ ├── kalman_filter.cpython-36.pyc │ │ ├── kalman_filter.cpython-37.pyc │ │ ├── linear_assignment.cpython-36.pyc │ │ ├── linear_assignment.cpython-37.pyc │ │ ├── nn_matching.cpython-36.pyc │ │ ├── nn_matching.cpython-37.pyc │ │ ├── preprocessing.cpython-36.pyc │ │ ├── preprocessing.cpython-37.pyc │ │ ├── track.cpython-36.pyc │ │ ├── track.cpython-37.pyc │ │ ├── tracker.cpython-36.pyc │ │ └── tracker.cpython-37.pyc │ │ ├── detection.py │ │ ├── iou_matching.py │ │ ├── kalman_filter.py │ │ ├── linear_assignment.py │ │ ├── nn_matching.py │ │ ├── preprocessing.py │ │ ├── track.py │ │ └── tracker.py └── utils │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── parser.cpython-36.pyc │ └── parser.cpython-37.pyc │ ├── asserts.py │ ├── draw.py │ ├── evaluation.py │ ├── io.py │ ├── json_logger.py │ ├── log.py │ ├── parser.py │ └── tools.py ├── demo.py ├── models ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── common.cpython-37.pyc │ ├── experimental.cpython-37.pyc │ └── yolo.cpython-37.pyc ├── common.py ├── experimental.py ├── yolo.py └── yolov5m.yaml ├── requirements.txt ├── tracker.py ├── utils ├── BaseDetector.py ├── __init__.py ├── __pycache__ │ ├── BaseDetector.cpython-37.pyc │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── autoanchor.cpython-37.pyc │ ├── datasets.cpython-37.pyc │ ├── general.cpython-36.pyc │ ├── general.cpython-37.pyc │ ├── google_utils.cpython-36.pyc │ ├── google_utils.cpython-37.pyc │ ├── loss.cpython-37.pyc │ ├── metrics.cpython-36.pyc │ ├── metrics.cpython-37.pyc │ ├── plots.cpython-37.pyc │ ├── torch_utils.cpython-36.pyc │ └── torch_utils.cpython-37.pyc ├── activations.py ├── autoanchor.py ├── general.py ├── google_utils.py ├── metrics.py └── torch_utils.py └── weights └── yolov5m.pt /AIDetector_pytorch.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from models.experimental import attempt_load 4 | from utils.general import non_max_suppression, scale_coords, letterbox 5 | from utils.torch_utils import select_device 6 | from utils.BaseDetector import baseDet 7 | 8 | 9 | class Detector(baseDet): 10 | 11 | def __init__(self): 12 | super(Detector, self).__init__() 13 | self.init_model() 14 | self.build_config() 15 | 16 | def init_model(self): 17 | 18 | self.weights = 'weights/yolov5m.pt' 19 | self.device = '0' if torch.cuda.is_available() else 'cpu' 20 | self.device = select_device(self.device) 21 | model = attempt_load(self.weights, map_location=self.device) 22 | model.to(self.device).eval() 23 | model.half() 24 | # torch.save(model, 'test.pt') 25 | self.m = model 26 | self.names = model.module.names if hasattr( 27 | model, 'module') else model.names 28 | 29 | def preprocess(self, img): 30 | 31 | img0 = img.copy() 32 | img = letterbox(img, new_shape=self.img_size)[0] 33 | img = img[:, :, ::-1].transpose(2, 0, 1) 34 | img = np.ascontiguousarray(img) 35 | img = torch.from_numpy(img).to(self.device) 36 | img = img.half() # 半精度 37 | img /= 255.0 # 图像归一化 38 | if img.ndimension() == 3: 39 | img = img.unsqueeze(0) 40 | 41 | return img0, img 42 | 43 | def detect(self, im): 44 | 45 | im0, img = self.preprocess(im) 46 | 47 | pred = self.m(img, augment=False)[0] 48 | pred = pred.float() 49 | pred = non_max_suppression(pred, self.threshold, 0.4) 50 | 51 | pred_boxes = [] 52 | for det in pred: 53 | 54 | if det is not None and len(det): 55 | det[:, :4] = scale_coords( 56 | img.shape[2:], det[:, :4], im0.shape).round() 57 | 58 | for *x, conf, cls_id in det: 59 | lbl = self.names[int(cls_id)] 60 | if not lbl in ['person', 'car', 'truck']: 61 | continue 62 | x1, y1 = int(x[0]), int(x[1]) 63 | x2, y2 = int(x[2]), int(x[3]) 64 | pred_boxes.append( 65 | (x1, y1, x2, y2, lbl, conf)) 66 | 67 | return im, pred_boxes 68 | 69 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # **YOLOv5 + DeepSort 用于目标跟踪与计数** 3 | 🚗🚶‍♂️ **使用 YOLOv5 和 DeepSort 实现车辆与行人实时跟踪与计数** 4 | 5 | [![GitHub stars](https://img.shields.io/github/stars/Sharpiless/Yolov5-deepsort-inference?style=social)](https://github.com/Sharpiless/Yolov5-deepsort-inference) [![GitHub forks](https://img.shields.io/github/forks/Sharpiless/Yolov5-deepsort-inference?style=social)](https://github.com/Sharpiless/Yolov5-deepsort-inference) [![License](https://img.shields.io/github/license/Sharpiless/Yolov5-deepsort-inference)](https://github.com/Sharpiless/Yolov5-deepsort-inference/blob/main/LICENSE) 6 | 7 | 最新版本：[https://github.com/Sharpiless/YOLOv11-DeepSort](https://github.com/Sharpiless/YOLOv11-DeepSort) 8 | 9 | --- 10 | 11 | ## **📌 项目简介** 12 | 13 | 本项目将 **YOLOv5** 与 **DeepSort** 相结合，实现了对目标的实时跟踪与计数。提供了一个封装的 `Detector` 类，方便将此功能嵌入到自定义项目中。 14 | 15 | 🔗 **阅读完整博客**：[【小白CV教程】YOLOv5+Deepsort实现车辆行人的检测、追踪和计数](https://blog.csdn.net/weixin_44936889/article/details/112002152) 16 | 17 | --- 18 | 19 | ## **🚀 核心功能** 20 | 21 | - **目标跟踪**：实时跟踪车辆与行人。 22 | - **计数功能**：轻松统计视频流中的车辆或行人数。 23 | - **封装式接口**：`Detector` 类封装了检测与跟踪逻辑，便于集成。 24 | - **高度自定义**：支持训练自己的 YOLOv5 模型并无缝接入框架。 25 | 26 | --- 27 | 28 | ## **🔧 使用说明** 29 | 30 | ### **安装依赖** 31 | ```bash 32 | pip install -r requirements.txt 33 | ``` 34 | 35 | 确保安装了 `requirements.txt` 文件中列出的所有依赖。 36 | ### **运行 Demo** 37 | ```bash 38 | python demo.py 39 | ``` 40 | --- 41 | 42 | ## **🛠️ 开发说明** 43 | 44 | ### **YOLOv5 检测器** 45 | 46 | ```python 47 | class Detector(baseDet): 48 | 49 | def __init__(self): 50 | super(Detector, self).__init__() 51 | self.init_model() 52 | self.build_config() 53 | 54 | def init_model(self): 55 | 56 | self.weights = 'weights/yolov5m.pt' 57 | self.device = '0' if torch.cuda.is_available() else 'cpu' 58 | self.device = select_device(self.device) 59 | model = attempt_load(self.weights, map_location=self.device) 60 | model.to(self.device).eval() 61 | model.half() 62 | # torch.save(model, 'test.pt') 63 | self.m = model 64 | self.names = model.module.names if hasattr( 65 | model, 'module') else model.names 66 | 67 | def preprocess(self, img): 68 | 69 | img0 = img.copy() 70 | img = letterbox(img, new_shape=self.img_size)[0] 71 | img = img[:, :, ::-1].transpose(2, 0, 1) 72 | img = np.ascontiguousarray(img) 73 | img = torch.from_numpy(img).to(self.device) 74 | img = img.half() # 半精度 75 | img /= 255.0 # 图像归一化 76 | if img.ndimension() == 3: 77 | img = img.unsqueeze(0) 78 | 79 | return img0, img 80 | 81 | def detect(self, im): 82 | 83 | im0, img = self.preprocess(im) 84 | 85 | pred = self.m(img, augment=False)[0] 86 | pred = pred.float() 87 | pred = non_max_suppression(pred, self.threshold, 0.4) 88 | 89 | pred_boxes = [] 90 | for det in pred: 91 | 92 | if det is not None and len(det): 93 | det[:, :4] = scale_coords( 94 | img.shape[2:], det[:, :4], im0.shape).round() 95 | 96 | for *x, conf, cls_id in det: 97 | lbl = self.names[int(cls_id)] 98 | if not lbl in ['person', 'car', 'truck']: 99 | continue 100 | x1, y1 = int(x[0]), int(x[1]) 101 | x2, y2 = int(x[2]), int(x[3]) 102 | pred_boxes.append( 103 | (x1, y1, x2, y2, lbl, conf)) 104 | 105 | return im, pred_boxes 106 | ``` 107 | - 调用 `self.detect()` 方法返回图像和预测结果 108 | ### **DeepSort 追踪器** 109 | 110 | ```python 111 | deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, 112 | max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, 113 | nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, 114 | max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, 115 | use_cuda=True) 116 | ``` 117 | - 调用 `self.update()` 方法更新追踪结果 118 | --- 119 | 120 | ## **📊 训练自己的模型** 121 | 122 | 如果需要训练自定义的 YOLOv5 模型，请参考以下教程： 123 | [【小白CV】手把手教你用YOLOv5训练自己的数据集（从Windows环境配置到模型部署）](https://blog.csdn.net/weixin_44936889/article/details/110661862) 124 | 125 | 训练完成后，将模型权重文件放置于 `weights` 文件夹中。 126 | 127 | --- 128 | 129 | ## **📦 API 调用** 130 | 131 | ### **初始化检测器** 132 | ```python 133 | from AIDetector_pytorch import Detector 134 | 135 | det = Detector() 136 | ``` 137 | 138 | ### **调用检测接口** 139 | ```python 140 | func_status = {} 141 | func_status['headpose'] = None 142 | 143 | result = det.feedCap(im, func_status) 144 | ``` 145 | 146 | - `im`: 输入的 BGR 图像。 147 | - `result['frame']`: 检测结果的可视化图像。 148 | 149 | --- 150 | 151 | ## **✨ 可视化效果** 152 | 153 | ![效果图](https://img-blog.csdnimg.cn/20201231090541223.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl80NDkzNjg4OQ==,size_16,color_FFFFFF,t_70) 154 | 155 | --- 156 | 157 | ## **📚 联系作者** 158 | - Bilibili: [https://space.bilibili.com/470550823](https://space.bilibili.com/470550823) 159 | - CSDN: [https://blog.csdn.net/weixin_44936889](https://blog.csdn.net/weixin_44936889) 160 | - AI Studio: [https://aistudio.baidu.com/aistudio/personalcenter/thirdview/67156](https://aistudio.baidu.com/aistudio/personalcenter/thirdview/67156) 161 | - GitHub: [https://github.com/Sharpiless](https://github.com/Sharpiless) 162 | 163 | --- 164 | 165 | 166 | 172 | 178 | Star History Chart 182 | 183 | 184 | ## **💡 许可证** 185 | 186 | 本项目遵循 **GNU General Public License v3.0** 协议。 187 | **标明目标检测部分来源**：[https://github.com/ultralytics/yolov5](https://github.com/ultralytics/yolov5) 188 | -------------------------------------------------------------------------------- /deep_sort/configs/deep_sort.yaml: -------------------------------------------------------------------------------- 1 | DEEPSORT: 2 | REID_CKPT: "deep_sort/deep_sort/deep/checkpoint/ckpt.t7" 3 | MAX_DIST: 0.2 4 | MIN_CONFIDENCE: 0.3 5 | NMS_MAX_OVERLAP: 0.5 6 | MAX_IOU_DISTANCE: 0.7 7 | MAX_AGE: 70 8 | N_INIT: 3 9 | NN_BUDGET: 100 10 | 11 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/README.md: -------------------------------------------------------------------------------- 1 | # Deep Sort 2 | 3 | This is the implemention of deep sort with pytorch. -------------------------------------------------------------------------------- /deep_sort/deep_sort/__init__.py: -------------------------------------------------------------------------------- 1 | from .deep_sort import DeepSort 2 | 3 | 4 | __all__ = ['DeepSort', 'build_tracker'] 5 | 6 | 7 | def build_tracker(cfg, use_cuda): 8 | return DeepSort(cfg.DEEPSORT.REID_CKPT, 9 | max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, 10 | nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, 11 | max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, use_cuda=use_cuda) 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/__pycache__/deep_sort.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/__pycache__/deep_sort.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/__pycache__/deep_sort.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/__pycache__/deep_sort.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__init__.py -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__pycache__/feature_extractor.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__pycache__/feature_extractor.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__pycache__/feature_extractor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__pycache__/feature_extractor.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__pycache__/model.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__pycache__/model.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/__pycache__/model.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/__pycache__/model.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/checkpoint/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/checkpoint/.gitkeep -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/checkpoint/ckpt.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/checkpoint/ckpt.t7 -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/evaluate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | features = torch.load("features.pth") 4 | qf = features["qf"] 5 | ql = features["ql"] 6 | gf = features["gf"] 7 | gl = features["gl"] 8 | 9 | scores = qf.mm(gf.t()) 10 | res = scores.topk(5, dim=1)[1][:,0] 11 | top1correct = gl[res].eq(ql).sum().item() 12 | 13 | print("Acc top1:{:.3f}".format(top1correct/ql.size(0))) 14 | 15 | 16 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/feature_extractor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision.transforms as transforms 3 | import numpy as np 4 | import cv2 5 | import logging 6 | 7 | from .model import Net 8 | 9 | class Extractor(object): 10 | def __init__(self, model_path, use_cuda=True): 11 | self.net = Net(reid=True) 12 | self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu" 13 | state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['net_dict'] 14 | self.net.load_state_dict(state_dict) 15 | logger = logging.getLogger("root.tracker") 16 | logger.info("Loading weights from {}... Done!".format(model_path)) 17 | self.net.to(self.device) 18 | self.size = (64, 128) 19 | self.norm = transforms.Compose([ 20 | transforms.ToTensor(), 21 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), 22 | ]) 23 | 24 | 25 | 26 | def _preprocess(self, im_crops): 27 | """ 28 | TODO: 29 | 1. to float with scale from 0 to 1 30 | 2. resize to (64, 128) as Market1501 dataset did 31 | 3. concatenate to a numpy array 32 | 3. to torch Tensor 33 | 4. normalize 34 | """ 35 | def _resize(im, size): 36 | return cv2.resize(im.astype(np.float32)/255., size) 37 | 38 | im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float() 39 | return im_batch 40 | 41 | 42 | def __call__(self, im_crops): 43 | im_batch = self._preprocess(im_crops) 44 | with torch.no_grad(): 45 | im_batch = im_batch.to(self.device) 46 | features = self.net(im_batch) 47 | return features.cpu().numpy() 48 | 49 | 50 | if __name__ == '__main__': 51 | img = cv2.imread("demo.jpg")[:,:,(2,1,0)] 52 | extr = Extractor("checkpoint/ckpt.t7") 53 | feature = extr(img) 54 | print(feature.shape) 55 | 56 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class BasicBlock(nn.Module): 6 | def __init__(self, c_in, c_out,is_downsample=False): 7 | super(BasicBlock,self).__init__() 8 | self.is_downsample = is_downsample 9 | if is_downsample: 10 | self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=2, padding=1, bias=False) 11 | else: 12 | self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=1, padding=1, bias=False) 13 | self.bn1 = nn.BatchNorm2d(c_out) 14 | self.relu = nn.ReLU(True) 15 | self.conv2 = nn.Conv2d(c_out,c_out,3,stride=1,padding=1, bias=False) 16 | self.bn2 = nn.BatchNorm2d(c_out) 17 | if is_downsample: 18 | self.downsample = nn.Sequential( 19 | nn.Conv2d(c_in, c_out, 1, stride=2, bias=False), 20 | nn.BatchNorm2d(c_out) 21 | ) 22 | elif c_in != c_out: 23 | self.downsample = nn.Sequential( 24 | nn.Conv2d(c_in, c_out, 1, stride=1, bias=False), 25 | nn.BatchNorm2d(c_out) 26 | ) 27 | self.is_downsample = True 28 | 29 | def forward(self,x): 30 | y = self.conv1(x) 31 | y = self.bn1(y) 32 | y = self.relu(y) 33 | y = self.conv2(y) 34 | y = self.bn2(y) 35 | if self.is_downsample: 36 | x = self.downsample(x) 37 | return F.relu(x.add(y),True) 38 | 39 | def make_layers(c_in,c_out,repeat_times, is_downsample=False): 40 | blocks = [] 41 | for i in range(repeat_times): 42 | if i ==0: 43 | blocks += [BasicBlock(c_in,c_out, is_downsample=is_downsample),] 44 | else: 45 | blocks += [BasicBlock(c_out,c_out),] 46 | return nn.Sequential(*blocks) 47 | 48 | class Net(nn.Module): 49 | def __init__(self, num_classes=751 ,reid=False): 50 | super(Net,self).__init__() 51 | # 3 128 64 52 | self.conv = nn.Sequential( 53 | nn.Conv2d(3,64,3,stride=1,padding=1), 54 | nn.BatchNorm2d(64), 55 | nn.ReLU(inplace=True), 56 | # nn.Conv2d(32,32,3,stride=1,padding=1), 57 | # nn.BatchNorm2d(32), 58 | # nn.ReLU(inplace=True), 59 | nn.MaxPool2d(3,2,padding=1), 60 | ) 61 | # 32 64 32 62 | self.layer1 = make_layers(64,64,2,False) 63 | # 32 64 32 64 | self.layer2 = make_layers(64,128,2,True) 65 | # 64 32 16 66 | self.layer3 = make_layers(128,256,2,True) 67 | # 128 16 8 68 | self.layer4 = make_layers(256,512,2,True) 69 | # 256 8 4 70 | self.avgpool = nn.AvgPool2d((8,4),1) 71 | # 256 1 1 72 | self.reid = reid 73 | self.classifier = nn.Sequential( 74 | nn.Linear(512, 256), 75 | nn.BatchNorm1d(256), 76 | nn.ReLU(inplace=True), 77 | nn.Dropout(), 78 | nn.Linear(256, num_classes), 79 | ) 80 | 81 | def forward(self, x): 82 | x = self.conv(x) 83 | x = self.layer1(x) 84 | x = self.layer2(x) 85 | x = self.layer3(x) 86 | x = self.layer4(x) 87 | x = self.avgpool(x) 88 | x = x.view(x.size(0),-1) 89 | # B x 128 90 | if self.reid: 91 | x = x.div(x.norm(p=2,dim=1,keepdim=True)) 92 | return x 93 | # classifier 94 | x = self.classifier(x) 95 | return x 96 | 97 | 98 | if __name__ == '__main__': 99 | net = Net() 100 | x = torch.randn(4,3,128,64) 101 | y = net(x) 102 | import ipdb; ipdb.set_trace() 103 | 104 | 105 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/original_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class BasicBlock(nn.Module): 6 | def __init__(self, c_in, c_out,is_downsample=False): 7 | super(BasicBlock,self).__init__() 8 | self.is_downsample = is_downsample 9 | if is_downsample: 10 | self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=2, padding=1, bias=False) 11 | else: 12 | self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=1, padding=1, bias=False) 13 | self.bn1 = nn.BatchNorm2d(c_out) 14 | self.relu = nn.ReLU(True) 15 | self.conv2 = nn.Conv2d(c_out,c_out,3,stride=1,padding=1, bias=False) 16 | self.bn2 = nn.BatchNorm2d(c_out) 17 | if is_downsample: 18 | self.downsample = nn.Sequential( 19 | nn.Conv2d(c_in, c_out, 1, stride=2, bias=False), 20 | nn.BatchNorm2d(c_out) 21 | ) 22 | elif c_in != c_out: 23 | self.downsample = nn.Sequential( 24 | nn.Conv2d(c_in, c_out, 1, stride=1, bias=False), 25 | nn.BatchNorm2d(c_out) 26 | ) 27 | self.is_downsample = True 28 | 29 | def forward(self,x): 30 | y = self.conv1(x) 31 | y = self.bn1(y) 32 | y = self.relu(y) 33 | y = self.conv2(y) 34 | y = self.bn2(y) 35 | if self.is_downsample: 36 | x = self.downsample(x) 37 | return F.relu(x.add(y),True) 38 | 39 | def make_layers(c_in,c_out,repeat_times, is_downsample=False): 40 | blocks = [] 41 | for i in range(repeat_times): 42 | if i ==0: 43 | blocks += [BasicBlock(c_in,c_out, is_downsample=is_downsample),] 44 | else: 45 | blocks += [BasicBlock(c_out,c_out),] 46 | return nn.Sequential(*blocks) 47 | 48 | class Net(nn.Module): 49 | def __init__(self, num_classes=625 ,reid=False): 50 | super(Net,self).__init__() 51 | # 3 128 64 52 | self.conv = nn.Sequential( 53 | nn.Conv2d(3,32,3,stride=1,padding=1), 54 | nn.BatchNorm2d(32), 55 | nn.ELU(inplace=True), 56 | nn.Conv2d(32,32,3,stride=1,padding=1), 57 | nn.BatchNorm2d(32), 58 | nn.ELU(inplace=True), 59 | nn.MaxPool2d(3,2,padding=1), 60 | ) 61 | # 32 64 32 62 | self.layer1 = make_layers(32,32,2,False) 63 | # 32 64 32 64 | self.layer2 = make_layers(32,64,2,True) 65 | # 64 32 16 66 | self.layer3 = make_layers(64,128,2,True) 67 | # 128 16 8 68 | self.dense = nn.Sequential( 69 | nn.Dropout(p=0.6), 70 | nn.Linear(128*16*8, 128), 71 | nn.BatchNorm1d(128), 72 | nn.ELU(inplace=True) 73 | ) 74 | # 256 1 1 75 | self.reid = reid 76 | self.batch_norm = nn.BatchNorm1d(128) 77 | self.classifier = nn.Sequential( 78 | nn.Linear(128, num_classes), 79 | ) 80 | 81 | def forward(self, x): 82 | x = self.conv(x) 83 | x = self.layer1(x) 84 | x = self.layer2(x) 85 | x = self.layer3(x) 86 | 87 | x = x.view(x.size(0),-1) 88 | if self.reid: 89 | x = self.dense[0](x) 90 | x = self.dense[1](x) 91 | x = x.div(x.norm(p=2,dim=1,keepdim=True)) 92 | return x 93 | x = self.dense(x) 94 | # B x 128 95 | # classifier 96 | x = self.classifier(x) 97 | return x 98 | 99 | 100 | if __name__ == '__main__': 101 | net = Net(reid=True) 102 | x = torch.randn(4,3,128,64) 103 | y = net(x) 104 | import ipdb; ipdb.set_trace() 105 | 106 | 107 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.backends.cudnn as cudnn 3 | import torchvision 4 | 5 | import argparse 6 | import os 7 | 8 | from model import Net 9 | 10 | parser = argparse.ArgumentParser(description="Train on market1501") 11 | parser.add_argument("--data-dir",default='data',type=str) 12 | parser.add_argument("--no-cuda",action="store_true") 13 | parser.add_argument("--gpu-id",default=0,type=int) 14 | args = parser.parse_args() 15 | 16 | # device 17 | device = "cuda:{}".format(args.gpu_id) if torch.cuda.is_available() and not args.no_cuda else "cpu" 18 | if torch.cuda.is_available() and not args.no_cuda: 19 | cudnn.benchmark = True 20 | 21 | # data loader 22 | root = args.data_dir 23 | query_dir = os.path.join(root,"query") 24 | gallery_dir = os.path.join(root,"gallery") 25 | transform = torchvision.transforms.Compose([ 26 | torchvision.transforms.Resize((128,64)), 27 | torchvision.transforms.ToTensor(), 28 | torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 29 | ]) 30 | queryloader = torch.utils.data.DataLoader( 31 | torchvision.datasets.ImageFolder(query_dir, transform=transform), 32 | batch_size=64, shuffle=False 33 | ) 34 | galleryloader = torch.utils.data.DataLoader( 35 | torchvision.datasets.ImageFolder(gallery_dir, transform=transform), 36 | batch_size=64, shuffle=False 37 | ) 38 | 39 | # net definition 40 | net = Net(reid=True) 41 | assert os.path.isfile("./checkpoint/ckpt.t7"), "Error: no checkpoint file found!" 42 | print('Loading from checkpoint/ckpt.t7') 43 | checkpoint = torch.load("./checkpoint/ckpt.t7") 44 | net_dict = checkpoint['net_dict'] 45 | net.load_state_dict(net_dict, strict=False) 46 | net.eval() 47 | net.to(device) 48 | 49 | # compute features 50 | query_features = torch.tensor([]).float() 51 | query_labels = torch.tensor([]).long() 52 | gallery_features = torch.tensor([]).float() 53 | gallery_labels = torch.tensor([]).long() 54 | 55 | with torch.no_grad(): 56 | for idx,(inputs,labels) in enumerate(queryloader): 57 | inputs = inputs.to(device) 58 | features = net(inputs).cpu() 59 | query_features = torch.cat((query_features, features), dim=0) 60 | query_labels = torch.cat((query_labels, labels)) 61 | 62 | for idx,(inputs,labels) in enumerate(galleryloader): 63 | inputs = inputs.to(device) 64 | features = net(inputs).cpu() 65 | gallery_features = torch.cat((gallery_features, features), dim=0) 66 | gallery_labels = torch.cat((gallery_labels, labels)) 67 | 68 | gallery_labels -= 2 69 | 70 | # save features 71 | features = { 72 | "qf": query_features, 73 | "ql": query_labels, 74 | "gf": gallery_features, 75 | "gl": gallery_labels 76 | } 77 | torch.save(features,"features.pth") -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/train.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/deep/train.jpg -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | import torch 8 | import torch.backends.cudnn as cudnn 9 | import torchvision 10 | 11 | from model import Net 12 | 13 | parser = argparse.ArgumentParser(description="Train on market1501") 14 | parser.add_argument("--data-dir",default='data',type=str) 15 | parser.add_argument("--no-cuda",action="store_true") 16 | parser.add_argument("--gpu-id",default=0,type=int) 17 | parser.add_argument("--lr",default=0.1, type=float) 18 | parser.add_argument("--interval",'-i',default=20,type=int) 19 | parser.add_argument('--resume', '-r',action='store_true') 20 | args = parser.parse_args() 21 | 22 | # device 23 | device = "cuda:{}".format(args.gpu_id) if torch.cuda.is_available() and not args.no_cuda else "cpu" 24 | if torch.cuda.is_available() and not args.no_cuda: 25 | cudnn.benchmark = True 26 | 27 | # data loading 28 | root = args.data_dir 29 | train_dir = os.path.join(root,"train") 30 | test_dir = os.path.join(root,"test") 31 | transform_train = torchvision.transforms.Compose([ 32 | torchvision.transforms.RandomCrop((128,64),padding=4), 33 | torchvision.transforms.RandomHorizontalFlip(), 34 | torchvision.transforms.ToTensor(), 35 | torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 36 | ]) 37 | transform_test = torchvision.transforms.Compose([ 38 | torchvision.transforms.Resize((128,64)), 39 | torchvision.transforms.ToTensor(), 40 | torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 41 | ]) 42 | trainloader = torch.utils.data.DataLoader( 43 | torchvision.datasets.ImageFolder(train_dir, transform=transform_train), 44 | batch_size=64,shuffle=True 45 | ) 46 | testloader = torch.utils.data.DataLoader( 47 | torchvision.datasets.ImageFolder(test_dir, transform=transform_test), 48 | batch_size=64,shuffle=True 49 | ) 50 | num_classes = max(len(trainloader.dataset.classes), len(testloader.dataset.classes)) 51 | 52 | # net definition 53 | start_epoch = 0 54 | net = Net(num_classes=num_classes) 55 | if args.resume: 56 | assert os.path.isfile("./checkpoint/ckpt.t7"), "Error: no checkpoint file found!" 57 | print('Loading from checkpoint/ckpt.t7') 58 | checkpoint = torch.load("./checkpoint/ckpt.t7") 59 | # import ipdb; ipdb.set_trace() 60 | net_dict = checkpoint['net_dict'] 61 | net.load_state_dict(net_dict) 62 | best_acc = checkpoint['acc'] 63 | start_epoch = checkpoint['epoch'] 64 | net.to(device) 65 | 66 | # loss and optimizer 67 | criterion = torch.nn.CrossEntropyLoss() 68 | optimizer = torch.optim.SGD(net.parameters(), args.lr, momentum=0.9, weight_decay=5e-4) 69 | best_acc = 0. 70 | 71 | # train function for each epoch 72 | def train(epoch): 73 | print("\nEpoch : %d"%(epoch+1)) 74 | net.train() 75 | training_loss = 0. 76 | train_loss = 0. 77 | correct = 0 78 | total = 0 79 | interval = args.interval 80 | start = time.time() 81 | for idx, (inputs, labels) in enumerate(trainloader): 82 | # forward 83 | inputs,labels = inputs.to(device),labels.to(device) 84 | outputs = net(inputs) 85 | loss = criterion(outputs, labels) 86 | 87 | # backward 88 | optimizer.zero_grad() 89 | loss.backward() 90 | optimizer.step() 91 | 92 | # accumurating 93 | training_loss += loss.item() 94 | train_loss += loss.item() 95 | correct += outputs.max(dim=1)[1].eq(labels).sum().item() 96 | total += labels.size(0) 97 | 98 | # print 99 | if (idx+1)%interval == 0: 100 | end = time.time() 101 | print("[progress:{:.1f}%]time:{:.2f}s Loss:{:.5f} Correct:{}/{} Acc:{:.3f}%".format( 102 | 100.*(idx+1)/len(trainloader), end-start, training_loss/interval, correct, total, 100.*correct/total 103 | )) 104 | training_loss = 0. 105 | start = time.time() 106 | 107 | return train_loss/len(trainloader), 1.- correct/total 108 | 109 | def test(epoch): 110 | global best_acc 111 | net.eval() 112 | test_loss = 0. 113 | correct = 0 114 | total = 0 115 | start = time.time() 116 | with torch.no_grad(): 117 | for idx, (inputs, labels) in enumerate(testloader): 118 | inputs, labels = inputs.to(device), labels.to(device) 119 | outputs = net(inputs) 120 | loss = criterion(outputs, labels) 121 | 122 | test_loss += loss.item() 123 | correct += outputs.max(dim=1)[1].eq(labels).sum().item() 124 | total += labels.size(0) 125 | 126 | print("Testing ...") 127 | end = time.time() 128 | print("[progress:{:.1f}%]time:{:.2f}s Loss:{:.5f} Correct:{}/{} Acc:{:.3f}%".format( 129 | 100.*(idx+1)/len(testloader), end-start, test_loss/len(testloader), correct, total, 100.*correct/total 130 | )) 131 | 132 | # saving checkpoint 133 | acc = 100.*correct/total 134 | if acc > best_acc: 135 | best_acc = acc 136 | print("Saving parameters to checkpoint/ckpt.t7") 137 | checkpoint = { 138 | 'net_dict':net.state_dict(), 139 | 'acc':acc, 140 | 'epoch':epoch, 141 | } 142 | if not os.path.isdir('checkpoint'): 143 | os.mkdir('checkpoint') 144 | torch.save(checkpoint, './checkpoint/ckpt.t7') 145 | 146 | return test_loss/len(testloader), 1.- correct/total 147 | 148 | # plot figure 149 | x_epoch = [] 150 | record = {'train_loss':[], 'train_err':[], 'test_loss':[], 'test_err':[]} 151 | fig = plt.figure() 152 | ax0 = fig.add_subplot(121, title="loss") 153 | ax1 = fig.add_subplot(122, title="top1err") 154 | def draw_curve(epoch, train_loss, train_err, test_loss, test_err): 155 | global record 156 | record['train_loss'].append(train_loss) 157 | record['train_err'].append(train_err) 158 | record['test_loss'].append(test_loss) 159 | record['test_err'].append(test_err) 160 | 161 | x_epoch.append(epoch) 162 | ax0.plot(x_epoch, record['train_loss'], 'bo-', label='train') 163 | ax0.plot(x_epoch, record['test_loss'], 'ro-', label='val') 164 | ax1.plot(x_epoch, record['train_err'], 'bo-', label='train') 165 | ax1.plot(x_epoch, record['test_err'], 'ro-', label='val') 166 | if epoch == 0: 167 | ax0.legend() 168 | ax1.legend() 169 | fig.savefig("train.jpg") 170 | 171 | # lr decay 172 | def lr_decay(): 173 | global optimizer 174 | for params in optimizer.param_groups: 175 | params['lr'] *= 0.1 176 | lr = params['lr'] 177 | print("Learning rate adjusted to {}".format(lr)) 178 | 179 | def main(): 180 | for epoch in range(start_epoch, start_epoch+40): 181 | train_loss, train_err = train(epoch) 182 | test_loss, test_err = test(epoch) 183 | draw_curve(epoch, train_loss, train_err, test_loss, test_err) 184 | if (epoch+1)%20==0: 185 | lr_decay() 186 | 187 | 188 | if __name__ == '__main__': 189 | main() 190 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/deep_sort.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from .deep.feature_extractor import Extractor 5 | from .sort.nn_matching import NearestNeighborDistanceMetric 6 | from .sort.preprocessing import non_max_suppression 7 | from .sort.detection import Detection 8 | from .sort.tracker import Tracker 9 | 10 | 11 | __all__ = ['DeepSort'] 12 | 13 | 14 | class DeepSort(object): 15 | def __init__(self, model_path, max_dist=0.2, min_confidence=0.3, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=70, n_init=3, nn_budget=100, use_cuda=True): 16 | self.min_confidence = min_confidence 17 | self.nms_max_overlap = nms_max_overlap 18 | 19 | self.extractor = Extractor(model_path, use_cuda=use_cuda) 20 | 21 | max_cosine_distance = max_dist 22 | nn_budget = 100 23 | metric = NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget) 24 | self.tracker = Tracker(metric, max_iou_distance=max_iou_distance, max_age=max_age, n_init=n_init) 25 | 26 | def update(self, bbox_xywh, confidences, ori_img): 27 | self.height, self.width = ori_img.shape[:2] 28 | # generate detections 29 | features = self._get_features(bbox_xywh, ori_img) 30 | bbox_tlwh = self._xywh_to_tlwh(bbox_xywh) 31 | detections = [Detection(bbox_tlwh[i], conf, features[i]) for i,conf in enumerate(confidences) if conf>self.min_confidence] 32 | 33 | # run on non-maximum supression 34 | boxes = np.array([d.tlwh for d in detections]) 35 | scores = np.array([d.confidence for d in detections]) 36 | indices = non_max_suppression(boxes, self.nms_max_overlap, scores) 37 | detections = [detections[i] for i in indices] 38 | 39 | # update tracker 40 | self.tracker.predict() 41 | self.tracker.update(detections) 42 | 43 | # output bbox identities 44 | outputs = [] 45 | for track in self.tracker.tracks: 46 | if not track.is_confirmed() or track.time_since_update > 1: 47 | continue 48 | box = track.to_tlwh() 49 | x1,y1,x2,y2 = self._tlwh_to_xyxy(box) 50 | track_id = track.track_id 51 | outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=np.int)) 52 | if len(outputs) > 0: 53 | outputs = np.stack(outputs,axis=0) 54 | return outputs 55 | 56 | 57 | """ 58 | TODO: 59 | Convert bbox from xc_yc_w_h to xtl_ytl_w_h 60 | Thanks JieChen91@github.com for reporting this bug! 61 | """ 62 | @staticmethod 63 | def _xywh_to_tlwh(bbox_xywh): 64 | if isinstance(bbox_xywh, np.ndarray): 65 | bbox_tlwh = bbox_xywh.copy() 66 | elif isinstance(bbox_xywh, torch.Tensor): 67 | bbox_tlwh = bbox_xywh.clone() 68 | bbox_tlwh[:,0] = bbox_xywh[:,0] - bbox_xywh[:,2]/2. 69 | bbox_tlwh[:,1] = bbox_xywh[:,1] - bbox_xywh[:,3]/2. 70 | return bbox_tlwh 71 | 72 | 73 | def _xywh_to_xyxy(self, bbox_xywh): 74 | x,y,w,h = bbox_xywh 75 | x1 = max(int(x-w/2),0) 76 | x2 = min(int(x+w/2),self.width-1) 77 | y1 = max(int(y-h/2),0) 78 | y2 = min(int(y+h/2),self.height-1) 79 | return x1,y1,x2,y2 80 | 81 | def _tlwh_to_xyxy(self, bbox_tlwh): 82 | """ 83 | TODO: 84 | Convert bbox from xtl_ytl_w_h to xc_yc_w_h 85 | Thanks JieChen91@github.com for reporting this bug! 86 | """ 87 | x,y,w,h = bbox_tlwh 88 | x1 = max(int(x),0) 89 | x2 = min(int(x+w),self.width-1) 90 | y1 = max(int(y),0) 91 | y2 = min(int(y+h),self.height-1) 92 | return x1,y1,x2,y2 93 | 94 | def _xyxy_to_tlwh(self, bbox_xyxy): 95 | x1,y1,x2,y2 = bbox_xyxy 96 | 97 | t = x1 98 | l = y1 99 | w = int(x2-x1) 100 | h = int(y2-y1) 101 | return t,l,w,h 102 | 103 | def _get_features(self, bbox_xywh, ori_img): 104 | im_crops = [] 105 | for box in bbox_xywh: 106 | x1,y1,x2,y2 = self._xywh_to_xyxy(box) 107 | im = ori_img[y1:y2,x1:x2] 108 | im_crops.append(im) 109 | if im_crops: 110 | features = self.extractor(im_crops) 111 | else: 112 | features = np.array([]) 113 | return features 114 | 115 | 116 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__init__.py -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/detection.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/detection.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/detection.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/detection.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/iou_matching.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/iou_matching.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/iou_matching.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/iou_matching.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/kalman_filter.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/kalman_filter.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/kalman_filter.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/kalman_filter.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/linear_assignment.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/linear_assignment.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/linear_assignment.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/linear_assignment.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/nn_matching.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/nn_matching.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/nn_matching.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/nn_matching.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/preprocessing.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/preprocessing.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/preprocessing.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/preprocessing.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/track.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/track.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/track.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/track.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/tracker.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/tracker.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/__pycache__/tracker.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/deep_sort/sort/__pycache__/tracker.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/detection.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | import numpy as np 3 | 4 | 5 | class Detection(object): 6 | """ 7 | This class represents a bounding box detection in a single image. 8 | 9 | Parameters 10 | ---------- 11 | tlwh : array_like 12 | Bounding box in format `(x, y, w, h)`. 13 | confidence : float 14 | Detector confidence score. 15 | feature : array_like 16 | A feature vector that describes the object contained in this image. 17 | 18 | Attributes 19 | ---------- 20 | tlwh : ndarray 21 | Bounding box in format `(top left x, top left y, width, height)`. 22 | confidence : ndarray 23 | Detector confidence score. 24 | feature : ndarray | NoneType 25 | A feature vector that describes the object contained in this image. 26 | 27 | """ 28 | 29 | def __init__(self, tlwh, confidence, feature): 30 | self.tlwh = np.asarray(tlwh, dtype=np.float) 31 | self.confidence = float(confidence) 32 | self.feature = np.asarray(feature, dtype=np.float32) 33 | 34 | def to_tlbr(self): 35 | """Convert bounding box to format `(min x, min y, max x, max y)`, i.e., 36 | `(top left, bottom right)`. 37 | """ 38 | ret = self.tlwh.copy() 39 | ret[2:] += ret[:2] 40 | return ret 41 | 42 | def to_xyah(self): 43 | """Convert bounding box to format `(center x, center y, aspect ratio, 44 | height)`, where the aspect ratio is `width / height`. 45 | """ 46 | ret = self.tlwh.copy() 47 | ret[:2] += ret[2:] / 2 48 | ret[2] /= ret[3] 49 | return ret 50 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/iou_matching.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | from __future__ import absolute_import 3 | import numpy as np 4 | from . import linear_assignment 5 | 6 | 7 | def iou(bbox, candidates): 8 | """Computer intersection over union. 9 | 10 | Parameters 11 | ---------- 12 | bbox : ndarray 13 | A bounding box in format `(top left x, top left y, width, height)`. 14 | candidates : ndarray 15 | A matrix of candidate bounding boxes (one per row) in the same format 16 | as `bbox`. 17 | 18 | Returns 19 | ------- 20 | ndarray 21 | The intersection over union in [0, 1] between the `bbox` and each 22 | candidate. A higher score means a larger fraction of the `bbox` is 23 | occluded by the candidate. 24 | 25 | """ 26 | bbox_tl, bbox_br = bbox[:2], bbox[:2] + bbox[2:] 27 | candidates_tl = candidates[:, :2] 28 | candidates_br = candidates[:, :2] + candidates[:, 2:] 29 | 30 | tl = np.c_[np.maximum(bbox_tl[0], candidates_tl[:, 0])[:, np.newaxis], 31 | np.maximum(bbox_tl[1], candidates_tl[:, 1])[:, np.newaxis]] 32 | br = np.c_[np.minimum(bbox_br[0], candidates_br[:, 0])[:, np.newaxis], 33 | np.minimum(bbox_br[1], candidates_br[:, 1])[:, np.newaxis]] 34 | wh = np.maximum(0., br - tl) 35 | 36 | area_intersection = wh.prod(axis=1) 37 | area_bbox = bbox[2:].prod() 38 | area_candidates = candidates[:, 2:].prod(axis=1) 39 | return area_intersection / (area_bbox + area_candidates - area_intersection) 40 | 41 | 42 | def iou_cost(tracks, detections, track_indices=None, 43 | detection_indices=None): 44 | """An intersection over union distance metric. 45 | 46 | Parameters 47 | ---------- 48 | tracks : List[deep_sort.track.Track] 49 | A list of tracks. 50 | detections : List[deep_sort.detection.Detection] 51 | A list of detections. 52 | track_indices : Optional[List[int]] 53 | A list of indices to tracks that should be matched. Defaults to 54 | all `tracks`. 55 | detection_indices : Optional[List[int]] 56 | A list of indices to detections that should be matched. Defaults 57 | to all `detections`. 58 | 59 | Returns 60 | ------- 61 | ndarray 62 | Returns a cost matrix of shape 63 | len(track_indices), len(detection_indices) where entry (i, j) is 64 | `1 - iou(tracks[track_indices[i]], detections[detection_indices[j]])`. 65 | 66 | """ 67 | if track_indices is None: 68 | track_indices = np.arange(len(tracks)) 69 | if detection_indices is None: 70 | detection_indices = np.arange(len(detections)) 71 | 72 | cost_matrix = np.zeros((len(track_indices), len(detection_indices))) 73 | for row, track_idx in enumerate(track_indices): 74 | if tracks[track_idx].time_since_update > 1: 75 | cost_matrix[row, :] = linear_assignment.INFTY_COST 76 | continue 77 | 78 | bbox = tracks[track_idx].to_tlwh() 79 | candidates = np.asarray([detections[i].tlwh for i in detection_indices]) 80 | cost_matrix[row, :] = 1. - iou(bbox, candidates) 81 | return cost_matrix 82 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/kalman_filter.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | import numpy as np 3 | import scipy.linalg 4 | 5 | 6 | """ 7 | Table for the 0.95 quantile of the chi-square distribution with N degrees of 8 | freedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv 9 | function and used as Mahalanobis gating threshold. 10 | """ 11 | chi2inv95 = { 12 | 1: 3.8415, 13 | 2: 5.9915, 14 | 3: 7.8147, 15 | 4: 9.4877, 16 | 5: 11.070, 17 | 6: 12.592, 18 | 7: 14.067, 19 | 8: 15.507, 20 | 9: 16.919} 21 | 22 | 23 | class KalmanFilter(object): 24 | """ 25 | A simple Kalman filter for tracking bounding boxes in image space. 26 | 27 | The 8-dimensional state space 28 | 29 | x, y, a, h, vx, vy, va, vh 30 | 31 | contains the bounding box center position (x, y), aspect ratio a, height h, 32 | and their respective velocities. 33 | 34 | Object motion follows a constant velocity model. The bounding box location 35 | (x, y, a, h) is taken as direct observation of the state space (linear 36 | observation model). 37 | 38 | """ 39 | 40 | def __init__(self): 41 | ndim, dt = 4, 1. 42 | 43 | # Create Kalman filter model matrices. 44 | self._motion_mat = np.eye(2 * ndim, 2 * ndim) 45 | for i in range(ndim): 46 | self._motion_mat[i, ndim + i] = dt 47 | self._update_mat = np.eye(ndim, 2 * ndim) 48 | 49 | # Motion and observation uncertainty are chosen relative to the current 50 | # state estimate. These weights control the amount of uncertainty in 51 | # the model. This is a bit hacky. 52 | self._std_weight_position = 1. / 20 53 | self._std_weight_velocity = 1. / 160 54 | 55 | def initiate(self, measurement): 56 | """Create track from unassociated measurement. 57 | 58 | Parameters 59 | ---------- 60 | measurement : ndarray 61 | Bounding box coordinates (x, y, a, h) with center position (x, y), 62 | aspect ratio a, and height h. 63 | 64 | Returns 65 | ------- 66 | (ndarray, ndarray) 67 | Returns the mean vector (8 dimensional) and covariance matrix (8x8 68 | dimensional) of the new track. Unobserved velocities are initialized 69 | to 0 mean. 70 | 71 | """ 72 | mean_pos = measurement 73 | mean_vel = np.zeros_like(mean_pos) 74 | mean = np.r_[mean_pos, mean_vel] 75 | 76 | std = [ 77 | 2 * self._std_weight_position * measurement[3], 78 | 2 * self._std_weight_position * measurement[3], 79 | 1e-2, 80 | 2 * self._std_weight_position * measurement[3], 81 | 10 * self._std_weight_velocity * measurement[3], 82 | 10 * self._std_weight_velocity * measurement[3], 83 | 1e-5, 84 | 10 * self._std_weight_velocity * measurement[3]] 85 | covariance = np.diag(np.square(std)) 86 | return mean, covariance 87 | 88 | def predict(self, mean, covariance): 89 | """Run Kalman filter prediction step. 90 | 91 | Parameters 92 | ---------- 93 | mean : ndarray 94 | The 8 dimensional mean vector of the object state at the previous 95 | time step. 96 | covariance : ndarray 97 | The 8x8 dimensional covariance matrix of the object state at the 98 | previous time step. 99 | 100 | Returns 101 | ------- 102 | (ndarray, ndarray) 103 | Returns the mean vector and covariance matrix of the predicted 104 | state. Unobserved velocities are initialized to 0 mean. 105 | 106 | """ 107 | std_pos = [ 108 | self._std_weight_position * mean[3], 109 | self._std_weight_position * mean[3], 110 | 1e-2, 111 | self._std_weight_position * mean[3]] 112 | std_vel = [ 113 | self._std_weight_velocity * mean[3], 114 | self._std_weight_velocity * mean[3], 115 | 1e-5, 116 | self._std_weight_velocity * mean[3]] 117 | motion_cov = np.diag(np.square(np.r_[std_pos, std_vel])) 118 | 119 | mean = np.dot(self._motion_mat, mean) 120 | covariance = np.linalg.multi_dot(( 121 | self._motion_mat, covariance, self._motion_mat.T)) + motion_cov 122 | 123 | return mean, covariance 124 | 125 | def project(self, mean, covariance): 126 | """Project state distribution to measurement space. 127 | 128 | Parameters 129 | ---------- 130 | mean : ndarray 131 | The state's mean vector (8 dimensional array). 132 | covariance : ndarray 133 | The state's covariance matrix (8x8 dimensional). 134 | 135 | Returns 136 | ------- 137 | (ndarray, ndarray) 138 | Returns the projected mean and covariance matrix of the given state 139 | estimate. 140 | 141 | """ 142 | std = [ 143 | self._std_weight_position * mean[3], 144 | self._std_weight_position * mean[3], 145 | 1e-1, 146 | self._std_weight_position * mean[3]] 147 | innovation_cov = np.diag(np.square(std)) 148 | 149 | mean = np.dot(self._update_mat, mean) 150 | covariance = np.linalg.multi_dot(( 151 | self._update_mat, covariance, self._update_mat.T)) 152 | return mean, covariance + innovation_cov 153 | 154 | def update(self, mean, covariance, measurement): 155 | """Run Kalman filter correction step. 156 | 157 | Parameters 158 | ---------- 159 | mean : ndarray 160 | The predicted state's mean vector (8 dimensional). 161 | covariance : ndarray 162 | The state's covariance matrix (8x8 dimensional). 163 | measurement : ndarray 164 | The 4 dimensional measurement vector (x, y, a, h), where (x, y) 165 | is the center position, a the aspect ratio, and h the height of the 166 | bounding box. 167 | 168 | Returns 169 | ------- 170 | (ndarray, ndarray) 171 | Returns the measurement-corrected state distribution. 172 | 173 | """ 174 | projected_mean, projected_cov = self.project(mean, covariance) 175 | 176 | chol_factor, lower = scipy.linalg.cho_factor( 177 | projected_cov, lower=True, check_finite=False) 178 | kalman_gain = scipy.linalg.cho_solve( 179 | (chol_factor, lower), np.dot(covariance, self._update_mat.T).T, 180 | check_finite=False).T 181 | innovation = measurement - projected_mean 182 | 183 | new_mean = mean + np.dot(innovation, kalman_gain.T) 184 | new_covariance = covariance - np.linalg.multi_dot(( 185 | kalman_gain, projected_cov, kalman_gain.T)) 186 | return new_mean, new_covariance 187 | 188 | def gating_distance(self, mean, covariance, measurements, 189 | only_position=False): 190 | """Compute gating distance between state distribution and measurements. 191 | 192 | A suitable distance threshold can be obtained from `chi2inv95`. If 193 | `only_position` is False, the chi-square distribution has 4 degrees of 194 | freedom, otherwise 2. 195 | 196 | Parameters 197 | ---------- 198 | mean : ndarray 199 | Mean vector over the state distribution (8 dimensional). 200 | covariance : ndarray 201 | Covariance of the state distribution (8x8 dimensional). 202 | measurements : ndarray 203 | An Nx4 dimensional matrix of N measurements, each in 204 | format (x, y, a, h) where (x, y) is the bounding box center 205 | position, a the aspect ratio, and h the height. 206 | only_position : Optional[bool] 207 | If True, distance computation is done with respect to the bounding 208 | box center position only. 209 | 210 | Returns 211 | ------- 212 | ndarray 213 | Returns an array of length N, where the i-th element contains the 214 | squared Mahalanobis distance between (mean, covariance) and 215 | `measurements[i]`. 216 | 217 | """ 218 | mean, covariance = self.project(mean, covariance) 219 | if only_position: 220 | mean, covariance = mean[:2], covariance[:2, :2] 221 | measurements = measurements[:, :2] 222 | 223 | cholesky_factor = np.linalg.cholesky(covariance) 224 | d = measurements - mean 225 | z = scipy.linalg.solve_triangular( 226 | cholesky_factor, d.T, lower=True, check_finite=False, 227 | overwrite_b=True) 228 | squared_maha = np.sum(z * z, axis=0) 229 | return squared_maha 230 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/linear_assignment.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | from __future__ import absolute_import 3 | import numpy as np 4 | # from sklearn.utils.linear_assignment_ import linear_assignment 5 | from scipy.optimize import linear_sum_assignment as linear_assignment 6 | from . import kalman_filter 7 | 8 | 9 | INFTY_COST = 1e+5 10 | 11 | 12 | def min_cost_matching( 13 | distance_metric, max_distance, tracks, detections, track_indices=None, 14 | detection_indices=None): 15 | """Solve linear assignment problem. 16 | 17 | Parameters 18 | ---------- 19 | distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray 20 | The distance metric is given a list of tracks and detections as well as 21 | a list of N track indices and M detection indices. The metric should 22 | return the NxM dimensional cost matrix, where element (i, j) is the 23 | association cost between the i-th track in the given track indices and 24 | the j-th detection in the given detection_indices. 25 | max_distance : float 26 | Gating threshold. Associations with cost larger than this value are 27 | disregarded. 28 | tracks : List[track.Track] 29 | A list of predicted tracks at the current time step. 30 | detections : List[detection.Detection] 31 | A list of detections at the current time step. 32 | track_indices : List[int] 33 | List of track indices that maps rows in `cost_matrix` to tracks in 34 | `tracks` (see description above). 35 | detection_indices : List[int] 36 | List of detection indices that maps columns in `cost_matrix` to 37 | detections in `detections` (see description above). 38 | 39 | Returns 40 | ------- 41 | (List[(int, int)], List[int], List[int]) 42 | Returns a tuple with the following three entries: 43 | * A list of matched track and detection indices. 44 | * A list of unmatched track indices. 45 | * A list of unmatched detection indices. 46 | 47 | """ 48 | if track_indices is None: 49 | track_indices = np.arange(len(tracks)) 50 | if detection_indices is None: 51 | detection_indices = np.arange(len(detections)) 52 | 53 | if len(detection_indices) == 0 or len(track_indices) == 0: 54 | return [], track_indices, detection_indices # Nothing to match. 55 | 56 | cost_matrix = distance_metric( 57 | tracks, detections, track_indices, detection_indices) 58 | cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5 59 | 60 | row_indices, col_indices = linear_assignment(cost_matrix) 61 | 62 | matches, unmatched_tracks, unmatched_detections = [], [], [] 63 | for col, detection_idx in enumerate(detection_indices): 64 | if col not in col_indices: 65 | unmatched_detections.append(detection_idx) 66 | for row, track_idx in enumerate(track_indices): 67 | if row not in row_indices: 68 | unmatched_tracks.append(track_idx) 69 | for row, col in zip(row_indices, col_indices): 70 | track_idx = track_indices[row] 71 | detection_idx = detection_indices[col] 72 | if cost_matrix[row, col] > max_distance: 73 | unmatched_tracks.append(track_idx) 74 | unmatched_detections.append(detection_idx) 75 | else: 76 | matches.append((track_idx, detection_idx)) 77 | return matches, unmatched_tracks, unmatched_detections 78 | 79 | 80 | def matching_cascade( 81 | distance_metric, max_distance, cascade_depth, tracks, detections, 82 | track_indices=None, detection_indices=None): 83 | """Run matching cascade. 84 | 85 | Parameters 86 | ---------- 87 | distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray 88 | The distance metric is given a list of tracks and detections as well as 89 | a list of N track indices and M detection indices. The metric should 90 | return the NxM dimensional cost matrix, where element (i, j) is the 91 | association cost between the i-th track in the given track indices and 92 | the j-th detection in the given detection indices. 93 | max_distance : float 94 | Gating threshold. Associations with cost larger than this value are 95 | disregarded. 96 | cascade_depth: int 97 | The cascade depth, should be se to the maximum track age. 98 | tracks : List[track.Track] 99 | A list of predicted tracks at the current time step. 100 | detections : List[detection.Detection] 101 | A list of detections at the current time step. 102 | track_indices : Optional[List[int]] 103 | List of track indices that maps rows in `cost_matrix` to tracks in 104 | `tracks` (see description above). Defaults to all tracks. 105 | detection_indices : Optional[List[int]] 106 | List of detection indices that maps columns in `cost_matrix` to 107 | detections in `detections` (see description above). Defaults to all 108 | detections. 109 | 110 | Returns 111 | ------- 112 | (List[(int, int)], List[int], List[int]) 113 | Returns a tuple with the following three entries: 114 | * A list of matched track and detection indices. 115 | * A list of unmatched track indices. 116 | * A list of unmatched detection indices. 117 | 118 | """ 119 | if track_indices is None: 120 | track_indices = list(range(len(tracks))) 121 | if detection_indices is None: 122 | detection_indices = list(range(len(detections))) 123 | 124 | unmatched_detections = detection_indices 125 | matches = [] 126 | for level in range(cascade_depth): 127 | if len(unmatched_detections) == 0: # No detections left 128 | break 129 | 130 | track_indices_l = [ 131 | k for k in track_indices 132 | if tracks[k].time_since_update == 1 + level 133 | ] 134 | if len(track_indices_l) == 0: # Nothing to match at this level 135 | continue 136 | 137 | matches_l, _, unmatched_detections = \ 138 | min_cost_matching( 139 | distance_metric, max_distance, tracks, detections, 140 | track_indices_l, unmatched_detections) 141 | matches += matches_l 142 | unmatched_tracks = list(set(track_indices) - set(k for k, _ in matches)) 143 | return matches, unmatched_tracks, unmatched_detections 144 | 145 | 146 | def gate_cost_matrix( 147 | kf, cost_matrix, tracks, detections, track_indices, detection_indices, 148 | gated_cost=INFTY_COST, only_position=False): 149 | """Invalidate infeasible entries in cost matrix based on the state 150 | distributions obtained by Kalman filtering. 151 | 152 | Parameters 153 | ---------- 154 | kf : The Kalman filter. 155 | cost_matrix : ndarray 156 | The NxM dimensional cost matrix, where N is the number of track indices 157 | and M is the number of detection indices, such that entry (i, j) is the 158 | association cost between `tracks[track_indices[i]]` and 159 | `detections[detection_indices[j]]`. 160 | tracks : List[track.Track] 161 | A list of predicted tracks at the current time step. 162 | detections : List[detection.Detection] 163 | A list of detections at the current time step. 164 | track_indices : List[int] 165 | List of track indices that maps rows in `cost_matrix` to tracks in 166 | `tracks` (see description above). 167 | detection_indices : List[int] 168 | List of detection indices that maps columns in `cost_matrix` to 169 | detections in `detections` (see description above). 170 | gated_cost : Optional[float] 171 | Entries in the cost matrix corresponding to infeasible associations are 172 | set this value. Defaults to a very large value. 173 | only_position : Optional[bool] 174 | If True, only the x, y position of the state distribution is considered 175 | during gating. Defaults to False. 176 | 177 | Returns 178 | ------- 179 | ndarray 180 | Returns the modified cost matrix. 181 | 182 | """ 183 | gating_dim = 2 if only_position else 4 184 | gating_threshold = kalman_filter.chi2inv95[gating_dim] 185 | measurements = np.asarray( 186 | [detections[i].to_xyah() for i in detection_indices]) 187 | for row, track_idx in enumerate(track_indices): 188 | track = tracks[track_idx] 189 | gating_distance = kf.gating_distance( 190 | track.mean, track.covariance, measurements, only_position) 191 | cost_matrix[row, gating_distance > gating_threshold] = gated_cost 192 | return cost_matrix 193 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/nn_matching.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | import numpy as np 3 | 4 | 5 | def _pdist(a, b): 6 | """Compute pair-wise squared distance between points in `a` and `b`. 7 | 8 | Parameters 9 | ---------- 10 | a : array_like 11 | An NxM matrix of N samples of dimensionality M. 12 | b : array_like 13 | An LxM matrix of L samples of dimensionality M. 14 | 15 | Returns 16 | ------- 17 | ndarray 18 | Returns a matrix of size len(a), len(b) such that eleement (i, j) 19 | contains the squared distance between `a[i]` and `b[j]`. 20 | 21 | """ 22 | a, b = np.asarray(a), np.asarray(b) 23 | if len(a) == 0 or len(b) == 0: 24 | return np.zeros((len(a), len(b))) 25 | a2, b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1) 26 | r2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :] 27 | r2 = np.clip(r2, 0., float(np.inf)) 28 | return r2 29 | 30 | 31 | def _cosine_distance(a, b, data_is_normalized=False): 32 | """Compute pair-wise cosine distance between points in `a` and `b`. 33 | 34 | Parameters 35 | ---------- 36 | a : array_like 37 | An NxM matrix of N samples of dimensionality M. 38 | b : array_like 39 | An LxM matrix of L samples of dimensionality M. 40 | data_is_normalized : Optional[bool] 41 | If True, assumes rows in a and b are unit length vectors. 42 | Otherwise, a and b are explicitly normalized to lenght 1. 43 | 44 | Returns 45 | ------- 46 | ndarray 47 | Returns a matrix of size len(a), len(b) such that eleement (i, j) 48 | contains the squared distance between `a[i]` and `b[j]`. 49 | 50 | """ 51 | if not data_is_normalized: 52 | a = np.asarray(a) / np.linalg.norm(a, axis=1, keepdims=True) 53 | b = np.asarray(b) / np.linalg.norm(b, axis=1, keepdims=True) 54 | return 1. - np.dot(a, b.T) 55 | 56 | 57 | def _nn_euclidean_distance(x, y): 58 | """ Helper function for nearest neighbor distance metric (Euclidean). 59 | 60 | Parameters 61 | ---------- 62 | x : ndarray 63 | A matrix of N row-vectors (sample points). 64 | y : ndarray 65 | A matrix of M row-vectors (query points). 66 | 67 | Returns 68 | ------- 69 | ndarray 70 | A vector of length M that contains for each entry in `y` the 71 | smallest Euclidean distance to a sample in `x`. 72 | 73 | """ 74 | distances = _pdist(x, y) 75 | return np.maximum(0.0, distances.min(axis=0)) 76 | 77 | 78 | def _nn_cosine_distance(x, y): 79 | """ Helper function for nearest neighbor distance metric (cosine). 80 | 81 | Parameters 82 | ---------- 83 | x : ndarray 84 | A matrix of N row-vectors (sample points). 85 | y : ndarray 86 | A matrix of M row-vectors (query points). 87 | 88 | Returns 89 | ------- 90 | ndarray 91 | A vector of length M that contains for each entry in `y` the 92 | smallest cosine distance to a sample in `x`. 93 | 94 | """ 95 | distances = _cosine_distance(x, y) 96 | return distances.min(axis=0) 97 | 98 | 99 | class NearestNeighborDistanceMetric(object): 100 | """ 101 | A nearest neighbor distance metric that, for each target, returns 102 | the closest distance to any sample that has been observed so far. 103 | 104 | Parameters 105 | ---------- 106 | metric : str 107 | Either "euclidean" or "cosine". 108 | matching_threshold: float 109 | The matching threshold. Samples with larger distance are considered an 110 | invalid match. 111 | budget : Optional[int] 112 | If not None, fix samples per class to at most this number. Removes 113 | the oldest samples when the budget is reached. 114 | 115 | Attributes 116 | ---------- 117 | samples : Dict[int -> List[ndarray]] 118 | A dictionary that maps from target identities to the list of samples 119 | that have been observed so far. 120 | 121 | """ 122 | 123 | def __init__(self, metric, matching_threshold, budget=None): 124 | 125 | 126 | if metric == "euclidean": 127 | self._metric = _nn_euclidean_distance 128 | elif metric == "cosine": 129 | self._metric = _nn_cosine_distance 130 | else: 131 | raise ValueError( 132 | "Invalid metric; must be either 'euclidean' or 'cosine'") 133 | self.matching_threshold = matching_threshold 134 | self.budget = budget 135 | self.samples = {} 136 | 137 | def partial_fit(self, features, targets, active_targets): 138 | """Update the distance metric with new data. 139 | 140 | Parameters 141 | ---------- 142 | features : ndarray 143 | An NxM matrix of N features of dimensionality M. 144 | targets : ndarray 145 | An integer array of associated target identities. 146 | active_targets : List[int] 147 | A list of targets that are currently present in the scene. 148 | 149 | """ 150 | for feature, target in zip(features, targets): 151 | self.samples.setdefault(target, []).append(feature) 152 | if self.budget is not None: 153 | self.samples[target] = self.samples[target][-self.budget:] 154 | self.samples = {k: self.samples[k] for k in active_targets} 155 | 156 | def distance(self, features, targets): 157 | """Compute distance between features and targets. 158 | 159 | Parameters 160 | ---------- 161 | features : ndarray 162 | An NxM matrix of N features of dimensionality M. 163 | targets : List[int] 164 | A list of targets to match the given `features` against. 165 | 166 | Returns 167 | ------- 168 | ndarray 169 | Returns a cost matrix of shape len(targets), len(features), where 170 | element (i, j) contains the closest squared distance between 171 | `targets[i]` and `features[j]`. 172 | 173 | """ 174 | cost_matrix = np.zeros((len(targets), len(features))) 175 | for i, target in enumerate(targets): 176 | cost_matrix[i, :] = self._metric(self.samples[target], features) 177 | return cost_matrix 178 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/preprocessing.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | import numpy as np 3 | import cv2 4 | 5 | 6 | def non_max_suppression(boxes, max_bbox_overlap, scores=None): 7 | """Suppress overlapping detections. 8 | 9 | Original code from [1]_ has been adapted to include confidence score. 10 | 11 | .. [1] http://www.pyimagesearch.com/2015/02/16/ 12 | faster-non-maximum-suppression-python/ 13 | 14 | Examples 15 | -------- 16 | 17 | >>> boxes = [d.roi for d in detections] 18 | >>> scores = [d.confidence for d in detections] 19 | >>> indices = non_max_suppression(boxes, max_bbox_overlap, scores) 20 | >>> detections = [detections[i] for i in indices] 21 | 22 | Parameters 23 | ---------- 24 | boxes : ndarray 25 | Array of ROIs (x, y, width, height). 26 | max_bbox_overlap : float 27 | ROIs that overlap more than this values are suppressed. 28 | scores : Optional[array_like] 29 | Detector confidence score. 30 | 31 | Returns 32 | ------- 33 | List[int] 34 | Returns indices of detections that have survived non-maxima suppression. 35 | 36 | """ 37 | if len(boxes) == 0: 38 | return [] 39 | 40 | boxes = boxes.astype(np.float) 41 | pick = [] 42 | 43 | x1 = boxes[:, 0] 44 | y1 = boxes[:, 1] 45 | x2 = boxes[:, 2] + boxes[:, 0] 46 | y2 = boxes[:, 3] + boxes[:, 1] 47 | 48 | area = (x2 - x1 + 1) * (y2 - y1 + 1) 49 | if scores is not None: 50 | idxs = np.argsort(scores) 51 | else: 52 | idxs = np.argsort(y2) 53 | 54 | while len(idxs) > 0: 55 | last = len(idxs) - 1 56 | i = idxs[last] 57 | pick.append(i) 58 | 59 | xx1 = np.maximum(x1[i], x1[idxs[:last]]) 60 | yy1 = np.maximum(y1[i], y1[idxs[:last]]) 61 | xx2 = np.minimum(x2[i], x2[idxs[:last]]) 62 | yy2 = np.minimum(y2[i], y2[idxs[:last]]) 63 | 64 | w = np.maximum(0, xx2 - xx1 + 1) 65 | h = np.maximum(0, yy2 - yy1 + 1) 66 | 67 | overlap = (w * h) / area[idxs[:last]] 68 | 69 | idxs = np.delete( 70 | idxs, np.concatenate( 71 | ([last], np.where(overlap > max_bbox_overlap)[0]))) 72 | 73 | return pick 74 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/track.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | 3 | 4 | class TrackState: 5 | """ 6 | Enumeration type for the single target track state. Newly created tracks are 7 | classified as `tentative` until enough evidence has been collected. Then, 8 | the track state is changed to `confirmed`. Tracks that are no longer alive 9 | are classified as `deleted` to mark them for removal from the set of active 10 | tracks. 11 | 12 | """ 13 | 14 | Tentative = 1 15 | Confirmed = 2 16 | Deleted = 3 17 | 18 | 19 | class Track: 20 | """ 21 | A single target track with state space `(x, y, a, h)` and associated 22 | velocities, where `(x, y)` is the center of the bounding box, `a` is the 23 | aspect ratio and `h` is the height. 24 | 25 | Parameters 26 | ---------- 27 | mean : ndarray 28 | Mean vector of the initial state distribution. 29 | covariance : ndarray 30 | Covariance matrix of the initial state distribution. 31 | track_id : int 32 | A unique track identifier. 33 | n_init : int 34 | Number of consecutive detections before the track is confirmed. The 35 | track state is set to `Deleted` if a miss occurs within the first 36 | `n_init` frames. 37 | max_age : int 38 | The maximum number of consecutive misses before the track state is 39 | set to `Deleted`. 40 | feature : Optional[ndarray] 41 | Feature vector of the detection this track originates from. If not None, 42 | this feature is added to the `features` cache. 43 | 44 | Attributes 45 | ---------- 46 | mean : ndarray 47 | Mean vector of the initial state distribution. 48 | covariance : ndarray 49 | Covariance matrix of the initial state distribution. 50 | track_id : int 51 | A unique track identifier. 52 | hits : int 53 | Total number of measurement updates. 54 | age : int 55 | Total number of frames since first occurance. 56 | time_since_update : int 57 | Total number of frames since last measurement update. 58 | state : TrackState 59 | The current track state. 60 | features : List[ndarray] 61 | A cache of features. On each measurement update, the associated feature 62 | vector is added to this list. 63 | 64 | """ 65 | 66 | def __init__(self, mean, covariance, track_id, n_init, max_age, 67 | feature=None): 68 | self.mean = mean 69 | self.covariance = covariance 70 | self.track_id = track_id 71 | self.hits = 1 72 | self.age = 1 73 | self.time_since_update = 0 74 | 75 | self.state = TrackState.Tentative 76 | self.features = [] 77 | if feature is not None: 78 | self.features.append(feature) 79 | 80 | self._n_init = n_init 81 | self._max_age = max_age 82 | 83 | def to_tlwh(self): 84 | """Get current position in bounding box format `(top left x, top left y, 85 | width, height)`. 86 | 87 | Returns 88 | ------- 89 | ndarray 90 | The bounding box. 91 | 92 | """ 93 | ret = self.mean[:4].copy() 94 | ret[2] *= ret[3] 95 | ret[:2] -= ret[2:] / 2 96 | return ret 97 | 98 | def to_tlbr(self): 99 | """Get current position in bounding box format `(min x, miny, max x, 100 | max y)`. 101 | 102 | Returns 103 | ------- 104 | ndarray 105 | The bounding box. 106 | 107 | """ 108 | ret = self.to_tlwh() 109 | ret[2:] = ret[:2] + ret[2:] 110 | return ret 111 | 112 | def predict(self, kf): 113 | """Propagate the state distribution to the current time step using a 114 | Kalman filter prediction step. 115 | 116 | Parameters 117 | ---------- 118 | kf : kalman_filter.KalmanFilter 119 | The Kalman filter. 120 | 121 | """ 122 | self.mean, self.covariance = kf.predict(self.mean, self.covariance) 123 | self.age += 1 124 | self.time_since_update += 1 125 | 126 | def update(self, kf, detection): 127 | """Perform Kalman filter measurement update step and update the feature 128 | cache. 129 | 130 | Parameters 131 | ---------- 132 | kf : kalman_filter.KalmanFilter 133 | The Kalman filter. 134 | detection : Detection 135 | The associated detection. 136 | 137 | """ 138 | self.mean, self.covariance = kf.update( 139 | self.mean, self.covariance, detection.to_xyah()) 140 | self.features.append(detection.feature) 141 | 142 | self.hits += 1 143 | self.time_since_update = 0 144 | if self.state == TrackState.Tentative and self.hits >= self._n_init: 145 | self.state = TrackState.Confirmed 146 | 147 | def mark_missed(self): 148 | """Mark this track as missed (no association at the current time step). 149 | """ 150 | if self.state == TrackState.Tentative: 151 | self.state = TrackState.Deleted 152 | elif self.time_since_update > self._max_age: 153 | self.state = TrackState.Deleted 154 | 155 | def is_tentative(self): 156 | """Returns True if this track is tentative (unconfirmed). 157 | """ 158 | return self.state == TrackState.Tentative 159 | 160 | def is_confirmed(self): 161 | """Returns True if this track is confirmed.""" 162 | return self.state == TrackState.Confirmed 163 | 164 | def is_deleted(self): 165 | """Returns True if this track is dead and should be deleted.""" 166 | return self.state == TrackState.Deleted 167 | -------------------------------------------------------------------------------- /deep_sort/deep_sort/sort/tracker.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | from __future__ import absolute_import 3 | import numpy as np 4 | from . import kalman_filter 5 | from . import linear_assignment 6 | from . import iou_matching 7 | from .track import Track 8 | 9 | 10 | class Tracker: 11 | """ 12 | This is the multi-target tracker. 13 | 14 | Parameters 15 | ---------- 16 | metric : nn_matching.NearestNeighborDistanceMetric 17 | A distance metric for measurement-to-track association. 18 | max_age : int 19 | Maximum number of missed misses before a track is deleted. 20 | n_init : int 21 | Number of consecutive detections before the track is confirmed. The 22 | track state is set to `Deleted` if a miss occurs within the first 23 | `n_init` frames. 24 | 25 | Attributes 26 | ---------- 27 | metric : nn_matching.NearestNeighborDistanceMetric 28 | The distance metric used for measurement to track association. 29 | max_age : int 30 | Maximum number of missed misses before a track is deleted. 31 | n_init : int 32 | Number of frames that a track remains in initialization phase. 33 | kf : kalman_filter.KalmanFilter 34 | A Kalman filter to filter target trajectories in image space. 35 | tracks : List[Track] 36 | The list of active tracks at the current time step. 37 | 38 | """ 39 | 40 | def __init__(self, metric, max_iou_distance=0.7, max_age=70, n_init=3): 41 | self.metric = metric 42 | self.max_iou_distance = max_iou_distance 43 | self.max_age = max_age 44 | self.n_init = n_init 45 | 46 | self.kf = kalman_filter.KalmanFilter() 47 | self.tracks = [] 48 | self._next_id = 1 49 | 50 | def predict(self): 51 | """Propagate track state distributions one time step forward. 52 | 53 | This function should be called once every time step, before `update`. 54 | """ 55 | for track in self.tracks: 56 | track.predict(self.kf) 57 | 58 | def update(self, detections): 59 | """Perform measurement update and track management. 60 | 61 | Parameters 62 | ---------- 63 | detections : List[deep_sort.detection.Detection] 64 | A list of detections at the current time step. 65 | 66 | """ 67 | # Run matching cascade. 68 | matches, unmatched_tracks, unmatched_detections = \ 69 | self._match(detections) 70 | 71 | # Update track set. 72 | for track_idx, detection_idx in matches: 73 | self.tracks[track_idx].update( 74 | self.kf, detections[detection_idx]) 75 | for track_idx in unmatched_tracks: 76 | self.tracks[track_idx].mark_missed() 77 | for detection_idx in unmatched_detections: 78 | self._initiate_track(detections[detection_idx]) 79 | self.tracks = [t for t in self.tracks if not t.is_deleted()] 80 | 81 | # Update distance metric. 82 | active_targets = [t.track_id for t in self.tracks if t.is_confirmed()] 83 | features, targets = [], [] 84 | for track in self.tracks: 85 | if not track.is_confirmed(): 86 | continue 87 | features += track.features 88 | targets += [track.track_id for _ in track.features] 89 | track.features = [] 90 | self.metric.partial_fit( 91 | np.asarray(features), np.asarray(targets), active_targets) 92 | 93 | def _match(self, detections): 94 | 95 | def gated_metric(tracks, dets, track_indices, detection_indices): 96 | features = np.array([dets[i].feature for i in detection_indices]) 97 | targets = np.array([tracks[i].track_id for i in track_indices]) 98 | cost_matrix = self.metric.distance(features, targets) 99 | cost_matrix = linear_assignment.gate_cost_matrix( 100 | self.kf, cost_matrix, tracks, dets, track_indices, 101 | detection_indices) 102 | 103 | return cost_matrix 104 | 105 | # Split track set into confirmed and unconfirmed tracks. 106 | confirmed_tracks = [ 107 | i for i, t in enumerate(self.tracks) if t.is_confirmed()] 108 | unconfirmed_tracks = [ 109 | i for i, t in enumerate(self.tracks) if not t.is_confirmed()] 110 | 111 | # Associate confirmed tracks using appearance features. 112 | matches_a, unmatched_tracks_a, unmatched_detections = \ 113 | linear_assignment.matching_cascade( 114 | gated_metric, self.metric.matching_threshold, self.max_age, 115 | self.tracks, detections, confirmed_tracks) 116 | 117 | # Associate remaining tracks together with unconfirmed tracks using IOU. 118 | iou_track_candidates = unconfirmed_tracks + [ 119 | k for k in unmatched_tracks_a if 120 | self.tracks[k].time_since_update == 1] 121 | unmatched_tracks_a = [ 122 | k for k in unmatched_tracks_a if 123 | self.tracks[k].time_since_update != 1] 124 | matches_b, unmatched_tracks_b, unmatched_detections = \ 125 | linear_assignment.min_cost_matching( 126 | iou_matching.iou_cost, self.max_iou_distance, self.tracks, 127 | detections, iou_track_candidates, unmatched_detections) 128 | 129 | matches = matches_a + matches_b 130 | unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b)) 131 | return matches, unmatched_tracks, unmatched_detections 132 | 133 | def _initiate_track(self, detection): 134 | mean, covariance = self.kf.initiate(detection.to_xyah()) 135 | self.tracks.append(Track( 136 | mean, covariance, self._next_id, self.n_init, self.max_age, 137 | detection.feature)) 138 | self._next_id += 1 139 | -------------------------------------------------------------------------------- /deep_sort/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/utils/__init__.py -------------------------------------------------------------------------------- /deep_sort/utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/utils/__pycache__/parser.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/utils/__pycache__/parser.cpython-36.pyc -------------------------------------------------------------------------------- /deep_sort/utils/__pycache__/parser.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/deep_sort/utils/__pycache__/parser.cpython-37.pyc -------------------------------------------------------------------------------- /deep_sort/utils/asserts.py: -------------------------------------------------------------------------------- 1 | from os import environ 2 | 3 | 4 | def assert_in(file, files_to_check): 5 | if file not in files_to_check: 6 | raise AssertionError("{} does not exist in the list".format(str(file))) 7 | return True 8 | 9 | 10 | def assert_in_env(check_list: list): 11 | for item in check_list: 12 | assert_in(item, environ.keys()) 13 | return True 14 | -------------------------------------------------------------------------------- /deep_sort/utils/draw.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) 5 | 6 | 7 | def compute_color_for_labels(label): 8 | """ 9 | Simple function that adds fixed color depending on the class 10 | """ 11 | color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette] 12 | return tuple(color) 13 | 14 | 15 | def draw_boxes(img, bbox, identities=None, offset=(0,0)): 16 | for i,box in enumerate(bbox): 17 | x1,y1,x2,y2 = [int(i) for i in box] 18 | x1 += offset[0] 19 | x2 += offset[0] 20 | y1 += offset[1] 21 | y2 += offset[1] 22 | # box text and bar 23 | id = int(identities[i]) if identities is not None else 0 24 | color = compute_color_for_labels(id) 25 | label = '{}{:d}'.format("", id) 26 | t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0] 27 | cv2.rectangle(img,(x1, y1),(x2,y2),color,3) 28 | cv2.rectangle(img,(x1, y1),(x1+t_size[0]+3,y1+t_size[1]+4), color,-1) 29 | cv2.putText(img,label,(x1,y1+t_size[1]+4), cv2.FONT_HERSHEY_PLAIN, 2, [255,255,255], 2) 30 | return img 31 | 32 | 33 | 34 | if __name__ == '__main__': 35 | for i in range(82): 36 | print(compute_color_for_labels(i)) 37 | -------------------------------------------------------------------------------- /deep_sort/utils/evaluation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import copy 4 | import motmetrics as mm 5 | mm.lap.default_solver = 'lap' 6 | from utils.io import read_results, unzip_objs 7 | 8 | 9 | class Evaluator(object): 10 | 11 | def __init__(self, data_root, seq_name, data_type): 12 | self.data_root = data_root 13 | self.seq_name = seq_name 14 | self.data_type = data_type 15 | 16 | self.load_annotations() 17 | self.reset_accumulator() 18 | 19 | def load_annotations(self): 20 | assert self.data_type == 'mot' 21 | 22 | gt_filename = os.path.join(self.data_root, self.seq_name, 'gt', 'gt.txt') 23 | self.gt_frame_dict = read_results(gt_filename, self.data_type, is_gt=True) 24 | self.gt_ignore_frame_dict = read_results(gt_filename, self.data_type, is_ignore=True) 25 | 26 | def reset_accumulator(self): 27 | self.acc = mm.MOTAccumulator(auto_id=True) 28 | 29 | def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False): 30 | # results 31 | trk_tlwhs = np.copy(trk_tlwhs) 32 | trk_ids = np.copy(trk_ids) 33 | 34 | # gts 35 | gt_objs = self.gt_frame_dict.get(frame_id, []) 36 | gt_tlwhs, gt_ids = unzip_objs(gt_objs)[:2] 37 | 38 | # ignore boxes 39 | ignore_objs = self.gt_ignore_frame_dict.get(frame_id, []) 40 | ignore_tlwhs = unzip_objs(ignore_objs)[0] 41 | 42 | 43 | # remove ignored results 44 | keep = np.ones(len(trk_tlwhs), dtype=bool) 45 | iou_distance = mm.distances.iou_matrix(ignore_tlwhs, trk_tlwhs, max_iou=0.5) 46 | if len(iou_distance) > 0: 47 | match_is, match_js = mm.lap.linear_sum_assignment(iou_distance) 48 | match_is, match_js = map(lambda a: np.asarray(a, dtype=int), [match_is, match_js]) 49 | match_ious = iou_distance[match_is, match_js] 50 | 51 | match_js = np.asarray(match_js, dtype=int) 52 | match_js = match_js[np.logical_not(np.isnan(match_ious))] 53 | keep[match_js] = False 54 | trk_tlwhs = trk_tlwhs[keep] 55 | trk_ids = trk_ids[keep] 56 | 57 | # get distance matrix 58 | iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5) 59 | 60 | # acc 61 | self.acc.update(gt_ids, trk_ids, iou_distance) 62 | 63 | if rtn_events and iou_distance.size > 0 and hasattr(self.acc, 'last_mot_events'): 64 | events = self.acc.last_mot_events # only supported by https://github.com/longcw/py-motmetrics 65 | else: 66 | events = None 67 | return events 68 | 69 | def eval_file(self, filename): 70 | self.reset_accumulator() 71 | 72 | result_frame_dict = read_results(filename, self.data_type, is_gt=False) 73 | frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys()))) 74 | for frame_id in frames: 75 | trk_objs = result_frame_dict.get(frame_id, []) 76 | trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2] 77 | self.eval_frame(frame_id, trk_tlwhs, trk_ids, rtn_events=False) 78 | 79 | return self.acc 80 | 81 | @staticmethod 82 | def get_summary(accs, names, metrics=('mota', 'num_switches', 'idp', 'idr', 'idf1', 'precision', 'recall')): 83 | names = copy.deepcopy(names) 84 | if metrics is None: 85 | metrics = mm.metrics.motchallenge_metrics 86 | metrics = copy.deepcopy(metrics) 87 | 88 | mh = mm.metrics.create() 89 | summary = mh.compute_many( 90 | accs, 91 | metrics=metrics, 92 | names=names, 93 | generate_overall=True 94 | ) 95 | 96 | return summary 97 | 98 | @staticmethod 99 | def save_summary(summary, filename): 100 | import pandas as pd 101 | writer = pd.ExcelWriter(filename) 102 | summary.to_excel(writer) 103 | writer.save() 104 | -------------------------------------------------------------------------------- /deep_sort/utils/io.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Dict 3 | import numpy as np 4 | 5 | # from utils.log import get_logger 6 | 7 | 8 | def write_results(filename, results, data_type): 9 | if data_type == 'mot': 10 | save_format = '{frame},{id},{x1},{y1},{w},{h},-1,-1,-1,-1\n' 11 | elif data_type == 'kitti': 12 | save_format = '{frame} {id} pedestrian 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n' 13 | else: 14 | raise ValueError(data_type) 15 | 16 | with open(filename, 'w') as f: 17 | for frame_id, tlwhs, track_ids in results: 18 | if data_type == 'kitti': 19 | frame_id -= 1 20 | for tlwh, track_id in zip(tlwhs, track_ids): 21 | if track_id < 0: 22 | continue 23 | x1, y1, w, h = tlwh 24 | x2, y2 = x1 + w, y1 + h 25 | line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h) 26 | f.write(line) 27 | 28 | 29 | # def write_results(filename, results_dict: Dict, data_type: str): 30 | # if not filename: 31 | # return 32 | # path = os.path.dirname(filename) 33 | # if not os.path.exists(path): 34 | # os.makedirs(path) 35 | 36 | # if data_type in ('mot', 'mcmot', 'lab'): 37 | # save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n' 38 | # elif data_type == 'kitti': 39 | # save_format = '{frame} {id} pedestrian -1 -1 -10 {x1} {y1} {x2} {y2} -1 -1 -1 -1000 -1000 -1000 -10 {score}\n' 40 | # else: 41 | # raise ValueError(data_type) 42 | 43 | # with open(filename, 'w') as f: 44 | # for frame_id, frame_data in results_dict.items(): 45 | # if data_type == 'kitti': 46 | # frame_id -= 1 47 | # for tlwh, track_id in frame_data: 48 | # if track_id < 0: 49 | # continue 50 | # x1, y1, w, h = tlwh 51 | # x2, y2 = x1 + w, y1 + h 52 | # line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h, score=1.0) 53 | # f.write(line) 54 | # logger.info('Save results to {}'.format(filename)) 55 | 56 | 57 | def read_results(filename, data_type: str, is_gt=False, is_ignore=False): 58 | if data_type in ('mot', 'lab'): 59 | read_fun = read_mot_results 60 | else: 61 | raise ValueError('Unknown data type: {}'.format(data_type)) 62 | 63 | return read_fun(filename, is_gt, is_ignore) 64 | 65 | 66 | """ 67 | labels={'ped', ... % 1 68 | 'person_on_vhcl', ... % 2 69 | 'car', ... % 3 70 | 'bicycle', ... % 4 71 | 'mbike', ... % 5 72 | 'non_mot_vhcl', ... % 6 73 | 'static_person', ... % 7 74 | 'distractor', ... % 8 75 | 'occluder', ... % 9 76 | 'occluder_on_grnd', ... %10 77 | 'occluder_full', ... % 11 78 | 'reflection', ... % 12 79 | 'crowd' ... % 13 80 | }; 81 | """ 82 | 83 | 84 | def read_mot_results(filename, is_gt, is_ignore): 85 | valid_labels = {1} 86 | ignore_labels = {2, 7, 8, 12} 87 | results_dict = dict() 88 | if os.path.isfile(filename): 89 | with open(filename, 'r') as f: 90 | for line in f.readlines(): 91 | linelist = line.split(',') 92 | if len(linelist) < 7: 93 | continue 94 | fid = int(linelist[0]) 95 | if fid < 1: 96 | continue 97 | results_dict.setdefault(fid, list()) 98 | 99 | if is_gt: 100 | if 'MOT16-' in filename or 'MOT17-' in filename: 101 | label = int(float(linelist[7])) 102 | mark = int(float(linelist[6])) 103 | if mark == 0 or label not in valid_labels: 104 | continue 105 | score = 1 106 | elif is_ignore: 107 | if 'MOT16-' in filename or 'MOT17-' in filename: 108 | label = int(float(linelist[7])) 109 | vis_ratio = float(linelist[8]) 110 | if label not in ignore_labels and vis_ratio >= 0: 111 | continue 112 | else: 113 | continue 114 | score = 1 115 | else: 116 | score = float(linelist[6]) 117 | 118 | tlwh = tuple(map(float, linelist[2:6])) 119 | target_id = int(linelist[1]) 120 | 121 | results_dict[fid].append((tlwh, target_id, score)) 122 | 123 | return results_dict 124 | 125 | 126 | def unzip_objs(objs): 127 | if len(objs) > 0: 128 | tlwhs, ids, scores = zip(*objs) 129 | else: 130 | tlwhs, ids, scores = [], [], [] 131 | tlwhs = np.asarray(tlwhs, dtype=float).reshape(-1, 4) 132 | 133 | return tlwhs, ids, scores -------------------------------------------------------------------------------- /deep_sort/utils/json_logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | References: 3 | https://medium.com/analytics-vidhya/creating-a-custom-logging-mechanism-for-real-time-object-detection-using-tdd-4ca2cfcd0a2f 4 | """ 5 | import json 6 | from os import makedirs 7 | from os.path import exists, join 8 | from datetime import datetime 9 | 10 | 11 | class JsonMeta(object): 12 | HOURS = 3 13 | MINUTES = 59 14 | SECONDS = 59 15 | PATH_TO_SAVE = 'LOGS' 16 | DEFAULT_FILE_NAME = 'remaining' 17 | 18 | 19 | class BaseJsonLogger(object): 20 | """ 21 | This is the base class that returns __dict__ of its own 22 | it also returns the dicts of objects in the attributes that are list instances 23 | 24 | """ 25 | 26 | def dic(self): 27 | # returns dicts of objects 28 | out = {} 29 | for k, v in self.__dict__.items(): 30 | if hasattr(v, 'dic'): 31 | out[k] = v.dic() 32 | elif isinstance(v, list): 33 | out[k] = self.list(v) 34 | else: 35 | out[k] = v 36 | return out 37 | 38 | @staticmethod 39 | def list(values): 40 | # applies the dic method on items in the list 41 | return [v.dic() if hasattr(v, 'dic') else v for v in values] 42 | 43 | 44 | class Label(BaseJsonLogger): 45 | """ 46 | For each bounding box there are various categories with confidences. Label class keeps track of that information. 47 | """ 48 | 49 | def __init__(self, category: str, confidence: float): 50 | self.category = category 51 | self.confidence = confidence 52 | 53 | 54 | class Bbox(BaseJsonLogger): 55 | """ 56 | This module stores the information for each frame and use them in JsonParser 57 | Attributes: 58 | labels (list): List of label module. 59 | top (int): 60 | left (int): 61 | width (int): 62 | height (int): 63 | 64 | Args: 65 | bbox_id (float): 66 | top (int): 67 | left (int): 68 | width (int): 69 | height (int): 70 | 71 | References: 72 | Check Label module for better understanding. 73 | 74 | 75 | """ 76 | 77 | def __init__(self, bbox_id, top, left, width, height): 78 | self.labels = [] 79 | self.bbox_id = bbox_id 80 | self.top = top 81 | self.left = left 82 | self.width = width 83 | self.height = height 84 | 85 | def add_label(self, category, confidence): 86 | # adds category and confidence only if top_k is not exceeded. 87 | self.labels.append(Label(category, confidence)) 88 | 89 | def labels_full(self, value): 90 | return len(self.labels) == value 91 | 92 | 93 | class Frame(BaseJsonLogger): 94 | """ 95 | This module stores the information for each frame and use them in JsonParser 96 | Attributes: 97 | timestamp (float): The elapsed time of captured frame 98 | frame_id (int): The frame number of the captured video 99 | bboxes (list of Bbox objects): Stores the list of bbox objects. 100 | 101 | References: 102 | Check Bbox class for better information 103 | 104 | Args: 105 | timestamp (float): 106 | frame_id (int): 107 | 108 | """ 109 | 110 | def __init__(self, frame_id: int, timestamp: float = None): 111 | self.frame_id = frame_id 112 | self.timestamp = timestamp 113 | self.bboxes = [] 114 | 115 | def add_bbox(self, bbox_id: int, top: int, left: int, width: int, height: int): 116 | bboxes_ids = [bbox.bbox_id for bbox in self.bboxes] 117 | if bbox_id not in bboxes_ids: 118 | self.bboxes.append(Bbox(bbox_id, top, left, width, height)) 119 | else: 120 | raise ValueError("Frame with id: {} already has a Bbox with id: {}".format(self.frame_id, bbox_id)) 121 | 122 | def add_label_to_bbox(self, bbox_id: int, category: str, confidence: float): 123 | bboxes = {bbox.id: bbox for bbox in self.bboxes} 124 | if bbox_id in bboxes.keys(): 125 | res = bboxes.get(bbox_id) 126 | res.add_label(category, confidence) 127 | else: 128 | raise ValueError('the bbox with id: {} does not exists!'.format(bbox_id)) 129 | 130 | 131 | class BboxToJsonLogger(BaseJsonLogger): 132 | """ 133 | ُ This module is designed to automate the task of logging jsons. An example json is used 134 | to show the contents of json file shortly 135 | Example: 136 | { 137 | "video_details": { 138 | "frame_width": 1920, 139 | "frame_height": 1080, 140 | "frame_rate": 20, 141 | "video_name": "/home/gpu/codes/MSD/pedestrian_2/project/public/camera1.avi" 142 | }, 143 | "frames": [ 144 | { 145 | "frame_id": 329, 146 | "timestamp": 3365.1254 147 | "bboxes": [ 148 | { 149 | "labels": [ 150 | { 151 | "category": "pedestrian", 152 | "confidence": 0.9 153 | } 154 | ], 155 | "bbox_id": 0, 156 | "top": 1257, 157 | "left": 138, 158 | "width": 68, 159 | "height": 109 160 | } 161 | ] 162 | }], 163 | 164 | Attributes: 165 | frames (dict): It's a dictionary that maps each frame_id to json attributes. 166 | video_details (dict): information about video file. 167 | top_k_labels (int): shows the allowed number of labels 168 | start_time (datetime object): we use it to automate the json output by time. 169 | 170 | Args: 171 | top_k_labels (int): shows the allowed number of labels 172 | 173 | """ 174 | 175 | def __init__(self, top_k_labels: int = 1): 176 | self.frames = {} 177 | self.video_details = self.video_details = dict(frame_width=None, frame_height=None, frame_rate=None, 178 | video_name=None) 179 | self.top_k_labels = top_k_labels 180 | self.start_time = datetime.now() 181 | 182 | def set_top_k(self, value): 183 | self.top_k_labels = value 184 | 185 | def frame_exists(self, frame_id: int) -> bool: 186 | """ 187 | Args: 188 | frame_id (int): 189 | 190 | Returns: 191 | bool: true if frame_id is recognized 192 | """ 193 | return frame_id in self.frames.keys() 194 | 195 | def add_frame(self, frame_id: int, timestamp: float = None) -> None: 196 | """ 197 | Args: 198 | frame_id (int): 199 | timestamp (float): opencv captured frame time property 200 | 201 | Raises: 202 | ValueError: if frame_id would not exist in class frames attribute 203 | 204 | Returns: 205 | None 206 | 207 | """ 208 | if not self.frame_exists(frame_id): 209 | self.frames[frame_id] = Frame(frame_id, timestamp) 210 | else: 211 | raise ValueError("Frame id: {} already exists".format(frame_id)) 212 | 213 | def bbox_exists(self, frame_id: int, bbox_id: int) -> bool: 214 | """ 215 | Args: 216 | frame_id: 217 | bbox_id: 218 | 219 | Returns: 220 | bool: if bbox exists in frame bboxes list 221 | """ 222 | bboxes = [] 223 | if self.frame_exists(frame_id=frame_id): 224 | bboxes = [bbox.bbox_id for bbox in self.frames[frame_id].bboxes] 225 | return bbox_id in bboxes 226 | 227 | def find_bbox(self, frame_id: int, bbox_id: int): 228 | """ 229 | 230 | Args: 231 | frame_id: 232 | bbox_id: 233 | 234 | Returns: 235 | bbox_id (int): 236 | 237 | Raises: 238 | ValueError: if bbox_id does not exist in the bbox list of specific frame. 239 | """ 240 | if not self.bbox_exists(frame_id, bbox_id): 241 | raise ValueError("frame with id: {} does not contain bbox with id: {}".format(frame_id, bbox_id)) 242 | bboxes = {bbox.bbox_id: bbox for bbox in self.frames[frame_id].bboxes} 243 | return bboxes.get(bbox_id) 244 | 245 | def add_bbox_to_frame(self, frame_id: int, bbox_id: int, top: int, left: int, width: int, height: int) -> None: 246 | """ 247 | 248 | Args: 249 | frame_id (int): 250 | bbox_id (int): 251 | top (int): 252 | left (int): 253 | width (int): 254 | height (int): 255 | 256 | Returns: 257 | None 258 | 259 | Raises: 260 | ValueError: if bbox_id already exist in frame information with frame_id 261 | ValueError: if frame_id does not exist in frames attribute 262 | """ 263 | if self.frame_exists(frame_id): 264 | frame = self.frames[frame_id] 265 | if not self.bbox_exists(frame_id, bbox_id): 266 | frame.add_bbox(bbox_id, top, left, width, height) 267 | else: 268 | raise ValueError( 269 | "frame with frame_id: {} already contains the bbox with id: {} ".format(frame_id, bbox_id)) 270 | else: 271 | raise ValueError("frame with frame_id: {} does not exist".format(frame_id)) 272 | 273 | def add_label_to_bbox(self, frame_id: int, bbox_id: int, category: str, confidence: float): 274 | """ 275 | Args: 276 | frame_id: 277 | bbox_id: 278 | category: 279 | confidence: the confidence value returned from yolo detection 280 | 281 | Returns: 282 | None 283 | 284 | Raises: 285 | ValueError: if labels quota (top_k_labels) exceeds. 286 | """ 287 | bbox = self.find_bbox(frame_id, bbox_id) 288 | if not bbox.labels_full(self.top_k_labels): 289 | bbox.add_label(category, confidence) 290 | else: 291 | raise ValueError("labels in frame_id: {}, bbox_id: {} is fulled".format(frame_id, bbox_id)) 292 | 293 | def add_video_details(self, frame_width: int = None, frame_height: int = None, frame_rate: int = None, 294 | video_name: str = None): 295 | self.video_details['frame_width'] = frame_width 296 | self.video_details['frame_height'] = frame_height 297 | self.video_details['frame_rate'] = frame_rate 298 | self.video_details['video_name'] = video_name 299 | 300 | def output(self): 301 | output = {'video_details': self.video_details} 302 | result = list(self.frames.values()) 303 | output['frames'] = [item.dic() for item in result] 304 | return output 305 | 306 | def json_output(self, output_name): 307 | """ 308 | Args: 309 | output_name: 310 | 311 | Returns: 312 | None 313 | 314 | Notes: 315 | It creates the json output with `output_name` name. 316 | """ 317 | if not output_name.endswith('.json'): 318 | output_name += '.json' 319 | with open(output_name, 'w') as file: 320 | json.dump(self.output(), file) 321 | file.close() 322 | 323 | def set_start(self): 324 | self.start_time = datetime.now() 325 | 326 | def schedule_output_by_time(self, output_dir=JsonMeta.PATH_TO_SAVE, hours: int = 0, minutes: int = 0, 327 | seconds: int = 60) -> None: 328 | """ 329 | Notes: 330 | Creates folder and then periodically stores the jsons on that address. 331 | 332 | Args: 333 | output_dir (str): the directory where output files will be stored 334 | hours (int): 335 | minutes (int): 336 | seconds (int): 337 | 338 | Returns: 339 | None 340 | 341 | """ 342 | end = datetime.now() 343 | interval = 0 344 | interval += abs(min([hours, JsonMeta.HOURS]) * 3600) 345 | interval += abs(min([minutes, JsonMeta.MINUTES]) * 60) 346 | interval += abs(min([seconds, JsonMeta.SECONDS])) 347 | diff = (end - self.start_time).seconds 348 | 349 | if diff > interval: 350 | output_name = self.start_time.strftime('%Y-%m-%d %H-%M-%S') + '.json' 351 | if not exists(output_dir): 352 | makedirs(output_dir) 353 | output = join(output_dir, output_name) 354 | self.json_output(output_name=output) 355 | self.frames = {} 356 | self.start_time = datetime.now() 357 | 358 | def schedule_output_by_frames(self, frames_quota, frame_counter, output_dir=JsonMeta.PATH_TO_SAVE): 359 | """ 360 | saves as the number of frames quota increases higher. 361 | :param frames_quota: 362 | :param frame_counter: 363 | :param output_dir: 364 | :return: 365 | """ 366 | pass 367 | 368 | def flush(self, output_dir): 369 | """ 370 | Notes: 371 | We use this function to output jsons whenever possible. 372 | like the time that we exit the while loop of opencv. 373 | 374 | Args: 375 | output_dir: 376 | 377 | Returns: 378 | None 379 | 380 | """ 381 | filename = self.start_time.strftime('%Y-%m-%d %H-%M-%S') + '-remaining.json' 382 | output = join(output_dir, filename) 383 | self.json_output(output_name=output) 384 | -------------------------------------------------------------------------------- /deep_sort/utils/log.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | 4 | def get_logger(name='root'): 5 | formatter = logging.Formatter( 6 | # fmt='%(asctime)s [%(levelname)s]: %(filename)s(%(funcName)s:%(lineno)s) >> %(message)s') 7 | fmt='%(asctime)s [%(levelname)s]: %(message)s', datefmt='%Y-%m-%d %H:%M:%S') 8 | 9 | handler = logging.StreamHandler() 10 | handler.setFormatter(formatter) 11 | 12 | logger = logging.getLogger(name) 13 | logger.setLevel(logging.INFO) 14 | logger.addHandler(handler) 15 | return logger 16 | 17 | 18 | -------------------------------------------------------------------------------- /deep_sort/utils/parser.py: -------------------------------------------------------------------------------- 1 | import os 2 | import yaml 3 | from easydict import EasyDict as edict 4 | 5 | class YamlParser(edict): 6 | """ 7 | This is yaml parser based on EasyDict. 8 | """ 9 | def __init__(self, cfg_dict=None, config_file=None): 10 | if cfg_dict is None: 11 | cfg_dict = {} 12 | 13 | if config_file is not None: 14 | assert(os.path.isfile(config_file)) 15 | with open(config_file, 'r') as fo: 16 | cfg_dict.update(yaml.load(fo.read())) 17 | 18 | super(YamlParser, self).__init__(cfg_dict) 19 | 20 | 21 | def merge_from_file(self, config_file): 22 | with open(config_file, 'r') as fo: 23 | self.update(yaml.load(fo.read())) 24 | 25 | 26 | def merge_from_dict(self, config_dict): 27 | self.update(config_dict) 28 | 29 | 30 | def get_config(config_file=None): 31 | return YamlParser(config_file=config_file) 32 | 33 | 34 | if __name__ == "__main__": 35 | cfg = YamlParser(config_file="../configs/yolov3.yaml") 36 | cfg.merge_from_file("../configs/deep_sort.yaml") 37 | 38 | import ipdb; ipdb.set_trace() -------------------------------------------------------------------------------- /deep_sort/utils/tools.py: -------------------------------------------------------------------------------- 1 | from functools import wraps 2 | from time import time 3 | 4 | 5 | def is_video(ext: str): 6 | """ 7 | Returns true if ext exists in 8 | allowed_exts for video files. 9 | 10 | Args: 11 | ext: 12 | 13 | Returns: 14 | 15 | """ 16 | 17 | allowed_exts = ('.mp4', '.webm', '.ogg', '.avi', '.wmv', '.mkv', '.3gp') 18 | return any((ext.endswith(x) for x in allowed_exts)) 19 | 20 | 21 | def tik_tok(func): 22 | """ 23 | keep track of time for each process. 24 | Args: 25 | func: 26 | 27 | Returns: 28 | 29 | """ 30 | @wraps(func) 31 | def _time_it(*args, **kwargs): 32 | start = time() 33 | try: 34 | return func(*args, **kwargs) 35 | finally: 36 | end_ = time() 37 | print("time: {:.03f}s, fps: {:.03f}".format(end_ - start, 1 / (end_ - start))) 38 | 39 | return _time_it 40 | -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | from AIDetector_pytorch import Detector 2 | import imutils 3 | import cv2 4 | 5 | def main(): 6 | 7 | func_status = {} 8 | func_status['headpose'] = None 9 | 10 | name = 'demo' 11 | 12 | det = Detector() 13 | cap = cv2.VideoCapture('E:/视频/行人监控/test01.mp4') 14 | fps = int(cap.get(5)) 15 | print('fps:', fps) 16 | t = int(1000/fps) 17 | 18 | size = None 19 | videoWriter = None 20 | 21 | while True: 22 | 23 | # try: 24 | _, im = cap.read() 25 | if im is None: 26 | break 27 | 28 | result = det.feedCap(im, func_status) 29 | result = result['frame'] 30 | result = imutils.resize(result, height=500) 31 | if videoWriter is None: 32 | fourcc = cv2.VideoWriter_fourcc( 33 | 'm', 'p', '4', 'v') # opencv3.0 34 | videoWriter = cv2.VideoWriter( 35 | 'result.mp4', fourcc, fps, (result.shape[1], result.shape[0])) 36 | 37 | videoWriter.write(result) 38 | cv2.imshow(name, result) 39 | cv2.waitKey(t) 40 | 41 | if cv2.getWindowProperty(name, cv2.WND_PROP_AUTOSIZE) < 1: 42 | # 点x退出 43 | break 44 | # except Exception as e: 45 | # print(e) 46 | # break 47 | 48 | cap.release() 49 | videoWriter.release() 50 | cv2.destroyAllWindows() 51 | 52 | if __name__ == '__main__': 53 | 54 | main() -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/models/__init__.py -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/common.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/models/__pycache__/common.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/experimental.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/models/__pycache__/experimental.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/yolo.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/models/__pycache__/yolo.cpython-37.pyc -------------------------------------------------------------------------------- /models/common.py: -------------------------------------------------------------------------------- 1 | # This file contains modules common to various models 2 | 3 | import math 4 | 5 | import torch 6 | import torch.nn as nn 7 | 8 | from utils.general import non_max_suppression 9 | 10 | def autopad(k, p=None): # kernel, padding 11 | # Pad to 'same' 12 | if p is None: 13 | p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad 14 | return p 15 | 16 | 17 | def DWConv(c1, c2, k=1, s=1, act=True): 18 | # Depthwise convolution 19 | return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) 20 | 21 | 22 | class Conv(nn.Module): 23 | # Standard convolution 24 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 25 | super(Conv, self).__init__() 26 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) 27 | self.bn = nn.BatchNorm2d(c2) 28 | self.act = nn.Hardswish() if act else nn.Identity() 29 | 30 | def forward(self, x): 31 | return self.act(self.bn(self.conv(x))) 32 | 33 | def fuseforward(self, x): 34 | return self.act(self.conv(x)) 35 | 36 | 37 | class Bottleneck(nn.Module): 38 | # Standard bottleneck 39 | def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion 40 | super(Bottleneck, self).__init__() 41 | c_ = int(c2 * e) # hidden channels 42 | self.cv1 = Conv(c1, c_, 1, 1) 43 | self.cv2 = Conv(c_, c2, 3, 1, g=g) 44 | self.add = shortcut and c1 == c2 45 | 46 | def forward(self, x): 47 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 48 | 49 | 50 | class BottleneckCSP(nn.Module): 51 | # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks 52 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 53 | super(BottleneckCSP, self).__init__() 54 | c_ = int(c2 * e) # hidden channels 55 | self.cv1 = Conv(c1, c_, 1, 1) 56 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 57 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 58 | self.cv4 = Conv(2 * c_, c2, 1, 1) 59 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 60 | self.act = nn.LeakyReLU(0.1, inplace=True) 61 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) 62 | 63 | def forward(self, x): 64 | y1 = self.cv3(self.m(self.cv1(x))) 65 | y2 = self.cv2(x) 66 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 67 | 68 | 69 | class SPP(nn.Module): 70 | # Spatial pyramid pooling layer used in YOLOv3-SPP 71 | def __init__(self, c1, c2, k=(5, 9, 13)): 72 | super(SPP, self).__init__() 73 | c_ = c1 // 2 # hidden channels 74 | self.cv1 = Conv(c1, c_, 1, 1) 75 | self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) 76 | self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) 77 | 78 | def forward(self, x): 79 | x = self.cv1(x) 80 | return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) 81 | 82 | 83 | class Focus(nn.Module): 84 | # Focus wh information into c-space 85 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 86 | super(Focus, self).__init__() 87 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 88 | 89 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) 90 | return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) 91 | 92 | 93 | class Concat(nn.Module): 94 | # Concatenate a list of tensors along dimension 95 | def __init__(self, dimension=1): 96 | super(Concat, self).__init__() 97 | self.d = dimension 98 | 99 | def forward(self, x): 100 | return torch.cat(x, self.d) 101 | 102 | 103 | class NMS(nn.Module): 104 | # Non-Maximum Suppression (NMS) module 105 | conf = 0.25 # confidence threshold 106 | iou = 0.45 # IoU threshold 107 | classes = None # (optional list) filter by class 108 | 109 | def __init__(self): 110 | super(NMS, self).__init__() 111 | 112 | def forward(self, x): 113 | return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) 114 | 115 | -------------------------------------------------------------------------------- /models/experimental.py: -------------------------------------------------------------------------------- 1 | # This file contains experimental modules 2 | 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | 7 | from models.common import Conv, DWConv 8 | from utils.google_utils import attempt_download 9 | 10 | 11 | class CrossConv(nn.Module): 12 | # Cross Convolution Downsample 13 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 14 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 15 | super(CrossConv, self).__init__() 16 | c_ = int(c2 * e) # hidden channels 17 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 18 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 19 | self.add = shortcut and c1 == c2 20 | 21 | def forward(self, x): 22 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 23 | 24 | 25 | class C3(nn.Module): 26 | # Cross Convolution CSP 27 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 28 | super(C3, self).__init__() 29 | c_ = int(c2 * e) # hidden channels 30 | self.cv1 = Conv(c1, c_, 1, 1) 31 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 32 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 33 | self.cv4 = Conv(2 * c_, c2, 1, 1) 34 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 35 | self.act = nn.LeakyReLU(0.1, inplace=True) 36 | self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)]) 37 | 38 | def forward(self, x): 39 | y1 = self.cv3(self.m(self.cv1(x))) 40 | y2 = self.cv2(x) 41 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 42 | 43 | 44 | class Sum(nn.Module): 45 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 46 | def __init__(self, n, weight=False): # n: number of inputs 47 | super(Sum, self).__init__() 48 | self.weight = weight # apply weights boolean 49 | self.iter = range(n - 1) # iter object 50 | if weight: 51 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights 52 | 53 | def forward(self, x): 54 | y = x[0] # no weight 55 | if self.weight: 56 | w = torch.sigmoid(self.w) * 2 57 | for i in self.iter: 58 | y = y + x[i + 1] * w[i] 59 | else: 60 | for i in self.iter: 61 | y = y + x[i + 1] 62 | return y 63 | 64 | 65 | class GhostConv(nn.Module): 66 | # Ghost Convolution https://github.com/huawei-noah/ghostnet 67 | def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups 68 | super(GhostConv, self).__init__() 69 | c_ = c2 // 2 # hidden channels 70 | self.cv1 = Conv(c1, c_, k, s, None, g, act) 71 | self.cv2 = Conv(c_, c_, 5, 1, None, c_, act) 72 | 73 | def forward(self, x): 74 | y = self.cv1(x) 75 | return torch.cat([y, self.cv2(y)], 1) 76 | 77 | 78 | class GhostBottleneck(nn.Module): 79 | # Ghost Bottleneck https://github.com/huawei-noah/ghostnet 80 | def __init__(self, c1, c2, k, s): 81 | super(GhostBottleneck, self).__init__() 82 | c_ = c2 // 2 83 | self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw 84 | DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw 85 | GhostConv(c_, c2, 1, 1, act=False)) # pw-linear 86 | self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), 87 | Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() 88 | 89 | def forward(self, x): 90 | return self.conv(x) + self.shortcut(x) 91 | 92 | 93 | class MixConv2d(nn.Module): 94 | # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595 95 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): 96 | super(MixConv2d, self).__init__() 97 | groups = len(k) 98 | if equal_ch: # equal c_ per group 99 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices 100 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels 101 | else: # equal weight.numel() per group 102 | b = [c2] + [0] * groups 103 | a = np.eye(groups + 1, groups, k=-1) 104 | a -= np.roll(a, 1, axis=1) 105 | a *= np.array(k) ** 2 106 | a[0] = 1 107 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 108 | 109 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) 110 | self.bn = nn.BatchNorm2d(c2) 111 | self.act = nn.LeakyReLU(0.1, inplace=True) 112 | 113 | def forward(self, x): 114 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 115 | 116 | 117 | class Ensemble(nn.ModuleList): 118 | # Ensemble of models 119 | def __init__(self): 120 | super(Ensemble, self).__init__() 121 | 122 | def forward(self, x, augment=False): 123 | y = [] 124 | for module in self: 125 | y.append(module(x, augment)[0]) 126 | # y = torch.stack(y).max(0)[0] # max ensemble 127 | # y = torch.cat(y, 1) # nms ensemble 128 | y = torch.stack(y).mean(0) # mean ensemble 129 | return y, None # inference, train output 130 | 131 | 132 | def attempt_load(weights, map_location=None): 133 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 134 | model = Ensemble() 135 | for w in weights if isinstance(weights, list) else [weights]: 136 | attempt_download(w) 137 | try: 138 | model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model 139 | except Exception: 140 | model.append(torch.load(w, map_location=map_location).float().fuse().eval()) # load FP32 model 141 | 142 | # Compatibility updates 143 | for m in model.modules(): 144 | if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 145 | m.inplace = True # pytorch 1.7.0 compatibility 146 | elif type(m) is Conv: 147 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 148 | 149 | if len(model) == 1: 150 | return model[-1] # return model 151 | else: 152 | print('Ensemble created with %s\n' % weights) 153 | for k in ['names', 'stride']: 154 | setattr(model, k, getattr(model[-1], k)) 155 | return model # return ensemble 156 | -------------------------------------------------------------------------------- /models/yolo.py: -------------------------------------------------------------------------------- 1 | import math 2 | import sys 3 | from copy import deepcopy 4 | from pathlib import Path 5 | 6 | sys.path.append('./') # to run '$ python *.py' files in subdirectories 7 | 8 | import torch 9 | import torch.nn as nn 10 | 11 | from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, NMS 12 | from models.experimental import MixConv2d, CrossConv, C3 13 | from utils.autoanchor import check_anchor_order 14 | from utils.general import make_divisible 15 | from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights 16 | 17 | thop = None 18 | 19 | 20 | class Detect(nn.Module): 21 | stride = None # strides computed during build 22 | export = False # onnx export 23 | 24 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 25 | super(Detect, self).__init__() 26 | self.nc = nc # number of classes 27 | self.no = nc + 5 # number of outputs per anchor 28 | self.nl = len(anchors) # number of detection layers 29 | self.na = len(anchors[0]) // 2 # number of anchors 30 | self.grid = [torch.zeros(1)] * self.nl # init grid 31 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 32 | self.register_buffer('anchors', a) # shape(nl,na,2) 33 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 34 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 35 | 36 | def forward(self, x): 37 | # x = x.copy() # for profiling 38 | z = [] # inference output 39 | self.training |= self.export 40 | for i in range(self.nl): 41 | x[i] = self.m[i](x[i]) # conv 42 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 43 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 44 | 45 | if not self.training: # inference 46 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 47 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 48 | 49 | y = x[i].sigmoid() 50 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 51 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 52 | z.append(y.view(bs, -1, self.no)) 53 | 54 | return x if self.training else (torch.cat(z, 1), x) 55 | 56 | @staticmethod 57 | def _make_grid(nx=20, ny=20): 58 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 59 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 60 | 61 | 62 | class Model(nn.Module): 63 | def __init__(self, cfg='yolov5m.yaml', ch=3, nc=None): # model, input channels, number of classes 64 | super(Model, self).__init__() 65 | if isinstance(cfg, dict): 66 | self.yaml = cfg # model dict 67 | else: # is *.yaml 68 | import yaml # for torch hub 69 | self.yaml_file = Path(cfg).name 70 | with open(cfg) as f: 71 | self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict 72 | 73 | # Define model 74 | if nc and nc != self.yaml['nc']: 75 | self.yaml['nc'] = nc # override yaml value 76 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out 77 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 78 | 79 | # Build strides, anchors 80 | m = self.model[-1] # Detect() 81 | if isinstance(m, Detect): 82 | s = 128 # 2x min stride 83 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 84 | m.anchors /= m.stride.view(-1, 1, 1) 85 | check_anchor_order(m) 86 | self.stride = m.stride 87 | self._initialize_biases() # only run once 88 | # print('Strides: %s' % m.stride.tolist()) 89 | 90 | # Init weights, biases 91 | initialize_weights(self) 92 | self.info() 93 | 94 | def forward(self, x, augment=False, profile=False): 95 | if augment: 96 | img_size = x.shape[-2:] # height, width 97 | s = [1, 0.83, 0.67] # scales 98 | f = [None, 3, None] # flips (2-ud, 3-lr) 99 | y = [] # outputs 100 | for si, fi in zip(s, f): 101 | xi = scale_img(x.flip(fi) if fi else x, si) 102 | yi = self.forward_once(xi)[0] # forward 103 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 104 | yi[..., :4] /= si # de-scale 105 | if fi == 2: 106 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 107 | elif fi == 3: 108 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 109 | y.append(yi) 110 | return torch.cat(y, 1), None # augmented inference, train 111 | else: 112 | return self.forward_once(x, profile) # single-scale inference, train 113 | 114 | def forward_once(self, x, profile=False): 115 | y, dt = [], [] # outputs 116 | for m in self.model: 117 | if m.f != -1: # if not from previous layer 118 | x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers 119 | 120 | if profile: 121 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPS 122 | t = time_synchronized() 123 | for _ in range(10): 124 | _ = m(x) 125 | dt.append((time_synchronized() - t) * 100) 126 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 127 | 128 | x = m(x) # run 129 | y.append(x if m.i in self.save else None) # save output 130 | 131 | if profile: 132 | print('%.1fms total' % sum(dt)) 133 | return x 134 | 135 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 136 | # https://arxiv.org/abs/1708.02002 section 3.3 137 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 138 | m = self.model[-1] # Detect() module 139 | for mi, s in zip(m.m, m.stride): # from 140 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 141 | b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 142 | b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 143 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 144 | 145 | def _print_biases(self): 146 | m = self.model[-1] # Detect() module 147 | for mi in m.m: # from 148 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 149 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 150 | 151 | # def _print_weights(self): 152 | # for m in self.model.modules(): 153 | # if type(m) is Bottleneck: 154 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 155 | 156 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 157 | print('Fusing layers... ') 158 | for m in self.model.modules(): 159 | if type(m) is Conv and hasattr(m, 'bn'): 160 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 161 | delattr(m, 'bn') # remove batchnorm 162 | m.forward = m.fuseforward # update forward 163 | self.info() 164 | return self 165 | 166 | def nms(self, mode=True): # add or remove NMS module 167 | present = type(self.model[-1]) is NMS # last layer is NMS 168 | if mode and not present: 169 | print('Adding NMS... ') 170 | m = NMS() # module 171 | m.f = -1 # from 172 | m.i = self.model[-1].i + 1 # index 173 | self.model.add_module(name='%s' % m.i, module=m) # add 174 | self.eval() 175 | elif not mode and present: 176 | print('Removing NMS... ') 177 | self.model = self.model[:-1] # remove 178 | return self 179 | 180 | def info(self, verbose=False, img_size=640): # print model information 181 | model_info(self, verbose, img_size) 182 | 183 | 184 | def parse_model(d, ch): # model_dict, input_channels(3) 185 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 186 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 187 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 188 | 189 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 190 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 191 | m = eval(m) if isinstance(m, str) else m # eval strings 192 | for j, a in enumerate(args): 193 | try: 194 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 195 | except: 196 | pass 197 | 198 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 199 | if m in [Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: 200 | c1, c2 = ch[f], args[0] 201 | 202 | # Normal 203 | # if i > 0 and args[0] != no: # channel expansion factor 204 | # ex = 1.75 # exponential (default 2.0) 205 | # e = math.log(c2 / ch[1]) / math.log(2) 206 | # c2 = int(ch[1] * ex ** e) 207 | # if m != Focus: 208 | 209 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 210 | 211 | # Experimental 212 | # if i > 0 and args[0] != no: # channel expansion factor 213 | # ex = 1 + gw # exponential (default 2.0) 214 | # ch1 = 32 # ch[1] 215 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n 216 | # c2 = int(ch1 * ex ** e) 217 | # if m != Focus: 218 | # c2 = make_divisible(c2, 8) if c2 != no else c2 219 | 220 | args = [c1, c2, *args[1:]] 221 | if m in [BottleneckCSP, C3]: 222 | args.insert(2, n) 223 | n = 1 224 | elif m is nn.BatchNorm2d: 225 | args = [ch[f]] 226 | elif m is Concat: 227 | c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) 228 | elif m is Detect: 229 | args.append([ch[x + 1] for x in f]) 230 | if isinstance(args[1], int): # number of anchors 231 | args[1] = [list(range(args[1] * 2))] * len(f) 232 | else: 233 | c2 = ch[f] 234 | 235 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 236 | t = str(m)[8:-2].replace('__main__.', '') # module type 237 | np = sum([x.numel() for x in m_.parameters()]) # number params 238 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 239 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 240 | layers.append(m_) 241 | ch.append(c2) 242 | return nn.Sequential(*layers), sorted(save) 243 | 244 | 245 | -------------------------------------------------------------------------------- /models/yolov5m.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 8 # number of classes 3 | depth_multiple: 0.67 # model depth multiple 4 | width_multiple: 0.75 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # base ---------------------------------------- 4 | Cython 5 | matplotlib>=3.2.2 6 | numpy>=1.18.5 7 | opencv-python>=4.1.2 8 | Pillow 9 | PyYAML>=5.3.1 10 | scipy>=1.4.1 11 | tensorboard>=2.2 12 | torch>=1.7.0 13 | torchvision>=0.8.1 14 | tqdm>=4.41.0 15 | 16 | # logging ------------------------------------- 17 | # wandb 18 | 19 | # plotting ------------------------------------ 20 | seaborn>=0.11.0 21 | pandas 22 | 23 | # export -------------------------------------- 24 | # coremltools>=4.1 25 | # onnx>=1.8.1 26 | # scikit-learn==0.19.2 # for coreml quantization 27 | 28 | # extras -------------------------------------- 29 | thop # FLOPS computation 30 | pycocotools>=2.0 # COCO mAP 31 | -------------------------------------------------------------------------------- /tracker.py: -------------------------------------------------------------------------------- 1 | from deep_sort.utils.parser import get_config 2 | from deep_sort.deep_sort import DeepSort 3 | import torch 4 | import cv2 5 | 6 | palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) 7 | cfg = get_config() 8 | cfg.merge_from_file("deep_sort/configs/deep_sort.yaml") 9 | deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT, 10 | max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE, 11 | nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE, 12 | max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, 13 | use_cuda=True) 14 | 15 | 16 | def plot_bboxes(image, bboxes, line_thickness=None): 17 | # Plots one bounding box on image img 18 | tl = line_thickness or round( 19 | 0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 # line/font thickness 20 | for (x1, y1, x2, y2, cls_id, pos_id) in bboxes: 21 | if cls_id in ['smoke', 'phone', 'eat']: 22 | color = (0, 0, 255) 23 | else: 24 | color = (0, 255, 0) 25 | if cls_id == 'eat': 26 | cls_id = 'eat-drink' 27 | c1, c2 = (x1, y1), (x2, y2) 28 | cv2.rectangle(image, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) 29 | tf = max(tl - 1, 1) # font thickness 30 | t_size = cv2.getTextSize(cls_id, 0, fontScale=tl / 3, thickness=tf)[0] 31 | c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 32 | cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA) # filled 33 | cv2.putText(image, '{} ID-{}'.format(cls_id, pos_id), (c1[0], c1[1] - 2), 0, tl / 3, 34 | [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) 35 | 36 | return image 37 | 38 | 39 | def update_tracker(target_detector, image): 40 | 41 | new_faces = [] 42 | _, bboxes = target_detector.detect(image) 43 | 44 | bbox_xywh = [] 45 | confs = [] 46 | bboxes2draw = [] 47 | face_bboxes = [] 48 | if len(bboxes): 49 | 50 | # Adapt detections to deep sort input format 51 | for x1, y1, x2, y2, _, conf in bboxes: 52 | 53 | obj = [ 54 | int((x1+x2)/2), int((y1+y2)/2), 55 | x2-x1, y2-y1 56 | ] 57 | bbox_xywh.append(obj) 58 | confs.append(conf) 59 | 60 | xywhs = torch.Tensor(bbox_xywh) 61 | confss = torch.Tensor(confs) 62 | 63 | # Pass detections to deepsort 64 | outputs = deepsort.update(xywhs, confss, image) 65 | 66 | for value in list(outputs): 67 | x1,y1,x2,y2,track_id = value 68 | bboxes2draw.append( 69 | (x1, y1, x2, y2, '', track_id) 70 | ) 71 | 72 | image = plot_bboxes(image, bboxes2draw) 73 | 74 | return image, new_faces, face_bboxes 75 | -------------------------------------------------------------------------------- /utils/BaseDetector.py: -------------------------------------------------------------------------------- 1 | from tracker import update_tracker 2 | import cv2 3 | 4 | 5 | class baseDet(object): 6 | 7 | def __init__(self): 8 | 9 | self.img_size = 640 10 | self.threshold = 0.3 11 | self.stride = 1 12 | 13 | def build_config(self): 14 | 15 | self.faceTracker = {} 16 | self.faceClasses = {} 17 | self.faceLocation1 = {} 18 | self.faceLocation2 = {} 19 | self.frameCounter = 0 20 | self.currentCarID = 0 21 | self.recorded = [] 22 | 23 | self.font = cv2.FONT_HERSHEY_SIMPLEX 24 | 25 | def feedCap(self, im, func_status): 26 | 27 | retDict = { 28 | 'frame': None, 29 | 'faces': None, 30 | 'list_of_ids': None, 31 | 'face_bboxes': [] 32 | } 33 | self.frameCounter += 1 34 | 35 | im, faces, face_bboxes = update_tracker(self, im) 36 | 37 | retDict['frame'] = im 38 | retDict['faces'] = faces 39 | retDict['face_bboxes'] = face_bboxes 40 | 41 | return retDict 42 | 43 | def init_model(self): 44 | raise EOFError("Undefined model type.") 45 | 46 | def preprocess(self): 47 | raise EOFError("Undefined model type.") 48 | 49 | def detect(self): 50 | raise EOFError("Undefined model type.") 51 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__init__.py -------------------------------------------------------------------------------- /utils/__pycache__/BaseDetector.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/BaseDetector.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/autoanchor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/autoanchor.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/datasets.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/datasets.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/general.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/general.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/general.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/general.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/google_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/google_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/google_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/google_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/loss.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/metrics.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/metrics.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plots.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/plots.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/torch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/torch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/torch_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/utils/__pycache__/torch_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/activations.py: -------------------------------------------------------------------------------- 1 | # Activation functions 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | # Swish https://arxiv.org/pdf/1905.02244.pdf --------------------------------------------------------------------------- 9 | class Swish(nn.Module): # 10 | @staticmethod 11 | def forward(x): 12 | return x * torch.sigmoid(x) 13 | 14 | 15 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish() 16 | @staticmethod 17 | def forward(x): 18 | # return x * F.hardsigmoid(x) # for torchscript and CoreML 19 | return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX 20 | 21 | 22 | class MemoryEfficientSwish(nn.Module): 23 | class F(torch.autograd.Function): 24 | @staticmethod 25 | def forward(ctx, x): 26 | ctx.save_for_backward(x) 27 | return x * torch.sigmoid(x) 28 | 29 | @staticmethod 30 | def backward(ctx, grad_output): 31 | x = ctx.saved_tensors[0] 32 | sx = torch.sigmoid(x) 33 | return grad_output * (sx * (1 + x * (1 - sx))) 34 | 35 | def forward(self, x): 36 | return self.F.apply(x) 37 | 38 | 39 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 40 | class Mish(nn.Module): 41 | @staticmethod 42 | def forward(x): 43 | return x * F.softplus(x).tanh() 44 | 45 | 46 | class MemoryEfficientMish(nn.Module): 47 | class F(torch.autograd.Function): 48 | @staticmethod 49 | def forward(ctx, x): 50 | ctx.save_for_backward(x) 51 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 52 | 53 | @staticmethod 54 | def backward(ctx, grad_output): 55 | x = ctx.saved_tensors[0] 56 | sx = torch.sigmoid(x) 57 | fx = F.softplus(x).tanh() 58 | return grad_output * (fx + x * sx * (1 - fx * fx)) 59 | 60 | def forward(self, x): 61 | return self.F.apply(x) 62 | 63 | 64 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 65 | class FReLU(nn.Module): 66 | def __init__(self, c1, k=3): # ch_in, kernel 67 | super().__init__() 68 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1) 69 | self.bn = nn.BatchNorm2d(c1) 70 | 71 | def forward(self, x): 72 | return torch.max(x, self.bn(self.conv(x))) 73 | -------------------------------------------------------------------------------- /utils/autoanchor.py: -------------------------------------------------------------------------------- 1 | # Auto-anchor utils 2 | 3 | import numpy as np 4 | import torch 5 | import yaml 6 | from scipy.cluster.vq import kmeans 7 | from tqdm import tqdm 8 | 9 | 10 | def check_anchor_order(m): 11 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary 12 | a = m.anchor_grid.prod(-1).view(-1) # anchor area 13 | da = a[-1] - a[0] # delta a 14 | ds = m.stride[-1] - m.stride[0] # delta s 15 | if da.sign() != ds.sign(): # same order 16 | print('Reversing anchor order') 17 | m.anchors[:] = m.anchors.flip(0) 18 | m.anchor_grid[:] = m.anchor_grid.flip(0) 19 | 20 | 21 | def check_anchors(dataset, model, thr=4.0, imgsz=640): 22 | # Check anchor fit to data, recompute if necessary 23 | print('\nAnalyzing anchors... ', end='') 24 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() 25 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) 26 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale 27 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh 28 | 29 | def metric(k): # compute metric 30 | r = wh[:, None] / k[None] 31 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 32 | best = x.max(1)[0] # best_x 33 | aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold 34 | bpr = (best > 1. / thr).float().mean() # best possible recall 35 | return bpr, aat 36 | 37 | bpr, aat = metric(m.anchor_grid.clone().cpu().view(-1, 2)) 38 | print('anchors/target = %.2f, Best Possible Recall (BPR) = %.4f' % (aat, bpr), end='') 39 | if bpr < 0.98: # threshold to recompute 40 | print('. Attempting to improve anchors, please wait...') 41 | na = m.anchor_grid.numel() // 2 # number of anchors 42 | new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) 43 | new_bpr = metric(new_anchors.reshape(-1, 2))[0] 44 | if new_bpr > bpr: # replace anchors 45 | new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors) 46 | m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference 47 | m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss 48 | check_anchor_order(m) 49 | print('New anchors saved to model. Update model *.yaml to use these anchors in the future.') 50 | else: 51 | print('Original anchors better than new anchors. Proceeding with original anchors.') 52 | print('') # newline 53 | 54 | 55 | def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): 56 | """ Creates kmeans-evolved anchors from training dataset 57 | 58 | Arguments: 59 | path: path to dataset *.yaml, or a loaded dataset 60 | n: number of anchors 61 | img_size: image size used for training 62 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 63 | gen: generations to evolve anchors using genetic algorithm 64 | verbose: print all results 65 | 66 | Return: 67 | k: kmeans evolved anchors 68 | 69 | Usage: 70 | from utils.autoanchor import *; _ = kmean_anchors() 71 | """ 72 | thr = 1. / thr 73 | 74 | def metric(k, wh): # compute metrics 75 | r = wh[:, None] / k[None] 76 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 77 | # x = wh_iou(wh, torch.tensor(k)) # iou metric 78 | return x, x.max(1)[0] # x, best_x 79 | 80 | def anchor_fitness(k): # mutation fitness 81 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh) 82 | return (best * (best > thr).float()).mean() # fitness 83 | 84 | def print_results(k): 85 | k = k[np.argsort(k.prod(1))] # sort small to large 86 | x, best = metric(k, wh0) 87 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr 88 | print('thr=%.2f: %.4f best possible recall, %.2f anchors past thr' % (thr, bpr, aat)) 89 | print('n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thr=%.3f-mean: ' % 90 | (n, img_size, x.mean(), best.mean(), x[x > thr].mean()), end='') 91 | for i, x in enumerate(k): 92 | print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg 93 | return k 94 | 95 | if isinstance(path, str): # *.yaml file 96 | with open(path) as f: 97 | data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict 98 | from utils.datasets import LoadImagesAndLabels 99 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) 100 | else: 101 | dataset = path # dataset 102 | 103 | # Get label wh 104 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) 105 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh 106 | 107 | # Filter 108 | i = (wh0 < 3.0).any(1).sum() 109 | if i: 110 | print('WARNING: Extremely small objects found. ' 111 | '%g of %g labels are < 3 pixels in width or height.' % (i, len(wh0))) 112 | wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels 113 | 114 | # Kmeans calculation 115 | print('Running kmeans for %g anchors on %g points...' % (n, len(wh))) 116 | s = wh.std(0) # sigmas for whitening 117 | k, dist = kmeans(wh / s, n, iter=30) # points, mean distance 118 | k *= s 119 | wh = torch.tensor(wh, dtype=torch.float32) # filtered 120 | wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered 121 | k = print_results(k) 122 | 123 | # Plot 124 | # k, d = [None] * 20, [None] * 20 125 | # for i in tqdm(range(1, 21)): 126 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance 127 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) 128 | # ax = ax.ravel() 129 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') 130 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh 131 | # ax[0].hist(wh[wh[:, 0]<100, 0],400) 132 | # ax[1].hist(wh[wh[:, 1]<100, 1],400) 133 | # fig.savefig('wh.png', dpi=200) 134 | 135 | # Evolve 136 | npr = np.random 137 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma 138 | pbar = tqdm(range(gen), desc='Evolving anchors with Genetic Algorithm') # progress bar 139 | for _ in pbar: 140 | v = np.ones(sh) 141 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates) 142 | v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) 143 | kg = (k.copy() * v).clip(min=2.0) 144 | fg = anchor_fitness(kg) 145 | if fg > f: 146 | f, k = fg, kg.copy() 147 | pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f 148 | if verbose: 149 | print_results(k) 150 | 151 | return print_results(k) 152 | -------------------------------------------------------------------------------- /utils/general.py: -------------------------------------------------------------------------------- 1 | # General utils 2 | 3 | import glob 4 | import os 5 | import platform 6 | import random 7 | import re 8 | import subprocess 9 | import time 10 | from pathlib import Path 11 | 12 | import cv2 13 | import math 14 | import numpy as np 15 | import torch 16 | import torchvision 17 | import yaml 18 | 19 | from utils.google_utils import gsutil_getsize 20 | from utils.metrics import fitness 21 | from utils.torch_utils import init_torch_seeds 22 | 23 | # Settings 24 | torch.set_printoptions(linewidth=320, precision=5, profile='long') 25 | # format short g, %precision=5 26 | np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) 27 | # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) 28 | cv2.setNumThreads(0) 29 | 30 | 31 | def init_seeds(seed=0): 32 | random.seed(seed) 33 | np.random.seed(seed) 34 | init_torch_seeds(seed) 35 | 36 | 37 | def get_latest_run(search_dir='.'): 38 | # Return path to most recent 'last.pt' in /runs (i.e. to --resume from) 39 | last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True) 40 | return max(last_list, key=os.path.getctime) if last_list else '' 41 | 42 | 43 | def check_git_status(): 44 | # Suggest 'git pull' if repo is out of date 45 | if platform.system() in ['Linux', 'Darwin'] and not os.path.isfile('/.dockerenv'): 46 | s = subprocess.check_output( 47 | 'if [ -d .git ]; then git fetch && git status -uno; fi', shell=True).decode('utf-8') 48 | if 'Your branch is behind' in s: 49 | print(s[s.find('Your branch is behind'):s.find('\n\n')] + '\n') 50 | 51 | 52 | def check_img_size(img_size, s=32): 53 | # Verify img_size is a multiple of stride s 54 | new_size = make_divisible(img_size, int(s)) # ceil gs-multiple 55 | if new_size != img_size: 56 | print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % 57 | (img_size, s, new_size)) 58 | return new_size 59 | 60 | 61 | def check_file(file): 62 | # Search for file if not found 63 | if os.path.isfile(file) or file == '': 64 | return file 65 | else: 66 | files = glob.glob('./**/' + file, recursive=True) # find file 67 | assert len(files), 'File Not Found: %s' % file # assert file was found 68 | assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % ( 69 | file, files) # assert unique 70 | return files[0] # return file 71 | 72 | 73 | def check_dataset(dict): 74 | # Download dataset if not found locally 75 | val, s = dict.get('val'), dict.get('download') 76 | if val and len(val): 77 | val = [Path(x).resolve() 78 | for x in (val if isinstance(val, list) else [val])] # val path 79 | if not all(x.exists() for x in val): 80 | print('\nWARNING: Dataset not found, nonexistent paths: %s' % 81 | [str(x) for x in val if not x.exists()]) 82 | if s and len(s): # download script 83 | print('Downloading %s ...' % s) 84 | if s.startswith('http') and s.endswith('.zip'): # URL 85 | f = Path(s).name # filename 86 | torch.hub.download_url_to_file(s, f) 87 | r = os.system('unzip -q %s -d ../ && rm %s' % 88 | (f, f)) # unzip 89 | else: # bash script 90 | r = os.system(s) 91 | print('Dataset autodownload %s\n' % ('success' if r == 92 | 0 else 'failure')) # analyze return value 93 | else: 94 | raise Exception('Dataset not found.') 95 | 96 | 97 | def make_divisible(x, divisor): 98 | # Returns x evenly divisible by divisor 99 | return math.ceil(x / divisor) * divisor 100 | 101 | 102 | def labels_to_class_weights(labels, nc=80): 103 | # Get class weights (inverse frequency) from training labels 104 | if labels[0] is None: # no labels loaded 105 | return torch.Tensor() 106 | 107 | labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO 108 | classes = labels[:, 0].astype(np.int) # labels = [class xywh] 109 | weights = np.bincount(classes, minlength=nc) # occurrences per class 110 | 111 | # Prepend gridpoint count (for uCE training) 112 | # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image 113 | # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start 114 | 115 | weights[weights == 0] = 1 # replace empty bins with 1 116 | weights = 1 / weights # number of targets per class 117 | weights /= weights.sum() # normalize 118 | return torch.from_numpy(weights) 119 | 120 | 121 | def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)): 122 | # Produces image weights based on class_weights and image contents 123 | class_counts = np.array( 124 | [np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels]) 125 | image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1) 126 | # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample 127 | return image_weights 128 | 129 | 130 | def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper) 131 | # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/ 132 | # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') 133 | # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') 134 | # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco 135 | # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet 136 | x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 137 | 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 138 | 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] 139 | return x 140 | 141 | 142 | def xyxy2xywh(x): 143 | # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right 144 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 145 | y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center 146 | y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center 147 | y[:, 2] = x[:, 2] - x[:, 0] # width 148 | y[:, 3] = x[:, 3] - x[:, 1] # height 149 | return y 150 | 151 | 152 | def xywh2xyxy(x): 153 | # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right 154 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 155 | y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x 156 | y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y 157 | y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x 158 | y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y 159 | return y 160 | 161 | 162 | def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): 163 | # Rescale coords (xyxy) from img1_shape to img0_shape 164 | if ratio_pad is None: # calculate from img0_shape 165 | gain = min(img1_shape[0] / img0_shape[0], 166 | img1_shape[1] / img0_shape[1]) # gain = old / new 167 | pad = (img1_shape[1] - img0_shape[1] * gain) / \ 168 | 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding 169 | else: 170 | gain = ratio_pad[0][0] 171 | pad = ratio_pad[1] 172 | 173 | coords[:, [0, 2]] -= pad[0] # x padding 174 | coords[:, [1, 3]] -= pad[1] # y padding 175 | coords[:, :4] /= gain 176 | clip_coords(coords, img0_shape) 177 | return coords 178 | 179 | 180 | def clip_coords(boxes, img_shape): 181 | # Clip bounding xyxy bounding boxes to image shape (height, width) 182 | boxes[:, 0].clamp_(0, img_shape[1]) # x1 183 | boxes[:, 1].clamp_(0, img_shape[0]) # y1 184 | boxes[:, 2].clamp_(0, img_shape[1]) # x2 185 | boxes[:, 3].clamp_(0, img_shape[0]) # y2 186 | 187 | 188 | def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9): 189 | # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 190 | box2 = box2.T 191 | 192 | # Get the coordinates of bounding boxes 193 | if x1y1x2y2: # x1, y1, x2, y2 = box1 194 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] 195 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] 196 | else: # transform from xywh to xyxy 197 | b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2 198 | b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2 199 | b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2 200 | b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 201 | 202 | # Intersection area 203 | inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ 204 | (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) 205 | 206 | # Union Area 207 | w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps 208 | w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps 209 | union = w1 * h1 + w2 * h2 - inter + eps 210 | 211 | iou = inter / union 212 | if GIoU or DIoU or CIoU: 213 | # convex (smallest enclosing box) width 214 | cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) 215 | ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height 216 | if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 217 | c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared 218 | rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + 219 | (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared 220 | if DIoU: 221 | return iou - rho2 / c2 # DIoU 222 | elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 223 | v = (4 / math.pi ** 2) * \ 224 | torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) 225 | with torch.no_grad(): 226 | alpha = v / ((1 + eps) - iou + v) 227 | return iou - (rho2 / c2 + v * alpha) # CIoU 228 | else: # GIoU https://arxiv.org/pdf/1902.09630.pdf 229 | c_area = cw * ch + eps # convex area 230 | return iou - (c_area - union) / c_area # GIoU 231 | else: 232 | return iou # IoU 233 | 234 | 235 | def box_iou(box1, box2): 236 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 237 | """ 238 | Return intersection-over-union (Jaccard index) of boxes. 239 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 240 | Arguments: 241 | box1 (Tensor[N, 4]) 242 | box2 (Tensor[M, 4]) 243 | Returns: 244 | iou (Tensor[N, M]): the NxM matrix containing the pairwise 245 | IoU values for every element in boxes1 and boxes2 246 | """ 247 | 248 | def box_area(box): 249 | # box = 4xn 250 | return (box[2] - box[0]) * (box[3] - box[1]) 251 | 252 | area1 = box_area(box1.T) 253 | area2 = box_area(box2.T) 254 | 255 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 256 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - 257 | torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 258 | # iou = inter / (area1 + area2 - inter) 259 | return inter / (area1[:, None] + area2 - inter) 260 | 261 | 262 | def wh_iou(wh1, wh2): 263 | # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2 264 | wh1 = wh1[:, None] # [N,1,2] 265 | wh2 = wh2[None] # [1,M,2] 266 | inter = torch.min(wh1, wh2).prod(2) # [N,M] 267 | # iou = inter / (area1 + area2 - inter) 268 | return inter / (wh1.prod(2) + wh2.prod(2) - inter) 269 | 270 | 271 | def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None, agnostic=False, labels=()): 272 | """Performs Non-Maximum Suppression (NMS) on inference results 273 | 274 | Returns: 275 | detections with shape: nx6 (x1, y1, x2, y2, conf, cls) 276 | """ 277 | 278 | nc = prediction[0].shape[1] - 5 # number of classes 279 | xc = prediction[..., 4] > conf_thres # candidates 280 | 281 | # Settings 282 | # (pixels) minimum and maximum box width and height 283 | min_wh, max_wh = 2, 4096 284 | max_det = 300 # maximum number of detections per image 285 | time_limit = 10.0 # seconds to quit after 286 | redundant = True # require redundant detections 287 | multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) 288 | merge = False # use merge-NMS 289 | 290 | t = time.time() 291 | output = [torch.zeros(0, 6)] * prediction.shape[0] 292 | for xi, x in enumerate(prediction): # image index, image inference 293 | # Apply constraints 294 | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height 295 | x = x[xc[xi]] # confidence 296 | 297 | # Cat apriori labels if autolabelling 298 | if labels and len(labels[xi]): 299 | l = labels[xi] 300 | v = torch.zeros((len(l), nc + 5), device=x.device) 301 | v[:, :4] = l[:, 1:5] # box 302 | v[:, 4] = 1.0 # conf 303 | v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls 304 | x = torch.cat((x, v), 0) 305 | 306 | # If none remain process next image 307 | if not x.shape[0]: 308 | continue 309 | 310 | # Compute conf 311 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 312 | 313 | # Box (center x, center y, width, height) to (x1, y1, x2, y2) 314 | box = xywh2xyxy(x[:, :4]) 315 | 316 | # Detections matrix nx6 (xyxy, conf, cls) 317 | if multi_label: 318 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 319 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 320 | else: # best class only 321 | conf, j = x[:, 5:].max(1, keepdim=True) 322 | x = torch.cat((box, conf, j.float()), 1)[ 323 | conf.view(-1) > conf_thres] 324 | 325 | # Filter by class 326 | if classes: 327 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 328 | 329 | # Apply finite constraint 330 | # if not torch.isfinite(x).all(): 331 | # x = x[torch.isfinite(x).all(1)] 332 | 333 | # If none remain process next image 334 | n = x.shape[0] # number of boxes 335 | if not n: 336 | continue 337 | 338 | # Sort by confidence 339 | # x = x[x[:, 4].argsort(descending=True)] 340 | 341 | # Batched NMS 342 | c = x[:, 5:6] * (0 if agnostic else max_wh) # classes 343 | # boxes (offset by class), scores 344 | boxes, scores = x[:, :4] + c, x[:, 4] 345 | i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS 346 | if i.shape[0] > max_det: # limit detections 347 | i = i[:max_det] 348 | if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) 349 | # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) 350 | iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix 351 | weights = iou * scores[None] # box weights 352 | x[i, :4] = torch.mm(weights, x[:, :4]).float( 353 | ) / weights.sum(1, keepdim=True) # merged boxes 354 | if redundant: 355 | i = i[iou.sum(1) > 1] # require redundancy 356 | 357 | output[xi] = x[i] 358 | if (time.time() - t) > time_limit: 359 | break # time limit exceeded 360 | 361 | return output 362 | 363 | 364 | # from utils.general import *; strip_optimizer() 365 | def strip_optimizer(f='weights/best.pt', s=''): 366 | # Strip optimizer from 'f' to finalize training, optionally save as 's' 367 | x = torch.load(f, map_location=torch.device('cpu')) 368 | x['optimizer'] = None 369 | x['training_results'] = None 370 | x['epoch'] = -1 371 | x['model'].half() # to FP16 372 | for p in x['model'].parameters(): 373 | p.requires_grad = False 374 | torch.save(x, s or f) 375 | mb = os.path.getsize(s or f) / 1E6 # filesize 376 | print('Optimizer stripped from %s,%s %.1fMB' % 377 | (f, (' saved as %s,' % s) if s else '', mb)) 378 | 379 | 380 | def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''): 381 | # Print mutation results to evolve.txt (for use with train.py --evolve) 382 | a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys 383 | b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values 384 | # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3) 385 | c = '%10.4g' * len(results) % results 386 | print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c)) 387 | 388 | if bucket: 389 | url = 'gs://%s/evolve.txt' % bucket 390 | if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0): 391 | # download evolve.txt if larger than local 392 | os.system('gsutil cp %s .' % url) 393 | 394 | with open('evolve.txt', 'a') as f: # append result 395 | f.write(c + b + '\n') 396 | x = np.unique(np.loadtxt('evolve.txt', ndmin=2), 397 | axis=0) # load unique rows 398 | x = x[np.argsort(-fitness(x))] # sort 399 | np.savetxt('evolve.txt', x, '%10.3g') # save sort by fitness 400 | 401 | # Save yaml 402 | for i, k in enumerate(hyp.keys()): 403 | hyp[k] = float(x[0, i + 7]) 404 | with open(yaml_file, 'w') as f: 405 | results = tuple(x[0, :7]) 406 | # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3) 407 | c = '%10.4g' * len(results) % results 408 | f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len( 409 | x) + c + '\n\n') 410 | yaml.dump(hyp, f, sort_keys=False) 411 | 412 | if bucket: 413 | os.system('gsutil cp evolve.txt %s gs://%s' % 414 | (yaml_file, bucket)) # upload 415 | 416 | 417 | def apply_classifier(x, model, img, im0): 418 | # applies a second stage classifier to yolo outputs 419 | im0 = [im0] if isinstance(im0, np.ndarray) else im0 420 | for i, d in enumerate(x): # per image 421 | if d is not None and len(d): 422 | d = d.clone() 423 | 424 | # Reshape and pad cutouts 425 | b = xyxy2xywh(d[:, :4]) # boxes 426 | b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square 427 | b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad 428 | d[:, :4] = xywh2xyxy(b).long() 429 | 430 | # Rescale boxes from img_size to im0 size 431 | scale_coords(img.shape[2:], d[:, :4], im0[i].shape) 432 | 433 | # Classes 434 | pred_cls1 = d[:, 5].long() 435 | ims = [] 436 | for j, a in enumerate(d): # per item 437 | cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])] 438 | im = cv2.resize(cutout, (224, 224)) # BGR 439 | # cv2.imwrite('test%i.jpg' % j, cutout) 440 | 441 | # BGR to RGB, to 3x416x416 442 | im = im[:, :, ::-1].transpose(2, 0, 1) 443 | im = np.ascontiguousarray( 444 | im, dtype=np.float32) # uint8 to float32 445 | im /= 255.0 # 0 - 255 to 0.0 - 1.0 446 | ims.append(im) 447 | 448 | pred_cls2 = model(torch.Tensor(ims).to(d.device) 449 | ).argmax(1) # classifier prediction 450 | # retain matching class detections 451 | x[i] = x[i][pred_cls1 == pred_cls2] 452 | 453 | return x 454 | 455 | 456 | def increment_path(path, exist_ok=True, sep=''): 457 | # Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc. 458 | path = Path(path) # os-agnostic 459 | if (path.exists() and exist_ok) or (not path.exists()): 460 | return str(path) 461 | else: 462 | dirs = glob.glob(f"{path}{sep}*") # similar paths 463 | matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs] 464 | i = [int(m.groups()[0]) for m in matches if m] # indices 465 | n = max(i) + 1 if i else 2 # increment number 466 | return f"{path}{sep}{n}" # update path 467 | 468 | 469 | def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): 470 | shape = img.shape[:2] # current shape [height, width] 471 | if isinstance(new_shape, int): 472 | new_shape = (new_shape, new_shape) 473 | 474 | # Scale ratio (new / old) 475 | r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) 476 | if not scaleup: # only scale down, do not scale up (for better test mAP) 477 | r = min(r, 1.0) 478 | 479 | # Compute padding 480 | ratio = r, r # width, height ratios 481 | new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) 482 | dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - \ 483 | new_unpad[1] # wh padding 484 | if auto: # minimum rectangle 485 | dw, dh = np.mod(dw, 32), np.mod(dh, 32) # wh padding 486 | elif scaleFill: # stretch 487 | dw, dh = 0.0, 0.0 488 | new_unpad = (new_shape[1], new_shape[0]) 489 | ratio = new_shape[1] / shape[1], new_shape[0] / \ 490 | shape[0] # width, height ratios 491 | 492 | dw /= 2 # divide padding into 2 sides 493 | dh /= 2 494 | 495 | if shape[::-1] != new_unpad: # resize 496 | img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) 497 | top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) 498 | left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) 499 | img = cv2.copyMakeBorder( 500 | img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border 501 | return img, ratio, (dw, dh) 502 | 503 | -------------------------------------------------------------------------------- /utils/google_utils.py: -------------------------------------------------------------------------------- 1 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries 2 | 3 | import os 4 | import platform 5 | import subprocess 6 | import time 7 | from pathlib import Path 8 | 9 | import torch 10 | 11 | 12 | def gsutil_getsize(url=''): 13 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 14 | s = subprocess.check_output('gsutil du %s' % url, shell=True).decode('utf-8') 15 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 16 | 17 | 18 | def attempt_download(weights): 19 | # Attempt to download pretrained weights if not found locally 20 | weights = weights.strip().replace("'", '') 21 | file = Path(weights).name.lower() 22 | 23 | msg = weights + ' missing, try downloading from https://github.com/ultralytics/yolov5/releases/' 24 | models = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'] # available models 25 | redundant = False # offer second download option 26 | 27 | if file in models and not os.path.isfile(weights): 28 | # Google Drive 29 | # d = {'yolov5s.pt': '1R5T6rIyy3lLwgFXNms8whc-387H0tMQO', 30 | # 'yolov5m.pt': '1vobuEExpWQVpXExsJ2w-Mbf3HJjWkQJr', 31 | # 'yolov5l.pt': '1hrlqD1Wdei7UT4OgT785BEk1JwnSvNEV', 32 | # 'yolov5x.pt': '1mM8aZJlWTxOg7BZJvNUMrTnA2AbeCVzS'} 33 | # r = gdrive_download(id=d[file], name=weights) if file in d else 1 34 | # if r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6: # check 35 | # return 36 | 37 | try: # GitHub 38 | url = 'https://github.com/ultralytics/yolov5/releases/download/v3.1/' + file 39 | print('Downloading %s to %s...' % (url, weights)) 40 | torch.hub.download_url_to_file(url, weights) 41 | assert os.path.exists(weights) and os.path.getsize(weights) > 1E6 # check 42 | except Exception as e: # GCP 43 | print('Download error: %s' % e) 44 | assert redundant, 'No secondary mirror' 45 | url = 'https://storage.googleapis.com/ultralytics/yolov5/ckpt/' + file 46 | print('Downloading %s to %s...' % (url, weights)) 47 | r = os.system('curl -L %s -o %s' % (url, weights)) # torch.hub.download_url_to_file(url, weights) 48 | finally: 49 | if not (os.path.exists(weights) and os.path.getsize(weights) > 1E6): # check 50 | os.remove(weights) if os.path.exists(weights) else None # remove partial downloads 51 | print('ERROR: Download failure: %s' % msg) 52 | print('') 53 | return 54 | 55 | 56 | def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'): 57 | # Downloads a file from Google Drive. from utils.google_utils import *; gdrive_download() 58 | t = time.time() 59 | 60 | print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='') 61 | os.remove(name) if os.path.exists(name) else None # remove existing 62 | os.remove('cookie') if os.path.exists('cookie') else None 63 | 64 | # Attempt file download 65 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 66 | os.system('curl -c ./cookie -s -L "drive.google.com/uc?export=download&id=%s" > %s ' % (id, out)) 67 | if os.path.exists('cookie'): # large file 68 | s = 'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm=%s&id=%s" -o %s' % (get_token(), id, name) 69 | else: # small file 70 | s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id) 71 | r = os.system(s) # execute, capture return 72 | os.remove('cookie') if os.path.exists('cookie') else None 73 | 74 | # Error check 75 | if r != 0: 76 | os.remove(name) if os.path.exists(name) else None # remove partial 77 | print('Download error ') # raise Exception('Download error') 78 | return r 79 | 80 | # Unzip if archive 81 | if name.endswith('.zip'): 82 | print('unzipping... ', end='') 83 | os.system('unzip -q %s' % name) # unzip 84 | os.remove(name) # remove zip to free space 85 | 86 | print('Done (%.1fs)' % (time.time() - t)) 87 | return r 88 | 89 | 90 | def get_token(cookie="./cookie"): 91 | with open(cookie) as f: 92 | for line in f: 93 | if "download" in line: 94 | return line.split()[-1] 95 | return "" 96 | 97 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 98 | # # Uploads a file to a bucket 99 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 100 | # 101 | # storage_client = storage.Client() 102 | # bucket = storage_client.get_bucket(bucket_name) 103 | # blob = bucket.blob(destination_blob_name) 104 | # 105 | # blob.upload_from_filename(source_file_name) 106 | # 107 | # print('File {} uploaded to {}.'.format( 108 | # source_file_name, 109 | # destination_blob_name)) 110 | # 111 | # 112 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 113 | # # Uploads a blob from a bucket 114 | # storage_client = storage.Client() 115 | # bucket = storage_client.get_bucket(bucket_name) 116 | # blob = bucket.blob(source_blob_name) 117 | # 118 | # blob.download_to_filename(destination_file_name) 119 | # 120 | # print('Blob {} downloaded to {}.'.format( 121 | # source_blob_name, 122 | # destination_file_name)) 123 | -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | # Model validation metrics 2 | 3 | from pathlib import Path 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | import torch 8 | 9 | from . import general 10 | 11 | 12 | def fitness(x): 13 | # Model fitness as a weighted combination of metrics 14 | w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95] 15 | return (x[:, :4] * w).sum(1) 16 | 17 | 18 | def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision-recall_curve.png', names=[]): 19 | """ Compute the average precision, given the recall and precision curves. 20 | Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. 21 | # Arguments 22 | tp: True positives (nparray, nx1 or nx10). 23 | conf: Objectness value from 0-1 (nparray). 24 | pred_cls: Predicted object classes (nparray). 25 | target_cls: True object classes (nparray). 26 | plot: Plot precision-recall curve at mAP@0.5 27 | save_dir: Plot save directory 28 | # Returns 29 | The average precision as computed in py-faster-rcnn. 30 | """ 31 | 32 | # Sort by objectness 33 | i = np.argsort(-conf) 34 | tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] 35 | 36 | # Find unique classes 37 | unique_classes = np.unique(target_cls) 38 | 39 | # Create Precision-Recall curve and compute AP for each class 40 | px, py = np.linspace(0, 1, 1000), [] # for plotting 41 | pr_score = 0.1 # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898 42 | s = [unique_classes.shape[0], tp.shape[1]] # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95) 43 | ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s) 44 | for ci, c in enumerate(unique_classes): 45 | i = pred_cls == c 46 | n_l = (target_cls == c).sum() # number of labels 47 | n_p = i.sum() # number of predictions 48 | 49 | if n_p == 0 or n_l == 0: 50 | continue 51 | else: 52 | # Accumulate FPs and TPs 53 | fpc = (1 - tp[i]).cumsum(0) 54 | tpc = tp[i].cumsum(0) 55 | 56 | # Recall 57 | recall = tpc / (n_l + 1e-16) # recall curve 58 | r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0]) # r at pr_score, negative x, xp because xp decreases 59 | 60 | # Precision 61 | precision = tpc / (tpc + fpc) # precision curve 62 | p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0]) # p at pr_score 63 | 64 | # AP from recall-precision curve 65 | for j in range(tp.shape[1]): 66 | ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j]) 67 | if plot and (j == 0): 68 | py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5 69 | 70 | # Compute F1 score (harmonic mean of precision and recall) 71 | f1 = 2 * p * r / (p + r + 1e-16) 72 | 73 | if plot: 74 | plot_pr_curve(px, py, ap, save_dir, names) 75 | 76 | return p, r, ap, f1, unique_classes.astype('int32') 77 | 78 | 79 | def compute_ap(recall, precision): 80 | """ Compute the average precision, given the recall and precision curves. 81 | Source: https://github.com/rbgirshick/py-faster-rcnn. 82 | # Arguments 83 | recall: The recall curve (list). 84 | precision: The precision curve (list). 85 | # Returns 86 | The average precision as computed in py-faster-rcnn. 87 | """ 88 | 89 | # Append sentinel values to beginning and end 90 | mrec = recall # np.concatenate(([0.], recall, [recall[-1] + 1E-3])) 91 | mpre = precision # np.concatenate(([0.], precision, [0.])) 92 | 93 | # Compute the precision envelope 94 | mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) 95 | 96 | # Integrate area under curve 97 | method = 'interp' # methods: 'continuous', 'interp' 98 | if method == 'interp': 99 | x = np.linspace(0, 1, 101) # 101-point interp (COCO) 100 | ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate 101 | else: # 'continuous' 102 | i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes 103 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve 104 | 105 | return ap, mpre, mrec 106 | 107 | 108 | class ConfusionMatrix: 109 | # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix 110 | def __init__(self, nc, conf=0.25, iou_thres=0.45): 111 | self.matrix = np.zeros((nc + 1, nc + 1)) 112 | self.nc = nc # number of classes 113 | self.conf = conf 114 | self.iou_thres = iou_thres 115 | 116 | def process_batch(self, detections, labels): 117 | """ 118 | Return intersection-over-union (Jaccard index) of boxes. 119 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 120 | Arguments: 121 | detections (Array[N, 6]), x1, y1, x2, y2, conf, class 122 | labels (Array[M, 5]), class, x1, y1, x2, y2 123 | Returns: 124 | None, updates confusion matrix accordingly 125 | """ 126 | detections = detections[detections[:, 4] > self.conf] 127 | gt_classes = labels[:, 0].int() 128 | detection_classes = detections[:, 5].int() 129 | iou = general.box_iou(labels[:, 1:], detections[:, :4]) 130 | 131 | x = torch.where(iou > self.iou_thres) 132 | if x[0].shape[0]: 133 | matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() 134 | if x[0].shape[0] > 1: 135 | matches = matches[matches[:, 2].argsort()[::-1]] 136 | matches = matches[np.unique(matches[:, 1], return_index=True)[1]] 137 | matches = matches[matches[:, 2].argsort()[::-1]] 138 | matches = matches[np.unique(matches[:, 0], return_index=True)[1]] 139 | else: 140 | matches = np.zeros((0, 3)) 141 | 142 | n = matches.shape[0] > 0 143 | m0, m1, _ = matches.transpose().astype(np.int16) 144 | for i, gc in enumerate(gt_classes): 145 | j = m0 == i 146 | if n and sum(j) == 1: 147 | self.matrix[gc, detection_classes[m1[j]]] += 1 # correct 148 | else: 149 | self.matrix[gc, self.nc] += 1 # background FP 150 | 151 | if n: 152 | for i, dc in enumerate(detection_classes): 153 | if not any(m1 == i): 154 | self.matrix[self.nc, dc] += 1 # background FN 155 | 156 | def matrix(self): 157 | return self.matrix 158 | 159 | def plot(self, save_dir='', names=()): 160 | try: 161 | import seaborn as sn 162 | 163 | array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6) # normalize 164 | array[array < 0.005] = np.nan # don't annotate (would appear as 0.00) 165 | 166 | fig = plt.figure(figsize=(12, 9), tight_layout=True) 167 | sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size 168 | labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels 169 | sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, 170 | xticklabels=names + ['background FN'] if labels else "auto", 171 | yticklabels=names + ['background FP'] if labels else "auto").set_facecolor((1, 1, 1)) 172 | fig.axes[0].set_xlabel('True') 173 | fig.axes[0].set_ylabel('Predicted') 174 | fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) 175 | except Exception as e: 176 | pass 177 | 178 | def print(self): 179 | for i in range(self.nc + 1): 180 | print(' '.join(map(str, self.matrix[i]))) 181 | 182 | 183 | # Plots ---------------------------------------------------------------------------------------------------------------- 184 | 185 | def plot_pr_curve(px, py, ap, save_dir='.', names=()): 186 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) 187 | py = np.stack(py, axis=1) 188 | 189 | if 0 < len(names) < 21: # show mAP in legend if < 10 classes 190 | for i, y in enumerate(py.T): 191 | ax.plot(px, y, linewidth=1, label=f'{names[i]} %.3f' % ap[i, 0]) # plot(recall, precision) 192 | else: 193 | ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision) 194 | 195 | ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean()) 196 | ax.set_xlabel('Recall') 197 | ax.set_ylabel('Precision') 198 | ax.set_xlim(0, 1) 199 | ax.set_ylim(0, 1) 200 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") 201 | fig.savefig(Path(save_dir) / 'precision_recall_curve.png', dpi=250) 202 | -------------------------------------------------------------------------------- /utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | # PyTorch utils 2 | 3 | import logging 4 | import math 5 | import os 6 | import time 7 | from contextlib import contextmanager 8 | from copy import deepcopy 9 | 10 | import torch 11 | import torch.backends.cudnn as cudnn 12 | import torch.nn as nn 13 | import torch.nn.functional as F 14 | import torchvision 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | @contextmanager 20 | def torch_distributed_zero_first(local_rank: int): 21 | """ 22 | Decorator to make all processes in distributed training wait for each local_master to do something. 23 | """ 24 | if local_rank not in [-1, 0]: 25 | torch.distributed.barrier() 26 | yield 27 | if local_rank == 0: 28 | torch.distributed.barrier() 29 | 30 | 31 | def init_torch_seeds(seed=0): 32 | # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html 33 | torch.manual_seed(seed) 34 | if seed == 0: # slower, more reproducible 35 | cudnn.deterministic = True 36 | cudnn.benchmark = False 37 | else: # faster, less reproducible 38 | cudnn.deterministic = False 39 | cudnn.benchmark = True 40 | 41 | 42 | def select_device(device='', batch_size=None): 43 | # device = 'cpu' or '0' or '0,1,2,3' 44 | cpu_request = device.lower() == 'cpu' 45 | if device and not cpu_request: # if device requested other than 'cpu' 46 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable 47 | assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity 48 | 49 | cuda = False if cpu_request else torch.cuda.is_available() 50 | if cuda: 51 | c = 1024 ** 2 # bytes to MB 52 | ng = torch.cuda.device_count() 53 | if ng > 1 and batch_size: # check that batch_size is compatible with device_count 54 | assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng) 55 | x = [torch.cuda.get_device_properties(i) for i in range(ng)] 56 | s = f'Using torch {torch.__version__} ' 57 | for i in range(0, ng): 58 | if i == 1: 59 | s = ' ' * len(s) 60 | logger.info("%sCUDA:%g (%s, %dMB)" % (s, i, x[i].name, x[i].total_memory / c)) 61 | else: 62 | logger.info(f'Using torch {torch.__version__} CPU') 63 | 64 | logger.info('') # skip a line 65 | return torch.device('cuda:0' if cuda else 'cpu') 66 | 67 | 68 | def time_synchronized(): 69 | torch.cuda.synchronize() if torch.cuda.is_available() else None 70 | return time.time() 71 | 72 | 73 | def is_parallel(model): 74 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 75 | 76 | 77 | def intersect_dicts(da, db, exclude=()): 78 | # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values 79 | return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape} 80 | 81 | 82 | def initialize_weights(model): 83 | for m in model.modules(): 84 | t = type(m) 85 | if t is nn.Conv2d: 86 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 87 | elif t is nn.BatchNorm2d: 88 | m.eps = 1e-3 89 | m.momentum = 0.03 90 | elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 91 | m.inplace = True 92 | 93 | 94 | def find_modules(model, mclass=nn.Conv2d): 95 | # Finds layer indices matching module class 'mclass' 96 | return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)] 97 | 98 | 99 | def sparsity(model): 100 | # Return global model sparsity 101 | a, b = 0., 0. 102 | for p in model.parameters(): 103 | a += p.numel() 104 | b += (p == 0).sum() 105 | return b / a 106 | 107 | 108 | def prune(model, amount=0.3): 109 | # Prune model to requested global sparsity 110 | import torch.nn.utils.prune as prune 111 | print('Pruning model... ', end='') 112 | for name, m in model.named_modules(): 113 | if isinstance(m, nn.Conv2d): 114 | prune.l1_unstructured(m, name='weight', amount=amount) # prune 115 | prune.remove(m, 'weight') # make permanent 116 | print(' %.3g global sparsity' % sparsity(model)) 117 | 118 | 119 | def fuse_conv_and_bn(conv, bn): 120 | # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/ 121 | fusedconv = nn.Conv2d(conv.in_channels, 122 | conv.out_channels, 123 | kernel_size=conv.kernel_size, 124 | stride=conv.stride, 125 | padding=conv.padding, 126 | groups=conv.groups, 127 | bias=True).requires_grad_(False).to(conv.weight.device) 128 | 129 | # prepare filters 130 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 131 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 132 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) 133 | 134 | # prepare spatial bias 135 | b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias 136 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) 137 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) 138 | 139 | return fusedconv 140 | 141 | 142 | def model_info(model, verbose=False, img_size=640): 143 | # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320] 144 | n_p = sum(x.numel() for x in model.parameters()) # number parameters 145 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients 146 | if verbose: 147 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) 148 | for i, (name, p) in enumerate(model.named_parameters()): 149 | name = name.replace('module_list.', '') 150 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' % 151 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) 152 | 153 | try: # FLOPS 154 | from thop import profile 155 | stride = int(model.stride.max()) 156 | img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device) # input 157 | flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride FLOPS 158 | img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float 159 | fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 FLOPS 160 | except (ImportError, Exception): 161 | fs = '' 162 | 163 | logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}") 164 | 165 | 166 | def load_classifier(name='resnet101', n=2): 167 | # Loads a pretrained model reshaped to n-class output 168 | model = torchvision.models.__dict__[name](pretrained=True) 169 | 170 | # ResNet model properties 171 | # input_size = [3, 224, 224] 172 | # input_space = 'RGB' 173 | # input_range = [0, 1] 174 | # mean = [0.485, 0.456, 0.406] 175 | # std = [0.229, 0.224, 0.225] 176 | 177 | # Reshape output to n classes 178 | filters = model.fc.weight.shape[1] 179 | model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True) 180 | model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True) 181 | model.fc.out_features = n 182 | return model 183 | 184 | 185 | def scale_img(img, ratio=1.0, same_shape=False): # img(16,3,256,416), r=ratio 186 | # scales img(bs,3,y,x) by ratio 187 | if ratio == 1.0: 188 | return img 189 | else: 190 | h, w = img.shape[2:] 191 | s = (int(h * ratio), int(w * ratio)) # new size 192 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize 193 | if not same_shape: # pad/crop img 194 | gs = 32 # (pixels) grid size 195 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] 196 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean 197 | 198 | 199 | def copy_attr(a, b, include=(), exclude=()): 200 | # Copy attributes from b to a, options to only include [...] and to exclude [...] 201 | for k, v in b.__dict__.items(): 202 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 203 | continue 204 | else: 205 | setattr(a, k, v) 206 | 207 | 208 | class ModelEMA: 209 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 210 | Keep a moving average of everything in the model state_dict (parameters and buffers). 211 | This is intended to allow functionality like 212 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 213 | A smoothed version of the weights is necessary for some training schemes to perform well. 214 | This class is sensitive where it is initialized in the sequence of model init, 215 | GPU assignment and distributed training wrappers. 216 | """ 217 | 218 | def __init__(self, model, decay=0.9999, updates=0): 219 | # Create EMA 220 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA 221 | # if next(model.parameters()).device.type != 'cpu': 222 | # self.ema.half() # FP16 EMA 223 | self.updates = updates # number of EMA updates 224 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) 225 | for p in self.ema.parameters(): 226 | p.requires_grad_(False) 227 | 228 | def update(self, model): 229 | # Update EMA parameters 230 | with torch.no_grad(): 231 | self.updates += 1 232 | d = self.decay(self.updates) 233 | 234 | msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict 235 | for k, v in self.ema.state_dict().items(): 236 | if v.dtype.is_floating_point: 237 | v *= d 238 | v += (1. - d) * msd[k].detach() 239 | 240 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): 241 | # Update EMA attributes 242 | copy_attr(self.ema, model, include, exclude) 243 | -------------------------------------------------------------------------------- /weights/yolov5m.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sharpiless/Yolov5-deepsort-inference/2644c110efa17e8fa12f4dff5ba4a9d229ddf638/weights/yolov5m.pt --------------------------------------------------------------------------------