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