def assign_detections_to_trackers(trackers, detections, iou_thrd = 0.3): IOU_mat = np.zeros((len(trackers), len(detections)), dtype=np.float32) for t, trk in enumerate(trackers): for d, det in enumerate(detections): IOU_mat[t, d] = Geometry.box_iou(trk, det) # solve the maximizing the sum of IOU assignment problem using the Hungarian algorithm matched_idx = linear_assignment(-IOU_mat) unmatched_trackers, unmatched_detections = [], [] for t, trk in enumerate(trackers): if t not in matched_idx[:, 0]: unmatched_trackers.append(t) for d, det in enumerate(detections): if d not in matched_idx[:, 1]: unmatched_detections.append(d) matches = [] for m in matched_idx: if IOU_mat[m[0], m[1]] < iou_thrd: unmatched_trackers.append(m[0]) unmatched_detections.append(m[1]) else: matches.append(m.reshape(1, 2)) if len(matches) == 0: matches = np.empty((0, 2), dtype=int) else: matches = np.concatenate(matches, axis=0) return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
def counting(trk, start, end, cnt, index, results, svm, frame): flag = 0 if trk.is_counted: return trk, cnt, results, flag cur_center = trk.cur_center pre_center = trk.pre_center #cls = trk.cls if pre_center == []: return trk, cnt, results, flag if Geometry.is_segment_cross(start, end, pre_center, cur_center): cnt[index] += 1 flag = 1 if not trk.is_crossed_first_line: trk.crossed_line[0] = index trk.is_crossed_first_line = True elif index != trk.crossed_line[0]: trk.crossed_line[1] = index trk.is_crossed_second_line = True print(trk.cur_box) top, bottom = min(trk.cur_box[1], trk.cur_box[3]), max(trk.cur_box[1], trk.cur_box[3]) left, right = min(trk.cur_box[0], trk.cur_box[2]), max(trk.cur_box[0], trk.cur_box[2]) img = frame[top:bottom, left:right] cls = int(svm_vehicle.get_label(svm, img)) trk.cls = cls print(cls) results[cls - 1][trk.crossed_line[0]][trk.crossed_line[1]] += 1 trk.is_counted = True return trk, cnt, results, flag
def counting(trk, start, end, cnt, index, results): flag = 0 if trk.is_counted: return trk, cnt, results, flag cur_center = trk.cur_center pre_center = trk.pre_center cls = trk.cls if pre_center == []: return trk, cnt, results, flag if Geometry.is_segment_cross(start, end, pre_center, cur_center): flag = 1 cnt[index] += 1 if not trk.is_crossed_first_line: trk.crossed_line[0] = index trk.is_crossed_first_line = True elif index != trk.crossed_line[0]: trk.crossed_line[1] = index trk.is_crossed_second_line = True results[cls - 1][trk.crossed_line[0]][trk.crossed_line[1]] += 1 trk.is_counted = True return trk, cnt, results, flag
def pipeline(frame, frame_num, lines, cnt, results, debug, model): global tracker_list global max_age global min_hits global track_id_list original_frame = frame # detect detections = DetectionRCNN.start_detect(frame, frame_num, model) # detections = detections[np.where(detections[:, 7] == 1)] # track z_box = [] x_box = [] clss = [] for i in range(len(detections)): z_box.append(detections[i][3:7]) clss.append(detections[i][7]) if len(tracker_list) > 0: for trk in tracker_list: x_box.append(trk.cur_box + [trk.id]) matched, unmatched_dets, unmatched_trks = MatchKalHun.assign_detections_to_trackers( x_box, z_box, iou_thrd=0.3) # Deal with matched detections if matched.size > 0: for trk_idx, det_idx in matched: z = z_box[det_idx] z = np.expand_dims(z, axis=0).T tmp_trk = tracker_list[trk_idx] tmp_trk.kalman_filter(z) xx = tmp_trk.x_state.T[0].tolist() xx = [xx[0], xx[2], xx[4], xx[6]] pre_box = tmp_trk.cur_box tmp_trk.pre_box = pre_box tmp_trk.cur_box = xx # tmp_trk.cls = clss[det_idx] if tmp_trk.cls != 5 and tmp_trk.cls != 6: tmp_trk.cls = clss[det_idx] tmp_trk.hits += 1 tmp_trk.no_losses = 0 # Deal with unmatched detections if len(unmatched_dets) > 0: for idx in unmatched_dets: z = z_box[idx] z = np.expand_dims(z, axis=0).T tmp_trk = MatchKalHun.Tracker() x = np.array([[z[0], 0, z[1], 0, z[2], 0, z[3], 0]]).T tmp_trk.x_state = x tmp_trk.predict_only() xx = tmp_trk.x_state xx = xx.T[0].tolist() xx = [xx[0], xx[2], xx[4], xx[6]] # tmp_trk.pre_box = z tmp_trk.cur_box = xx tmp_trk.cls = clss[idx] tmp_trk.id = track_id_list.popleft() # new id add into the track list tracker_list.append(tmp_trk) # Deal with unmatched tracks if len(unmatched_trks) > 0: for trk_idx in unmatched_trks: tmp_trk = tracker_list[trk_idx] tmp_trk.no_losses += 1 tmp_trk.predict_only() xx = tmp_trk.x_state xx = xx.T[0].tolist() xx = [xx[0], xx[2], xx[4], xx[6]] pre_box = tmp_trk.cur_box tmp_trk.pre_box = pre_box tmp_trk.cur_box = xx # count and show good_tracker_list = [] flag = np.zeros(len(lines)) for trk in tracker_list: if (trk.hits >= min_hits) and (trk.no_losses <= max_age): if debug: x_cv2 = trk.cur_box id = trk.id cls = trk.cls frame = Geometry.draw_box_label(frame, x_cv2, id, cls) trk.get_center() for j in range(len(lines)): trk, cnt, results, flag_temp = Count.counting( trk, lines[j][0:2], lines[j][2:4], cnt, j, results) flag[j] += flag_temp good_tracker_list.append(trk) if debug: frame = Geometry.draw_lines(frame, lines, flag) tracker_list = [x for x in tracker_list if x.no_losses <= max_age] return frame, cnt, results