├── Demo_Tracking.py ├── README.md ├── coco.names └── sort.py /Demo_Tracking.py: -------------------------------------------------------------------------------- 1 | from sort import * 2 | import random 3 | import cv2 4 | 5 | def load_class_names(namesfile): 6 | class_names = [] 7 | with open(namesfile, 'r') as fp: 8 | lines = fp.readlines() 9 | for line in lines: 10 | line = line.rstrip() 11 | class_names.append(line) 12 | return class_names 13 | 14 | 15 | def detect_cv2_camera(cfgfile, weightfile): 16 | #Заграузка сети 17 | net = cv2.dnn.readNetFromDarknet(cfgfile, weightfile) 18 | ln = net.getLayerNames() 19 | ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()] 20 | #Инициализируем трекер 21 | mot_tracker = Sort() 22 | #Используем YOLO 23 | namesfile = 'coco.names' 24 | class_names = load_class_names(namesfile) 25 | #Камера 26 | cap = cv2.VideoCapture(0) 27 | #Генерим коробочкам разные цвета 28 | color_list = [] 29 | for j in range(1000): 30 | color_list.append(((int)(random.randrange(255)),(int)(random.randrange(255)),(int)(random.randrange(255)))) 31 | 32 | #go! 33 | ret = True 34 | while ret: 35 | ret, img = cap.read() 36 | if ret: 37 | #Запустим сеть по картинке 38 | blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (416, 416), 39 | swapRB=True, crop=False) 40 | 41 | h,w,_ = img.shape 42 | net.setInput(blob) 43 | layerOutputs = net.forward(ln) 44 | #Разберём все выходы 45 | boxes=[] 46 | confidences=[] 47 | classIDs=[] 48 | for output in layerOutputs: 49 | for detection in output: 50 | scores = detection[5:] 51 | classID = np.argmax(scores) 52 | confidence = scores[classID] 53 | if confidence > 0.5: 54 | box = detection[0:4] * np.array([w, h, w, h]) 55 | (centerX, centerY, width, height) = box.astype("int") 56 | x = int(centerX - (width / 2)) 57 | y = int(centerY - (height / 2)) 58 | boxes.append([x, y, int(width), int(height)]) 59 | confidences.append(float(confidence)) 60 | classIDs.append(classID) 61 | 62 | idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 63 | 0.3) 64 | 65 | 66 | result_img = np.copy(img) 67 | dets = [] 68 | count_detection=0 69 | for j in range(len(idxs)): 70 | name = class_names[classIDs[idxs[j][0]]] 71 | if name == 'person': 72 | count_detection+=1 73 | if count_detection>0: 74 | detects = np.zeros((count_detection,5)) 75 | count=0 76 | #Подготовим в формат который ест трекер 77 | for j in range(len(idxs)): 78 | b = boxes[j] 79 | name = class_names[classIDs[idxs[j][0]]] 80 | if name == 'person': 81 | x1 = int(b[0]) 82 | y1 = int(b[1]) 83 | x2 = int((b[0] + b[2])) 84 | y2 = int((b[1] + b[3])) 85 | box = np.array([x1,y1,x2,y2,confidences[idxs[j][0]]]) 86 | detects[count,:] = box[:] 87 | count+=1 88 | #Подаём в трекер и сразу имеем результат! 89 | if len(detects)!=0: 90 | trackers = mot_tracker.update(detects) 91 | for d in trackers: 92 | result_img = cv2.rectangle(result_img, ((int)(d[0]), (int)(d[1])), ((int)(d[2]), (int)(d[3])), color_list[(int)(d[4])], 2) 93 | cv2.imshow('demo', result_img) 94 | cv2.waitKey(1) 95 | cap.release() 96 | 97 | 98 | 99 | if __name__ == '__main__': 100 | detect_cv2_camera('yolov3-tiny.cfg', 'yolov3-tiny.weights') -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SimpleTrackSample 2 | Very easy sample for tracking algoritm 3 | 4 | Sort algorithm from here - https://github.com/abewley/sort 5 | Net weight from here - https://pjreddie.com/darknet/yolo/ 6 | More about (in russian) - http://cv-blog.ru/ 7 | -------------------------------------------------------------------------------- /coco.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | aeroplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | sofa 59 | pottedplant 60 | bed 61 | diningtable 62 | toilet 63 | tvmonitor 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush 81 | -------------------------------------------------------------------------------- /sort.py: -------------------------------------------------------------------------------- 1 | """ 2 | SORT: A Simple, Online and Realtime Tracker 3 | Copyright (C) 2016-2020 Alex Bewley alex@bewley.ai 4 | 5 | This program is free software: you can redistribute it and/or modify 6 | it under the terms of the GNU General Public License as published by 7 | the Free Software Foundation, either version 3 of the License, or 8 | (at your option) any later version. 9 | 10 | This program is distributed in the hope that it will be useful, 11 | but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 | GNU General Public License for more details. 14 | 15 | You should have received a copy of the GNU General Public License 16 | along with this program. If not, see . 17 | """ 18 | 19 | import numpy as np 20 | 21 | try: 22 | from numba import jit 23 | except: 24 | def jit(func): 25 | return func 26 | 27 | np.random.seed(0) 28 | 29 | from filterpy.kalman import KalmanFilter 30 | 31 | 32 | def linear_assignment(cost_matrix): 33 | try: 34 | import lap 35 | _, x, y = lap.lapjv(cost_matrix, extend_cost=True) 36 | return np.array([[y[i],i] for i in x if i >= 0]) # 37 | except ImportError: 38 | from scipy.optimize import linear_sum_assignment 39 | x, y = linear_sum_assignment(cost_matrix) 40 | return np.array(list(zip(x, y))) 41 | 42 | 43 | @jit 44 | def iou(bb_test, bb_gt): 45 | """ 46 | Computes IUO between two bboxes in the form [x1,y1,x2,y2] 47 | """ 48 | xx1 = np.maximum(bb_test[0], bb_gt[0]) 49 | yy1 = np.maximum(bb_test[1], bb_gt[1]) 50 | xx2 = np.minimum(bb_test[2], bb_gt[2]) 51 | yy2 = np.minimum(bb_test[3], bb_gt[3]) 52 | w = np.maximum(0., xx2 - xx1) 53 | h = np.maximum(0., yy2 - yy1) 54 | wh = w * h 55 | o = wh / ((bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1]) 56 | + (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1]) - wh) 57 | return(o) 58 | 59 | 60 | def convert_bbox_to_z(bbox): 61 | """ 62 | Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form 63 | [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is 64 | the aspect ratio 65 | """ 66 | w = bbox[2] - bbox[0] 67 | h = bbox[3] - bbox[1] 68 | x = bbox[0] + w/2. 69 | y = bbox[1] + h/2. 70 | s = w * h #scale is just area 71 | r = w / float(h) 72 | return np.array([x, y, s, r]).reshape((4, 1)) 73 | 74 | 75 | def convert_x_to_bbox(x,score=None): 76 | """ 77 | Takes a bounding box in the centre form [x,y,s,r] and returns it in the form 78 | [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right 79 | """ 80 | w = np.sqrt(x[2] * x[3]) 81 | h = x[2] / w 82 | if(score==None): 83 | return np.array([x[0]-w/2.,x[1]-h/2.,x[0]+w/2.,x[1]+h/2.]).reshape((1,4)) 84 | else: 85 | return np.array([x[0]-w/2.,x[1]-h/2.,x[0]+w/2.,x[1]+h/2.,score]).reshape((1,5)) 86 | 87 | 88 | class KalmanBoxTracker(object): 89 | """ 90 | This class represents the internal state of individual tracked objects observed as bbox. 91 | """ 92 | count = 0 93 | def __init__(self,bbox): 94 | """ 95 | Initialises a tracker using initial bounding box. 96 | """ 97 | #define constant velocity model 98 | self.kf = KalmanFilter(dim_x=7, dim_z=4) 99 | self.kf.F = np.array([[1,0,0,0,1,0,0],[0,1,0,0,0,1,0],[0,0,1,0,0,0,1],[0,0,0,1,0,0,0], [0,0,0,0,1,0,0],[0,0,0,0,0,1,0],[0,0,0,0,0,0,1]]) 100 | self.kf.H = np.array([[1,0,0,0,0,0,0],[0,1,0,0,0,0,0],[0,0,1,0,0,0,0],[0,0,0,1,0,0,0]]) 101 | 102 | self.kf.R[2:,2:] *= 10. 103 | self.kf.P[4:,4:] *= 1000. #give high uncertainty to the unobservable initial velocities 104 | self.kf.P *= 10. 105 | self.kf.Q[-1,-1] *= 0.01 106 | self.kf.Q[4:,4:] *= 0.01 107 | 108 | self.kf.x[:4] = convert_bbox_to_z(bbox) 109 | self.time_since_update = 0 110 | self.id = KalmanBoxTracker.count 111 | KalmanBoxTracker.count += 1 112 | self.history = [] 113 | self.hits = 0 114 | self.hit_streak = 0 115 | self.age = 0 116 | 117 | def update(self,bbox): 118 | """ 119 | Updates the state vector with observed bbox. 120 | """ 121 | self.time_since_update = 0 122 | self.history = [] 123 | self.hits += 1 124 | self.hit_streak += 1 125 | self.kf.update(convert_bbox_to_z(bbox)) 126 | 127 | def predict(self): 128 | """ 129 | Advances the state vector and returns the predicted bounding box estimate. 130 | """ 131 | if((self.kf.x[6]+self.kf.x[2])<=0): 132 | self.kf.x[6] *= 0.0 133 | self.kf.predict() 134 | self.age += 1 135 | if(self.time_since_update>0): 136 | self.hit_streak = 0 137 | self.time_since_update += 1 138 | self.history.append(convert_x_to_bbox(self.kf.x)) 139 | return self.history[-1] 140 | 141 | def get_state(self): 142 | """ 143 | Returns the current bounding box estimate. 144 | """ 145 | return convert_x_to_bbox(self.kf.x) 146 | 147 | 148 | def associate_detections_to_trackers(detections,trackers,iou_threshold = 0.3): 149 | """ 150 | Assigns detections to tracked object (both represented as bounding boxes) 151 | 152 | Returns 3 lists of matches, unmatched_detections and unmatched_trackers 153 | """ 154 | if(len(trackers)==0): 155 | return np.empty((0,2),dtype=int), np.arange(len(detections)), np.empty((0,5),dtype=int) 156 | iou_matrix = np.zeros((len(detections),len(trackers)),dtype=np.float32) 157 | 158 | for d,det in enumerate(detections): 159 | for t,trk in enumerate(trackers): 160 | iou_matrix[d,t] = iou(det,trk) 161 | 162 | if min(iou_matrix.shape) > 0: 163 | a = (iou_matrix > iou_threshold).astype(np.int32) 164 | if a.sum(1).max() == 1 and a.sum(0).max() == 1: 165 | matched_indices = np.stack(np.where(a), axis=1) 166 | else: 167 | matched_indices = linear_assignment(-iou_matrix) 168 | else: 169 | matched_indices = np.empty(shape=(0,2)) 170 | 171 | unmatched_detections = [] 172 | for d, det in enumerate(detections): 173 | if(d not in matched_indices[:,0]): 174 | unmatched_detections.append(d) 175 | unmatched_trackers = [] 176 | for t, trk in enumerate(trackers): 177 | if(t not in matched_indices[:,1]): 178 | unmatched_trackers.append(t) 179 | 180 | #filter out matched with low IOU 181 | matches = [] 182 | for m in matched_indices: 183 | if(iou_matrix[m[0], m[1]]= self.min_hits or self.frame_count <= self.min_hits): 242 | ret.append(np.concatenate((d,[trk.id+1])).reshape(1,-1)) # +1 as MOT benchmark requires positive 243 | i -= 1 244 | # remove dead tracklet 245 | if(trk.time_since_update > self.max_age): 246 | self.trackers.pop(i) 247 | if(len(ret)>0): 248 | return np.concatenate(ret) 249 | return np.empty((0,5)) 250 | 251 | --------------------------------------------------------------------------------