def main(video_path,
model_path,
track_target=0,
visualize=True):
"""run video prediction
Args:
video_path: video path
model_path: model path
track_target: 0-person; 1-bicycle; 2-car; 7-truck
visualize: whether visualize tracking list
"""
detector = Detector(model_path=model_path)
kalman_filter = KalmanFilter()
capture = cv2.VideoCapture(video_path)
height = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)
width = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
# tracking list
tracking_list = []
label_count = 0
is_first_frame = True
while True:
success, frame = capture.read()
if not success:
capture.release()
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# convert to Image object
frame_pil = Image.fromarray(np.uint8(frame))
new_frame = letterbox_image(frame_pil, INPUT_SIZE)
image_array = np.expand_dims(np.array(new_frame, dtype='float32') / 255.0, axis=0)
image_shape = np.expand_dims(np.array([height, width], dtype='float32'), axis=0)
image_constant = tf.constant(image_array, dtype=tf.float32)
image_shape = tf.constant(image_shape, dtype=tf.float32)
# detect image
results = detector.detect(image_constant, image_shape)
pred_results = []
for key, value in results.items():
pred_results.append(value)
boxes = pred_results[0].numpy()
# scores = scores.numpy
classes = pred_results[2].numpy()
# find tracking targets
track_id = np.where(classes == track_target)[0]
track_boxes = boxes[track_id]
num_tracks = len(track_boxes)
if num_tracks > 0:
track_boxes = box2xyah(track_boxes)
track_boxes = [track_box for track_box in track_boxes]
if not is_first_frame:
# start tracking
tracking_list, label_count = matching_cascade(tracking_list, track_boxes,
kalman_filter, label_count)
if is_first_frame and (num_tracks > 0):
is_first_frame = False
for i in range(num_tracks):
# initialize first frame
mean_init, cov_init = kalman_filter.initiate(measurement=track_boxes[i])
# create tracker
new_tracker = create_tracker(mean=mean_init,
cov=cov_init,
detection=track_boxes[i])
tracking_list.append(new_tracker)
if visualize:
# visulize results
img = visualize_results(tracking_list, height, frame)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imshow('avoid invasion', img)
key = cv2.waitKey(30) & 0xff
if key == 27:
capture.release()
break
class Tracker(object):
"""
This is the multi-target tracker.
Parameters
----------
metric : nn_matching.NearestNeighborDistanceMetric
A distance metric for measurement-to-track association.
max_age : int
Maximum number of missed misses before a track is deleted.
n_init : int
Number of consecutive detections before the track is confirmed. The
track state is set to `Deleted` if a miss occurs within the first
`n_init` frames.
Attributes
----------
metric : nn_matching.NearestNeighborDistanceMetric
The distance metric used for measurement to track association.
max_age : int
Maximum number of missed misses before a track is deleted.
n_init : int
Number of frames that a track remains in initialization phase.
kf : kalman_filter.KalmanFilter
A Kalman filter to filter target trajectories in image space.
tracks : List[Track]
The list of active tracks at the current time step.
"""
def __init__(self, metric, max_iou_distance=0.7, max_age=70, n_init=3):
self.metric = metric
self.max_iou_distance = max_iou_distance
self.max_age = max_age
self.n_init = n_init
self.kf = KalmanFilter()
self.tracks = []
self._next_id = 1
self.last_confirm_id = 1
def predict(self):
"""Propagate track state distributions one time step forward.
This function should be called once every time step, before `update`.
"""
if cfgs.debug:
print('tracker pred:', len(self.tracks))
for track in self.tracks:
track.predict(self.kf)
#print('run predict')
def update(self, detections):
"""Perform measurement update and track management.
Parameters
----------
detections : List[deep_sort.detection.Detection]
A list of detections at the current time step.
"""
if cfgs.debug:
print('updata_detect:', len(detections))
# Run matching cascade.
matches, unmatched_tracks, unmatched_detections = \
self._match(detections)
if cfgs.debug:
print('unmatch_track for delete:', len(unmatched_tracks))
print('unmatch_det for init track:', len(unmatched_detections))
# Update track set.
for track_idx, detection_idx in matches:
#print("run match")
self.last_confirm_id = self.tracks[track_idx].update(
self.kf, detections[detection_idx], self.last_confirm_id)
for track_idx in unmatched_tracks:
self.tracks[track_idx].mark_missed()
#print("run unmatch_track for delete")
for detection_idx in unmatched_detections:
#print("run unmatch_det for init track")
self._initiate_track(detections[detection_idx])
self.tracks = [t for t in self.tracks if not t.is_deleted()]
# Update distance metric.
active_targets = [t.track_id for t in self.tracks if t.is_confirmed()]
if cfgs.debug:
print('activate:', len(active_targets))
features, targets = [], []
for track in self.tracks:
#print('track_feature',len(track.features))
if not track.is_confirmed():
continue
features += track.features
targets += [track.track_id for _ in track.features]
track.features = []
self.metric.partial_fit(np.asarray(features), np.asarray(targets),
active_targets)
def _match(self, detections):
# Split track set into confirmed and unconfirmed tracks.
confirmed_tracks = [
i for i, t in enumerate(self.tracks) if t.is_confirmed()
]
unconfirmed_tracks = [
i for i, t in enumerate(self.tracks) if not t.is_confirmed()
]
if cfgs.debug:
print('confirm and unconfirm:', len(confirmed_tracks),
len(unconfirmed_tracks))
# Associate confirmed tracks using appearance features.
matches_a, unmatched_tracks_a, unmatched_detections = \
matching_cascade(self.metric,self.kf, self.metric.matching_threshold, self.max_age,
self.tracks, detections, confirmed_tracks)
# Associate remaining tracks together with unconfirmed tracks using IOU.
iou_track_candidates = unconfirmed_tracks + [
k for k in unmatched_tracks_a
if self.tracks[k].time_since_update == 1
]
unmatched_tracks_a = [
k for k in unmatched_tracks_a
if self.tracks[k].time_since_update != 1
]
cost_matrix = iou_cost(self.tracks, detections, iou_track_candidates,
unmatched_detections)
matches_b, unmatched_tracks_b, unmatched_detections = \
min_cost_matching(
cost_matrix, self.max_iou_distance, self.tracks,
detections, iou_track_candidates, unmatched_detections)
if cfgs.debug:
print('iou mach tracks:: unconfirm and unmatch_track:',
iou_track_candidates)
print("matcha and matchb", len(matches_a), len(matches_b))
print('a:', matches_a)
print('b:', matches_b)
matches = matches_a + matches_b
unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b))
return matches, unmatched_tracks, unmatched_detections
def _initiate_track(self, detection):
mean, covariance = self.kf.initiate(detection.to_xyah())
self.tracks.append(
Track(mean, covariance, self._next_id, self.n_init, self.max_age,
detection.feature))
self._next_id += 1
class Track:
def __init__(self, bb, id_, max_miss_):
self.track = [bb]
self.alive = True
self.kf = KalmanFilter()
self.id_track = id_
self.mean, self.cov = self.kf.initiate(bb.asp_ratio())
self.num_miss = 0
self.num_max_miss = max_miss_
self.project_mean = 0
self.color = (int(random.random() * 256), int(random.random() * 256),
int(random.random() * 256))
def last(self):
return self.track[-1]
def get_last_mean(self):
return self.mean[0:4]
def append(self, bb):
self.kalman_steps(bb.asp_ratio())
aproxBB = BoudingBox(coord=_restoreStandardValue(self.get_last_mean()),
frame=bb.getIDFrame())
self.track.append(aproxBB)
def kalman_steps(self, dets):
self.mean, self.cov = self.kf.predict(self.mean, self.cov)
self.mean, self.cov = self.kf.update(self.mean, self.cov, dets)
def getTrack(self):
return self.track
def getBBFrame(self, frame_id):
for bb in self.track:
if bb.getIDFrame() == frame_id:
return bb
def __len__(self):
return len(self.track)
def is_alive(self):
return self.alive
def ressurect(self):
self.alive = True
def project_position(self):
self.append(
BoudingBox(_restoreStandardValue(self.mean[:4]),
self.last().getIDFrame() + 1))
def missed(self):
self.num_miss += 1
self.project_position()
if self.num_miss > self.num_max_miss:
self.kill()
def kill(self):
self.alive = False
def getID(self):
return self.id_track
def getcolor(self):
return self.color
